import os, logging
import glob
import logging
import re
import hashlib
import numpy as np
import asyncio
from typing import List, Optional, Any
import shutil
import secrets

# --- Pydantic & FastAPI ---
from pydantic import BaseModel, Field
import uvicorn
from fastapi import FastAPI, HTTPException, APIRouter, UploadFile, File, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.security import HTTPBasic, HTTPBasicCredentials

# --- LangChain & RAG Core ---
from PyPDF2 import PdfReader

try:
    from docx import Document as DocxDocument
except ImportError:
    DocxDocument = None
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.retrievers import EnsembleRetriever
from langchain.chains.combine_documents import (
    create_stuff_documents_chain,
)  # <-- DEPRECATION FIX
from langchain.chains import create_retrieval_chain
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.llms.base import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain.docstore.document import Document
from langchain_core.retrievers import BaseRetriever
from langchain.prompts import PromptTemplate

# --- Scikit-Learn for TF-IDF ---
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# --- TruLens for Evaluation ---
# from trulens_eval import Feedback, Tru, TruChain
# from trulens_eval.feedback.provider import LiteLLM

# tru = Tru()

# --- Mistral AI ---
from mistralai import Mistral as MistralClient
from dotenv import load_dotenv

load_dotenv()

# --- Global Variables & Configuration ---
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Define the original, read-only data folder from the repo
REPO_DATA_FOLDER = "data"

# Use persistent storage if available (on Hugging Face Spaces)
if os.path.exists("/data"):
    logger.info("Persistent storage '/data' found. Using it for PDFs and cache.")
    DATA_FOLDER = "/data/pdf_storage"
    CACHE_DIR = "/data/cache"
else:
    logger.info(
        "No persistent storage '/data' found. Using local 'data' folder and '/tmp/cache'."
    )
    DATA_FOLDER = REPO_DATA_FOLDER
    CACHE_DIR = "/tmp/cache"

os.makedirs(DATA_FOLDER, exist_ok=True)
os.makedirs(CACHE_DIR, exist_ok=True)
retrieval_chain = None

HF_CACHE = "/tmp/hf"
os.environ["HF_HOME"] = HF_CACHE
os.environ["HF_HOME"] = HF_CACHE
os.environ["HF_DATASETS_CACHE"] = HF_CACHE
os.environ["HF_MODULES_CACHE"] = HF_CACHE

retrieval_chain = None

db_path = "/data/default.sqlite" if os.path.exists("/data") else "default.sqlite"

# ✅ Only run the dashboard locally (not in Hugging Face Spaces)
if os.environ.get("SPACE_ID") is None:
    # tru.run_dashboard()
    pass


# --- Pydantic Models ---
class QuestionRequest(BaseModel):
    question: str = Field(..., min_length=1)


class AnswerResponse(BaseModel):
    answer: str
    sources: List[str]


# --- Greeting Handler ---
GREETING_RESPONSES = {
    "hi": "Hello! How can I assist you?",
    "hello": "Hello! How can I help you today?",
    "bye": "Bye-bye! See you soon.",
    "ty": "Your Welcome!",
}

# --- Enhanced Prompt ---
UNIVERSITY_ASSISTANT_PROMPT = PromptTemplate.from_template(
    """
You are a meticulous Guru Nanak Dev University Academic Assistant.
Your task is to provide accurate, specific, and concise answers based ONLY on the provided Context.
**Instructions:**
1.  **Analyze the Context:** Carefully read the user's question and the provided context to synthesize an answer.
2.  **Extract Specifics:** If asked for a fee, number, or date, provide the EXACT information.
3.  **Cite Sources:** After your answer, cite the source document and page number: `(Source: Prospectus.pdf, Page: 15)`.
4.  **Handle Missing Information:** If the info is not in the context, state: "Based on the provided documents, I could not find specific information about this topic."
---
**Context:**
{context}
---
**Question:**
{input}
**Answer (with citations):**
"""
)


# --- Custom Scikit-Learn TF-IDF Retriever ---
class ScikitLearnTfidfRetriever(BaseRetriever):
    docs: List[Document]
    tfidf_vectorizer: TfidfVectorizer
    tfidf_matrix: Any
    k: int = 15

    class Config:
        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str) -> List[Document]:
        query_vector = self.tfidf_vectorizer.transform([query])
        similarity_scores = cosine_similarity(query_vector, self.tfidf_matrix).flatten()
        top_k_indices = np.argsort(similarity_scores)[-self.k :][::-1]
        return [self.docs[i] for i in top_k_indices]


# --- MistralLLM Class ---
class MistralLLM(LLM):
    client: Any = None
    model: str = "mistral-small-latest"
    temperature: float = 0.0
    api_key: str = ""

    def __init__(
        self,
        api_key: str,
        model: str = "mistral-small-latest",
        temperature: float = 0.0,
    ):
        super().__init__()
        self.api_key = api_key
        self.model = model
        self.temperature = temperature
        self.client = MistralClient(api_key=self.api_key)

    @property
    def _llm_type(self) -> str:
        return "mistral"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        messages = [{"role": "user", "content": prompt}]
        response = self.client.chat.complete(
            model=self.model, messages=messages, temperature=self.temperature
        )
        return response.choices[0].message.content.strip()


