import os
import json
from PyPDF2 import PdfReader  # FIXED: Correct case (was "PyPDF2")
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_groq import ChatGroq
from langchain.chains.retrieval_qa.base import RetrievalQA  # FIXED: Updated import path

# Setup working directory
working_dir = os.path.dirname(os.path.abspath(__file__))

os.environ["GROQ_API_KEY"] = os.getenv('GROQ_API_KEY')

# Load embedding model
embedding = HuggingFaceEmbeddings()

# Initialize LLM from Groq
llm = ChatGroq(
    model="deepseek-r1-distill-llama-70b",
    temperature=0
)


def extract_text_from_pdf(file_path):
    """
    Extract text content from a PDF file using PyPDF2.
    """
    try:
        reader = PdfReader(file_path)
        text = ""
        for page in reader.pages:
            extracted = page.extract_text()
            if extracted:
                text += extracted + "\n"
        if not text.strip():
            raise ValueError(f"⚠️ No text extracted from '{os.path.basename(file_path)}'. The file might be empty or image-based.")
        return text
    except Exception as e:
        raise RuntimeError(f"⚠️ Error extracting text from PDF: {e}")


def process_document_to_chroma_db(directory_path):
    """
    Process all PDF documents in the given directory, split their text, 
    and store embeddings in a persistent ChromaDB.
    """
    try:
        all_texts = []
        # Iterate through all PDF files in the directory
        for file_name in os.listdir(directory_path):
            if file_name.endswith(".pdf"):
                file_path = os.path.join(directory_path, file_name)
                print(f"📂 Processing document: {file_name}")

                # Extract text from the PDF
                text = extract_text_from_pdf(file_path)

                # Split text into chunks
                print("🔄 Splitting document into smaller chunks...")
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=2000,
                    chunk_overlap=200
                )
                texts = text_splitter.split_text(text)
                all_texts.extend(texts)

        if not all_texts:
            raise ValueError("⚠️ No PDF files with extractable text found in the directory.")

        # Create a persistent ChromaDB instance with all texts
        print("💾 Storing embeddings in ChromaDB...")
        vectordb = Chroma.from_texts(
            texts=all_texts,
            embedding=embedding,
            persist_directory=os.path.join(working_dir, "doc_vectorstore")
        )
        # No need to call persist() explicitly with Chroma 0.4+

        print("✅ All documents successfully processed and stored in ChromaDB.")
        return "✅ Documents successfully processed and stored in ChromaDB."

    except Exception as e:
        raise RuntimeError(f"⚠️ Error processing documents: {e}")


def answer_question(user_question):
    """
    Retrieve and generate an answer for the given user question 
    based on the stored document embeddings.
    """
    try:
        # Load the persistent vector database
        vectordb_path = os.path.join(working_dir, "doc_vectorstore")
        if not os.path.exists(vectordb_path):
            raise FileNotFoundError("⚠️ ChromaDB vector store not found. Please process a document first.")

        print("📂 Loading vector database...")
        vectordb = Chroma(
            persist_directory=vectordb_path,
            embedding_function=embedding
        )

        # Create a retriever from the vector database
        retriever = vectordb.as_retriever(search_kwargs={"k": 3})

        # Create a QA chain with DeepSeek-R1
        print("🤖 Initializing Retrieval QA chain...")
        qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True  # Optional: helps debugging
        )

        # Invoke the QA chain with the user question
        print("💬 Generating answer...")
        response = qa_chain.invoke({"query": user_question})
        answer = response.get("result", "⚠️ No response generated.")

        return answer

    except Exception as e:
        raise RuntimeError(f"⚠️ Error generating response: {e}")