# app/rag.py
from __future__ import annotations

from pathlib import Path
from typing import List, Dict, Any

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
from langchain_core.prompts import ChatPromptTemplate

from ecdallm.app.search_engine import ECDAStore
from ecdallm.app.vector import DEFAULT_EMBED_MODEL

from ecdallm.app.paths import UPLOAD_DIR, ensure_dirs


ensure_dirs()


ALLOWED_EXTS = {".pdf", ".txt", ".docx"}


RAG_PROMPT = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are ECDA CHATBOT.\n\n"
            "Rules:\n"
            "- Answer the user's question in the SAME LANGUAGE as the question.\n"
            "- Use ONLY the information available in the provided CONTEXT.\n"
            "- If the answer cannot be found in the CONTEXT, respond exactly with:\n"
            "\"I don't know!\"\n"
            "- IMPORTANT: Your answer must be MAXIMUM 150 words.\n"
            "- If the user asks for a longer answer, reply that you can only provide up to 250 words and offer to answer in multiple parts.\n\n"
            "- Keep it short: maximum 5 sentences.\n"
            "- Do not add filler like 'Overall' or 'I'm excited...'.\n"
            "When answering using document information, include citations whenever possible like:\n"
            "(source: filename, page N)."
        ),
        ("human", "QUESTION:\n{question}\n\nCONTEXT:\n{context}\n\nAnswer:"),
    ]
)


def _splitter() -> RecursiveCharacterTextSplitter:
    return RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=80,
        separators=["\n\n", "\n", ".", " ", ""],
    )



def _display_name_from_stored(stored: str) -> str:
    p = Path(stored)
    if "__" in p.stem:
        stem = p.stem.split("__", 1)[0]
        return f"{stem}{p.suffix}"
    return stored


def _load_all_documents() -> List[Document]:
    docs: List[Document] = []

    for p in sorted(UPLOAD_DIR.glob("*")):
        if not p.is_file():
            continue

        ext = p.suffix.lower()
        if ext not in ALLOWED_EXTS:
            continue

        try:
            if ext == ".pdf":
                loader = PyPDFLoader(str(p))
                loaded = loader.load()  # includes page metadata
            elif ext == ".txt":
                loader = TextLoader(str(p), encoding="utf-8", autodetect_encoding=True)
                loaded = loader.load()
            elif ext == ".docx":
                loader = Docx2txtLoader(str(p))
                loaded = loader.load()
            else:
                loaded = []
        except Exception:
            loaded = []

        for d in loaded:
            src = d.metadata.get("source", "")
            d.metadata["source_file"] = Path(src).name if src else p.name
            d.metadata["display_file"] = _display_name_from_stored(Path(d.metadata["source_file"]).name)

        docs.extend(loaded)

    return docs


class LangChainRAG:
    def __init__(self):
        # Use ECDAStore (Chroma + FastEmbed persistence layer)
        self.engine = ECDAStore()
        self.splitter = _splitter()

    def load(self, *, embedding_model: str | None = None) -> None:
        self.engine.load(embedding_model=embedding_model or DEFAULT_EMBED_MODEL)


    def index_new_files(self, *, embedding_model: str | None = None) -> Dict[str, Any]:
        raw_docs = _load_all_documents()
        chunks = self.splitter.split_documents(raw_docs)
        info = self.engine.add_documents(chunks, embedding_model=embedding_model or DEFAULT_EMBED_MODEL)
        files = len({c.metadata.get("display_file", "") for c in chunks}) if chunks else 0
        return {"ok": True, "files": files, **info}


    def reindex(self, *, embedding_model: str | None = None) -> Dict[str, Any]:
        raw_docs = _load_all_documents()
        chunks = self.splitter.split_documents(raw_docs)

        info = self.engine.rebuild(chunks, embedding_model=embedding_model or DEFAULT_EMBED_MODEL)

        files = len({c.metadata.get("display_file", "") for c in chunks}) if chunks else 0
        return {"ok": True, "files": files, "chunks": info.get("chunks", 0)}

    def retrieve(self, query: str, k: int = 5) -> List[Document]:
        return self.engine.search(query, k=k)

    def build_prompt_messages(self, *, question: str, docs: List[Document]) -> List[dict]:
        blocks = []
        for d in docs:
            display_file = d.metadata.get("display_file") or d.metadata.get("source_file") or "unknown"
            page = d.metadata.get("page", None)
            cite = f"{display_file}" + (f", page {page}" if page is not None else "")
            blocks.append(f"[source: {cite}]\n{d.page_content}")

        context = "\n\n".join(blocks) if blocks else "NO CONTEXT PROVIDED."
        msgs = RAG_PROMPT.format_messages(question=question, context=context)

        role_map = {"human": "user", "ai": "assistant", "system": "system"}
        return [{"role": role_map.get(m.type, m.type), "content": m.content} for m in msgs]


rag = LangChainRAG()