# --- Caching and Data Processing Functions ---
def get_data_hash(folder_path=DATA_FOLDER) -> str:
    """Generates a hash based on the filenames and modification times of supported files."""
    hasher = hashlib.md5()
    if not os.path.exists(folder_path):
        return "no_data"

    supported_extensions = [".pdf", ".txt", ".md", ".docx"]
    files_to_hash = []
    for ext in supported_extensions:
        files_to_hash.extend(glob.glob(os.path.join(folder_path, f"*{ext}")))

    for file_path in sorted(files_to_hash):
        try:
            stat = os.stat(file_path)
            file_info = f"{os.path.basename(file_path)}-{stat.st_mtime}-{stat.st_size}"
            hasher.update(file_info.encode("utf-8"))
        except FileNotFoundError:
            continue  # File might have been deleted during the process

    return hasher.hexdigest()


def get_documents_from_files(folder_path=DATA_FOLDER) -> List[Document]:
    docs = []
    supported_extensions = [".pdf", ".txt", ".md", ".docx"]

    for file_path in glob.glob(os.path.join(folder_path, "*")):
        file_name = os.path.basename(file_path)
        file_ext = os.path.splitext(file_name)[1].lower()

        if file_ext not in supported_extensions:
            continue

        if file_ext == ".pdf":
            pdf_reader = PdfReader(file_path)
            for i, page in enumerate(pdf_reader.pages):
                text = page.extract_text() or ""
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": i + 1}
                    )
                )
        elif file_ext in [".txt", ".md"]:
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": 1}
                    )
                )
        elif file_ext == ".docx":
            if DocxDocument is None:
                logger.warning(
                    "Skipping .docx file because 'python-docx' is not installed. Please install it with 'pip install python-docx'."
                )
                continue
            try:
                doc = DocxDocument(file_path)
                text = "\n".join([para.text for para in doc.paragraphs])
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": 1}
                    )
                )
            except Exception as e:
                logger.error(f"Error processing .docx file {file_name}: {e}")

    return docs


async def initialize_rag_system():
    global retrieval_chain
    logger.info("🚀 Initializing RAG system...")
    os.makedirs(CACHE_DIR, exist_ok=True)

    data_hash = get_data_hash()
    cache_path = os.path.join(CACHE_DIR, f"vectorstore_{data_hash}.faiss")

    embeddings = HuggingFaceEmbeddings(
        model_name="BAAI/bge-base-en-v1.5",
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )

    if os.path.exists(cache_path):
        logger.info(f"✅ Loading vector store from cache: {cache_path}")
        vectorstore = FAISS.load_local(
            cache_path, embeddings, allow_dangerous_deserialization=True
        )
        documents = get_documents_from_files()
        text_chunks = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=150
        ).split_documents(documents)
    else:
        logger.warning(
            "No cache found or data has changed. Building new vector store..."
        )
        documents = get_documents_from_files()
        if not documents:
            logger.warning(
                "No documents found in /data folder. RAG system will be inactive."
            )
            retrieval_chain = None  # Ensure chain is None
            return

        text_chunks = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=150
        ).split_documents(documents)

        logger.info(
            f"Embedding {len(text_chunks)} chunks... (This will be slow on the first run)"
        )
        vectorstore = FAISS.from_documents(documents=text_chunks, embedding=embeddings)

        logger.info(f"💾 Saving new vector store to cache: {cache_path}")
        vectorstore.save_local(cache_path)

    # --- Chain Setup ---
    mistral_api_key = os.getenv("MISTRAL_API_KEY")
    if not mistral_api_key:
        raise ValueError("MISTRAL_API_KEY not found")

    llm = MistralLLM(api_key=mistral_api_key)

    # --- DEPRECATION FIX: Use create_stuff_documents_chain ---
    question_answer_chain = create_stuff_documents_chain(
        llm, UNIVERSITY_ASSISTANT_PROMPT
    )

    faiss_retriever = vectorstore.as_retriever(
        search_type="similarity", search_kwargs={"k": 15}
    )

    tfidf_vectorizer = TfidfVectorizer(stop_words="english")

    # Ensure text_chunks is not empty before fitting
    if not text_chunks:
        logger.warning("No text chunks to process for TF-IDF.")
        # Handle case with no documents gracefully
        tfidf_retriever = ScikitLearnTfidfRetriever(
            docs=[],
            tfidf_vectorizer=tfidf_vectorizer,  # an empty vectorizer
            tfidf_matrix=tfidf_vectorizer.fit_transform([]),  # an empty matrix
            k=15,
        )
    else:
        tfidf_matrix = tfidf_vectorizer.fit_transform(
            [doc.page_content for doc in text_chunks]
        )
        tfidf_retriever = ScikitLearnTfidfRetriever(
            docs=text_chunks,
            tfidf_vectorizer=tfidf_vectorizer,
            tfidf_matrix=tfidf_matrix,
            k=15,
        )

    ensemble_retriever = EnsembleRetriever(
        retrievers=[tfidf_retriever, faiss_retriever], weights=[0.5, 0.5]
    )

    # Create the final retrieval chain that combines the retriever and the QA chain
    retrieval_chain = create_retrieval_chain(ensemble_retriever, question_answer_chain)

    logger.info("✅ RAG system initialized successfully.")


