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
from contextlib import asynccontextmanager
from nps_agent import run_nps_task

# --- 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_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_classic.retrievers import EnsembleRetriever
from langchain_classic.chains.combine_documents import (
    create_stuff_documents_chain,
)
from langchain_classic.chains import create_retrieval_chain
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.language_models import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from langchain_core.prompts import PromptTemplate

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

# --- 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_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"


# --- Pydantic Models ---

# --- Pydantic Models ---
class QuestionRequest(BaseModel):
    question: str = Field(..., min_length=1)
    mode: str = Field("chat", description="Mode of operation: 'chat' or 'agent'")
    language: str = Field("English", description="Language of the question and expected answer")


class AgentRequest(BaseModel):
    task: str = Field(..., min_length=5, description="The task for the browser agent to perform")


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]

    async def _aget_relevant_documents(
        self, query: str, *, run_manager=None
    ) -> List[Document]:
        return self._get_relevant_documents(query)


# --- 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)

    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")

    if not text_chunks:
        logger.warning("No text chunks to process for TF-IDF.")
        tfidf_retriever = ScikitLearnTfidfRetriever(
            docs=[],
            tfidf_vectorizer=tfidf_vectorizer,
            tfidf_matrix=tfidf_vectorizer.fit_transform([]),
            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]
    )

    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
            )
        )
    )

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 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"]
        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()
    yield
    # Shutdown event (if needed in future)


# --- FastAPI Application ---
app = FastAPI(title="GNDU Assistant (Cached)", version="4.0.0", lifespan=lifespan)
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", "admin")
    )
    if not (correct_username and correct_password):
        raise HTTPException(
            status_code=401,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username


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."""
    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")


async def translate_text(text: str, target_language: str, api_key: str) -> str:
    """Translates text to the target language using Mistral."""
    if target_language.lower() == "english":
        return text

    client = MistralClient(api_key=api_key)
    prompt = f"Translate the following text to {target_language}. Return ONLY the translated text, no unique prefix/suffix.\n\nText: {text}"
    
    try:
        messages = [{"role": "user", "content": prompt}]
        response = await client.chat.complete_async(
            model="mistral-small-latest", messages=messages, temperature=0.0
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        logger.error(f"Translation failed: {e}")
        return text  # Fallback to original text

@app.post("/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest):
    normalized_question = request.question.strip().lower()
    mistral_api_key = os.getenv("MISTRAL_API_KEY")
    
    # 1. Translate Query to English if needed
    english_question = request.question
    if request.language.lower() != "english":
        try:
             # Use a simple synchronous call for now if async is tricky with the client version, 
             # but let's try to stick to the pattern. 
             # IMPORTANT check: MistralClient in this code might be sync or async. 
             # The existing MistralLLM uses sync `client.chat.complete`.
             # We will use sync here for safety to match MistralLLM usage, or wrap in run_in_executor if needed.
             # For simplicity in this step, let's use the same sync client pattern as MistralLLM.
            
            client = MistralClient(api_key=mistral_api_key)
            trans_prompt = f"Translate the following text to English. Return ONLY the translated text.\n\nText: {request.question}"
            trans_resp = client.chat.complete(
                model="mistral-small-latest", messages=[{"role": "user", "content": trans_prompt}], temperature=0.0
            )
            english_question = trans_resp.choices[0].message.content.strip()
            logger.info(f"Translated '{request.question}' ({request.language}) -> '{english_question}' (English)")
        except Exception as e:
            logger.error(f"Query translation failed: {e}")

    # 2. Process Request
    answer_text = ""
    sources_list = []

    if request.mode == "agent":
        try:
            logger.info(f"🚀 Executing Agent Task: {english_question}")
            answer_text = await run_nps_task(english_question)
            sources_list = ["Autonomus Agent"]
        except Exception as e:
            logger.error(f"Error running agent task: {e}")
            raise HTTPException(status_code=500, detail=f"Agent Error: {str(e)}")

    elif normalized_question in GREETING_RESPONSES:
         answer_text = GREETING_RESPONSES[normalized_question]
         sources_list = ["Conversational Greeting"]
         
    else:
        # RAG Logic
        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:
            # Modify the input to specificy the output language in the prompt if it's RAG
            # Actually, RAG prompt expects 'input' and 'context'. 
            # We will rely on post-translation for consistency across Agent/Chat
            result = await retrieval_chain.ainvoke({"input": english_question})
            answer_text = result.get("answer", "No answer could be generated.")
            sources_list = extract_sources_from_docs(result.get("context", []))

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

    # 3. Translate Answer back to Target Language
    if request.language.lower() != "english":
        try:
            client = MistralClient(api_key=mistral_api_key)
            back_prompt = f"Translate the following text to {request.language}. Return ONLY the translated text.\n\nText: {answer_text}"
            back_resp = client.chat.complete(
                model="mistral-small-latest", messages=[{"role": "user", "content": back_prompt}], temperature=0.0
            )
            answer_text = back_resp.choices[0].message.content.strip()
        except Exception as e:
            logger.error(f"Answer translation failed: {e}")

    return AnswerResponse(answer=answer_text, sources=sources_list)


@app.post("/agent/run")
async def run_agent(request: AgentRequest):
    """
    Executes a task using the NPS Nanobot Agent.
    """
    try:
        logger.info(f"🚀 Received agent task: {request.task}")
        result = await run_nps_task(request.task)
        return {"result": result}
    except Exception as e:
        logger.error(f"Error running agent task: {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)