import os
import requests
import tempfile
from typing import List
from langchain_community.document_loaders import PyPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain_core.tools import Tool
from pinecone import Pinecone, ServerlessSpec
from docx import Document as DocxDocument
from dotenv import load_dotenv

load_dotenv()

PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_INDEX_NAME = os.getenv("document-processor-index-1") or "policy-index"

pc = Pinecone(api_key=PINECONE_API_KEY)
embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

def download_file(url: str) -> str:
    try:
        response = requests.get(url)
        response.raise_for_status()
        suffix = ".pdf" if ".pdf" in url else ".docx"
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_file.write(response.content)
            return temp_file.name
    except requests.RequestException as e:
        print(f"Error downloading file from {url}: {e}")
        raise

def load_docx_documents(file_path: str) -> List[Document]:
    doc = DocxDocument(file_path)
    full_text = [para.text for para in doc.paragraphs if para.text.strip()]
    content = "\n".join(full_text)
    return [Document(page_content=content)]

def process_documents(document_url: str) -> PineconeVectorStore:
    print("Starting data ingestion pipeline...")
    temp_file_path = download_file(document_url)
    
    if temp_file_path.endswith(".pdf"):
        loader = PyPDFLoader(temp_file_path)
        docs = loader.load()
    elif temp_file_path.endswith(".docx"):
        docs = load_docx_documents(temp_file_path)
    else:
        raise ValueError("Unsupported document format provided.")
    
    print(f"Loaded {len(docs)} document pages.")

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunked_docs = text_splitter.split_documents(docs)
    print(f"Split document into {len(chunked_docs)} chunks.")

    if PINECONE_INDEX_NAME not in pc.list_indexes().names():
        print(f"Creating new Pinecone index: {PINECONE_INDEX_NAME}")
        pc.create_index(
            name=PINECONE_INDEX_NAME,
            dimension=384,
            metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="us-east-1")
        )

    print("Upserting document chunks to Pinecone...")
    vector_store = PineconeVectorStore.from_documents(
        chunked_docs, embeddings_model, index_name=PINECONE_INDEX_NAME
    )
    print("Document processing complete.")
    
    os.remove(temp_file_path)
    return vector_store

def create_retriever_tool(vector_store: PineconeVectorStore) -> Tool:
    retriever = vector_store.as_retriever(search_kwargs={"k": 5})
    def retrieve_context(query: str) -> str:
        retrieved_docs = retriever.invoke(query)
        return "\n\n---\n\n".join([doc.page_content for doc in retrieved_docs])
    return Tool(
        name="document_retriever",
        func=retrieve_context,
        description="Use this tool to retrieve relevant context from the policy document to answer a user's question."
    )