def extract_sources_from_docs(docs: List[Document]) -> List[str]:
    """Extracts unique source and page numbers from a list of documents."""
    return sorted(
        list(
            set(
                f"({doc.metadata.get('source', 'N/A')}, Page: {doc.metadata.get('page', 'N/A')})"
                for doc in docs
            )
        )
    )


# --- FastAPI Application ---
app = FastAPI(title="GNDU Assistant (Cached)", version="4.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

security = HTTPBasic()


def get_current_username(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = secrets.compare_digest(
        credentials.username, os.environ.get("ADMIN_USERNAME", "admin")
    )
    correct_password = secrets.compare_digest(
        credentials.password, os.environ.get("ADMIN_PASSWORD", "password")
    )
    if not (correct_username and correct_password):
        raise HTTPException(
            status_code=401,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username


@app.on_event("startup")
async def startup_event():
    # On Hugging Face, copy initial files from the read-only repo folder to persistent storage
    if DATA_FOLDER != REPO_DATA_FOLDER and os.path.exists(REPO_DATA_FOLDER):
        supported_extensions = [".pdf", ".txt", ".md", ".docx"]
        # Check if the destination is empty of supported files before copying
        if not any(
            any(f.endswith(ext) for ext in supported_extensions)
            for f in os.listdir(DATA_FOLDER)
        ):
            logger.info(
                f"Copying initial files from '{REPO_DATA_FOLDER}' to '{DATA_FOLDER}'..."
            )
            for filename in os.listdir(REPO_DATA_FOLDER):
                if any(filename.endswith(ext) for ext in supported_extensions):
                    src_path = os.path.join(REPO_DATA_FOLDER, filename)
                    dest_path = os.path.join(DATA_FOLDER, filename)
                    if not os.path.exists(dest_path):
                        shutil.copy(src_path, dest_path)

    await initialize_rag_system()


admin_router = APIRouter()


@admin_router.get("/files")
async def get_files(_: str = Depends(get_current_username)):
    """Lists all supported files in the data directory."""
    if not os.path.exists(DATA_FOLDER):
        return []
    supported_extensions = [".pdf", ".txt", ".md", ".docx"]
    files = sorted(
        [
            f
            for f in os.listdir(DATA_FOLDER)
            if any(f.endswith(ext) for ext in supported_extensions)
        ]
    )
    return files


@admin_router.post("/files")
async def upload_file(
    file: UploadFile = File(...), _: str = Depends(get_current_username)
):
    """Uploads a new PDF file."""
    file_path = os.path.join(DATA_FOLDER, file.filename)
    if os.path.exists(file_path):
        raise HTTPException(
            status_code=409, detail="File with this name already exists."
        )

    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    finally:
        file.file.close()

    await initialize_rag_system()

    return {"filename": file.filename}


@admin_router.delete("/files/{filename}")
async def delete_file(filename: str, _: str = Depends(get_current_username)):
    """Deletes a PDF file."""
    # Basic security check
    if ".." in filename:
        raise HTTPException(status_code=400, detail="Invalid filename.")

    file_path = os.path.join(DATA_FOLDER, filename)

    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found.")

    os.remove(file_path)

    await initialize_rag_system()

    return {"message": f"File '{filename}' deleted successfully."}


app.include_router(admin_router, prefix="/admin")


@app.post("/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest):
    normalized_question = request.question.strip().lower()
    if normalized_question in GREETING_RESPONSES:
        return AnswerResponse(
            answer=GREETING_RESPONSES[normalized_question],
            sources=["Conversational Greeting"],
        )

    if not retrieval_chain:
        raise HTTPException(
            status_code=503,
            detail="RAG system is not initialized. Please add PDF files to the /data folder and restart.",
        )

    try:
        # The new chain structure returns a dictionary with 'answer' and 'context'
        result = await retrieval_chain.ainvoke({"input": request.question})

        answer = result.get("answer", "No answer could be generated.")
        sources = extract_sources_from_docs(result.get("context", []))

        return AnswerResponse(answer=answer, sources=sources)

    except Exception as e:
        logger.error(f"Error during RAG chain invocation: {e}")
        raise HTTPException(status_code=500, detail=str(e))


app.mount("/", StaticFiles(directory="frontend", html=True), name="static")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)