from __future__ import annotations

import shutil
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

from langchain_core.documents import Document
from langchain_core.chat_history import InMemoryChatMessageHistory

from src.chunking import get_chunker
from src.config import settings
from src.embedding_model import get_embeddings
from src.logger import logger
from src.prompt import RAG_SYSTEM_MESSAGE
from src.guardrails import run_input_guardrails, run_output_guardrails
from src.indexing import build_index
from src.llm_client import get_openai_response_with_trace


class SessionState:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.mode: str = "none"  # none | direct | rag
        self.filename: Optional[str] = None
        self.file_size: int = 0
        self.direct_text: Optional[str] = None
        self.vectorstore = None
        self.collection_name: Optional[str] = None
        self.chunk_count: int = 0
        self.persist_dir: Optional[Path] = None
        self.pages: int = 0
        self.chunks: List[Document] = []
        self.memory = InMemoryChatMessageHistory()

        self.llm_provider: str = "openai"
        self.llm_model: str = getattr(settings, "model_name", "gpt-4o-mini")
        self.api_key: Optional[str] = None
        self.api_base_url: Optional[str] = None
        self.llm_client = None

        self.embedding_model_name: Optional[str] = None
        self.vector_db_type_used: Optional[str] = None
        self.chunk_strategy_used: Optional[str] = None

        self.enable_guardrails: bool = True
        self.guardrail_injection: bool = True
        self.guardrail_pii: bool = True
        self.guardrail_numeric: bool = True

class RAGEngine:
    def __init__(self):
        self.sessions: Dict[str, SessionState] = {}
        logger.info("Initializing RAG engine")

        try:
            self._embeddings = get_embeddings()
            logger.info(
                "Embeddings ready",
                context={"model": settings.embedding_model},
            )
        except Exception as e:
            logger.error("Failed to initialize embeddings", context={"error": str(e)})
            self._embeddings = None

        # No global LLM client required at startup (API keys come per session)
        self._client = None

    # ------------------------------------------------------------------
    # Session management
    # ------------------------------------------------------------------

    def create_session(self) -> str:
        session_id = str(uuid.uuid4())
        self.sessions[session_id] = SessionState(session_id)
        logger.info("Session created", context={"session_id": session_id})
        return session_id

    # Replace the entire get_session method in RAGEngine class:
    def get_session(self, session_id: str) -> Optional[SessionState]:
        state = self.sessions.get(session_id)
        if state is not None:
            return state

        import pickle
        from langchain_community.vectorstores import FAISS

        # Try to recover from FAISS first (since that's what you're using)
        target = Path(settings.faiss_persist_directory) / session_id
        chunks_path = target / "chunks.pkl"
        meta_path = target / "session_meta.pkl"
        
        if chunks_path.exists():
            state = SessionState(session_id)
            try:
                with open(chunks_path, "rb") as f:
                    state.chunks = pickle.load(f)
            except Exception as e:
                logger.error("Failed to load chunks", context={"error": str(e)})
                return None
                
            meta = {}
            if meta_path.exists():
                try:
                    with open(meta_path, "rb") as f:
                        meta = pickle.load(f)
                except Exception as e:
                    logger.error("Failed to load metadata", context={"error": str(e)})
            
            state.mode = meta.get("mode", "rag")
            state.filename = meta.get("filename")
            state.vector_db_type_used = meta.get("vector_db_type", "faiss")
            state.embedding_model_name = meta.get("embedding_model")
            state.chunk_count = len(state.chunks)
            state.persist_dir = target

            # Restore vectorstore
            try:
                emb = get_embeddings(state.embedding_model_name or settings.embedding_model)
                if state.vector_db_type_used == "faiss":
                    state.vectorstore = FAISS.load_local(
                        str(target),
                        emb,
                        allow_dangerous_deserialization=True,
                    )
            except Exception as e:
                logger.error("Failed to reload vectorstore", context={"error": str(e)})
                state.vectorstore = None

            self.sessions[session_id] = state
            logger.info("Recovered session from disk", context={"session_id": session_id})
            return state

        # Also check for Chroma recovery (fallback)
        target_chroma = Path(settings.chroma_persist_directory) / session_id
        chunks_path_chroma = target_chroma / "chunks.pkl"
        meta_path_chroma = target_chroma / "session_meta.pkl"
        
        if chunks_path_chroma.exists():
            state = SessionState(session_id)
            try:
                with open(chunks_path_chroma, "rb") as f:
                    state.chunks = pickle.load(f)
            except Exception as e:
                logger.error("Failed to load chunks from chroma", context={"error": str(e)})
                return None
                
            meta = {}
            if meta_path_chroma.exists():
                try:
                    with open(meta_path_chroma, "rb") as f:
                        meta = pickle.load(f)
                except Exception as e:
                    logger.error("Failed to load metadata from chroma", context={"error": str(e)})
            
            state.mode = meta.get("mode", "rag")
            state.filename = meta.get("filename")
            state.vector_db_type_used = meta.get("vector_db_type", "chroma")
            state.embedding_model_name = meta.get("embedding_model")
            state.chunk_count = len(state.chunks)
            state.persist_dir = target_chroma

            self.sessions[session_id] = state
            logger.info("Recovered session from disk (chroma)", context={"session_id": session_id})
            return state

        return None



    def delete_session(self, session_id: str) -> bool:
        state = self.sessions.pop(session_id, None)
        if state is None:
            return False

        if state.persist_dir and state.persist_dir.exists():
            try:
                shutil.rmtree(state.persist_dir, ignore_errors=True)
                logger.info("Cleaned persist dir", context={"path": str(state.persist_dir)})
            except Exception as e:
                logger.error("Failed to clean persist dir", context={"error": str(e)})

        logger.info("Session deleted", context={"session_id": session_id})
        return True

    def list_sessions(self) -> List[str]:
        return list(self.sessions.keys())

    # ------------------------------------------------------------------
    # Ingestion
    # ------------------------------------------------------------------

    def ingest_pdf(
        self,
        session_id: str,
        file_path: Path,
        original_filename: str,
        chunk_strategy: Optional[str] = None,
        embedding_model: Optional[str] = None,
        vector_db_type: Optional[str] = None,
        llm_provider: Optional[str] = None,
        llm_model: Optional[str] = None,
        api_key: Optional[str] = None,
        api_base_url: Optional[str] = None,
        enable_guardrails: Optional[bool] = None,
        guardrail_injection: Optional[bool] = None,
        guardrail_pii: Optional[bool] = None,
        guardrail_numeric: Optional[bool] = None,
    ) -> Dict[str, Any]:
        import pickle
        from src.llm_client import get_client_for_provider

        state = self.get_session(session_id)
        if state is None:
            raise ValueError(f"Unknown session: {session_id}")

        if llm_provider:
            state.llm_provider = llm_provider.lower()
        if llm_model:
            state.llm_model = llm_model
        if api_key:
            state.api_key = api_key
        if api_base_url is not None:
            state.api_base_url = api_base_url.strip() or None
        if embedding_model:
            state.embedding_model_name = embedding_model
        if enable_guardrails is not None:
            state.enable_guardrails = enable_guardrails
        if guardrail_injection is not None:
            state.guardrail_injection = guardrail_injection
        if guardrail_pii is not None:
            state.guardrail_pii = guardrail_pii
        if guardrail_numeric is not None:
            state.guardrail_numeric = guardrail_numeric

        state.llm_client = get_client_for_provider(
            state.llm_provider,
            state.api_key,
            base_url=state.api_base_url,
        )

        emb_name = embedding_model or state.embedding_model_name or settings.embedding_model
        emb = get_embeddings(emb_name)

        steps: List[str] = []
        steps.append(f"LLM: {state.llm_provider} / {state.llm_model}")
        steps.append(f"Embedding model: {emb_name}")

        file_size = file_path.stat().st_size
        state.filename = original_filename
        state.file_size = file_size

        # ------------------------------------------------------------------
        # 1. Load PDF
        # ------------------------------------------------------------------
        if getattr(settings, "enable_multimodal", False):
            from src.pdf_processor import load_pdf_multimodal
            from src.llm_client import describe_image_with_vision

            pages: List[Document] = load_pdf_multimodal(
                file_path,
                vision_fn=describe_image_with_vision,
            )
            steps.append("Loaded PDF with multimodal processor (text + tables + vision)")
        else:
            from langchain_community.document_loaders import PyPDFLoader

            pages = PyPDFLoader(str(file_path)).load()
            steps.append("Loaded PDF with PyPDFLoader")

        if not pages:
            raise ValueError("PDF is empty or unreadable")

        state.pages = len(pages)
        full_text = "\n\n".join(p.page_content for p in pages)
        steps.append(f"Loaded PDF ({len(pages)} pages)")

        # ------------------------------------------------------------------
        # 2. Direct context for small files
        # ------------------------------------------------------------------
        if file_size <= settings.direct_context_threshold_bytes:
            state.mode = "direct"
            state.direct_text = full_text
            state.chunk_count = 0
            state.chunks = []
            state.vectorstore = None
            state.collection_name = None
            state.persist_dir = None
            steps.append("Small file → direct context mode (no indexing)")

            # Persist minimal meta so ZeroGPU chat worker can recover
            meta_dir = Path(settings.upload_path) / session_id
            meta_dir.mkdir(parents=True, exist_ok=True)
            with open(meta_dir / "session_meta.pkl", "wb") as f:
                pickle.dump(
                    {
                        "mode": "direct",
                        "filename": original_filename,
                        "direct_text": full_text,
                        "embedding_model": emb_name,
                    },
                    f,
                )

            return {
                "session_id": session_id,
                "mode": "direct",
                "filename": original_filename,
                "file_size_bytes": file_size,
                "pages": len(pages),
                "chunks": 0,
                "message": (
                    f"Small file processed in direct context mode "
                    f"(≤ {settings.direct_context_threshold_kb} KB)."
                ),
                "pipeline_steps": steps,
                "chunk_strategy_used": None,
                "embedding_model_used": emb_name,
                "vector_db_type_used": None,
            }

        # ------------------------------------------------------------------
        # 3. Chunk
        # ------------------------------------------------------------------
        state.mode = "rag"
        strategy = (chunk_strategy or settings.chunk_strategy).lower()
        state.chunk_strategy_used = strategy
        steps.append(f"Chunk strategy: {strategy}")

        splitter = get_chunker(strategy=strategy)
        chunks = splitter.split_documents(pages)
        state.chunk_count = len(chunks)
        state.chunks = chunks
        steps.append(f"Created {len(chunks)} chunks")

        # ------------------------------------------------------------------
        # 4. Index
        # ------------------------------------------------------------------
        index_type = (vector_db_type or settings.vector_db_type).lower()
        state.vector_db_type_used = index_type
        steps.append(f"Vector DB: {index_type}")

        vector_store = build_index(
            documents=chunks,
            embedding_model=emb,
            session_id=session_id,
            filename=original_filename,
            index_type=index_type,
        )
        state.vectorstore = vector_store
        steps.append(f"Indexed into {index_type}")

        if index_type == "chroma":
            state.persist_dir = Path(settings.chroma_persist_directory) / session_id
            state.collection_name = f"session_{session_id}"
        elif index_type == "faiss":
            state.persist_dir = Path(settings.faiss_persist_directory) / session_id
            state.collection_name = f"session_{session_id}"
        else:
            state.persist_dir = Path(settings.upload_path) / session_id
            state.collection_name = None

        # ------------------------------------------------------------------
        # 5. Persist chunks + meta for ZeroGPU worker recovery  ← STEP 2
        # ------------------------------------------------------------------
        meta_dir = state.persist_dir or (Path(settings.upload_path) / session_id)
        meta_dir.mkdir(parents=True, exist_ok=True)

        with open(meta_dir / "chunks.pkl", "wb") as f:
            pickle.dump(chunks, f)

        with open(meta_dir / "session_meta.pkl", "wb") as f:
            pickle.dump(
                {
                    "mode": state.mode,
                    "filename": original_filename,
                    "vector_db_type": index_type,
                    "embedding_model": emb_name,
                    "chunk_strategy": strategy,
                    "pages": state.pages,
                },
                f,
            )
        steps.append(f"Persisted session artifacts to {meta_dir}")

        return {
            "session_id": session_id,
            "mode": "rag",
            "filename": original_filename,
            "file_size_bytes": file_size,
            "pages": len(pages),
            "chunks": len(chunks),
            "message": "PDF ingested into vector store.",
            "pipeline_steps": steps,
            "chunk_strategy_used": strategy,
            "embedding_model_used": emb_name,
            "vector_db_type_used": index_type,
        }

    def process_uploaded_file(
        self,
        session_id: str,
        file_path: Path,
        original_filename: Optional[str] = None,
        chunk_strategy: Optional[str] = None,
        embedding_model: Optional[str] = None,
        vector_db_type: Optional[str] = None,
        llm_provider: Optional[str] = None,
        llm_model: Optional[str] = None,
        api_key: Optional[str] = None,
        api_base_url: Optional[str] = None,
        enable_guardrails: Optional[bool] = None,
        guardrail_injection: Optional[bool] = None,
        guardrail_pii: Optional[bool] = None,
        guardrail_numeric: Optional[bool] = None,
    ) -> Dict[str, Any]:
        return self.ingest_pdf(
            session_id,
            file_path,
            original_filename or file_path.name,
            chunk_strategy=chunk_strategy,
            embedding_model=embedding_model,
            vector_db_type=vector_db_type,
            llm_provider=llm_provider,
            llm_model=llm_model,
            api_key=api_key,
            api_base_url=api_base_url,
            enable_guardrails=enable_guardrails,
            guardrail_injection=guardrail_injection,
            guardrail_pii=guardrail_pii,
            guardrail_numeric=guardrail_numeric,
        )

    # ------------------------------------------------------------------
    # Query
    # ------------------------------------------------------------------

    def query(
        self,
        session_id: str,
        question: str,
        k: Optional[int] = None,
        retrieval_type: Optional[str] = None,
        llm_provider: Optional[str] = None,
        llm_model: Optional[str] = None,
        api_key: Optional[str] = None,
        api_base_url: Optional[str] = None,
    ) -> Dict[str, Any]:
        from src.llm_client import get_client_for_provider

        state = self.get_session(session_id)
        if state is None:
            raise ValueError("Session not found")
        if state.mode == "none":
            raise ValueError("No document uploaded for this session")

        # ----- Per-request overrides (required on ZeroGPU: keys don't survive workers) -----
        if llm_provider:
            state.llm_provider = llm_provider.lower()
        if llm_model:
            state.llm_model = llm_model
        if api_key:
            state.api_key = api_key
        if api_base_url is not None:
            state.api_base_url = (api_base_url or "").strip() or None

        steps: List[str] = []

        client = get_client_for_provider(
            state.llm_provider,
            state.api_key,
            base_url=state.api_base_url,
        )
        state.llm_client = client

        model = state.llm_model or settings.model_name
        rtype = (
            retrieval_type
            or getattr(settings, "retrieval_type", None)
            or "semantic"
        ).lower()
        k = k or settings.num_retrieved_chunks

        steps.append(f"LLM: {state.llm_provider} / {model}")
        steps.append(f"Retrieval type: {rtype}")
        steps.append(f"Top-k: {k}")

        # ------------------------------------------------------------------
        # 1. Input guardrails
        # ------------------------------------------------------------------
        use_guardrails = getattr(state, "enable_guardrails", True)
        if use_guardrails:
            steps.append("Input guardrails: enabled")
            input_guard = run_input_guardrails(question)
            
            if not getattr(state, "guardrail_injection", True) and input_guard.injection_detected:
                input_guard.injection_detected = False
                input_guard.overall_passed = True
                input_guard.status = "passed"
                input_guard.details = [d for d in input_guard.details if "Injection" not in d]
                input_guard.details.append("Injection check: skipped by user")

            if not getattr(state, "guardrail_pii", True) and input_guard.pii_detected:
                input_guard.pii_detected = False
                input_guard.pii_entities = []
                input_guard.redacted_query = None
                input_guard.warnings = [w for w in input_guard.warnings if "PII" not in w]
                input_guard.details.append("PII check: skipped by user")
                if input_guard.overall_passed and not input_guard.warnings:
                    input_guard.status = "passed"
            
            if not input_guard.overall_passed:
                steps.append("Input guardrails: BLOCKED")
                return {
                    "mode": "blocked",
                    "answer": (
                        "Your question was blocked by safety guardrails.\n"
                        + "\n".join(input_guard.details)
                    ),
                    "sources": [],
                    "guardrails": {
                        "input": input_guard.to_dict(),
                        "output": None,
                        "status": "blocked",
                        "warnings": input_guard.warnings,
                        "ungrounded_numbers": [],
                    },
                    "reasoning_trace": None,
                    "pipeline_steps": steps,
                    "retrieval_type_used": rtype,
                    "history_turns": len(state.memory.messages) // 2,
                }
            safe_question = input_guard.redacted_query or question
            if input_guard.warnings:
                steps.append(f"Input guardrails: warned ({len(input_guard.warnings)})")
            else:
                steps.append("Input guardrails: passed")
        else:
            steps.append("Input guardrails: skipped")

            class _Pass:
                overall_passed = True
                redacted_query = question
                warnings: List[str] = []
                details: List[str] = []
                ungrounded_numbers: List[str] = []

                def to_dict(self):
                    return {
                        "overall_passed": True,
                        "injection_detected": False,
                        "pii_detected": False,
                        "pii_entities": [],
                        "warnings": [],
                        "details": ["Guardrails disabled for this session"],
                    }

            input_guard = _Pass()
            safe_question = question

        # ------------------------------------------------------------------
        # 2. Retrieve context
        # ------------------------------------------------------------------
        from src.retrieval import get_retriever, retrieve_documents

        if state.mode == "direct":
            context = state.direct_text or ""
            context_texts = [context] if context else []
            sources: List[Dict[str, Any]] = []
            steps.append("Mode: direct context (no retrieval)")
        else:
            if state.vectorstore is None and rtype in ("semantic", "hybrid"):
                if rtype != "bm25":
                    raise ValueError(
                        "Vector store not available for semantic/hybrid retrieval"
                    )
            if not state.chunks:
                raise ValueError("Chunks not available")

            retriever = get_retriever(
                vectorstore=state.vectorstore,
                documents=state.chunks,
                k=k,
                retrieval_type=rtype,
            )
            docs = retrieve_documents(retriever, safe_question, k=k)
            context_texts = [doc.page_content for doc in docs]
            context = "\n\n".join(
                f"[Chunk {i + 1}]\n{doc.page_content}" for i, doc in enumerate(docs)
            )
            sources = [
                {
                    "content": doc.page_content[:300]
                    + ("..." if len(doc.page_content) > 300 else ""),
                    "metadata": doc.metadata,
                }
                for doc in docs
            ]
            steps.append(f"Retrieved {len(docs)} chunks")

        # ------------------------------------------------------------------
        # 3. Prompt
        # ------------------------------------------------------------------
        history_text = ""
        if state.memory.messages:
            history_text = "Previous conversation:\n"
            for msg in state.memory.messages:
                role = "User" if msg.type == "human" else "Assistant"
                history_text += f"{role}: {msg.content}\n"
            history_text += "\n"
            steps.append(f"History turns: {len(state.memory.messages) // 2}")

        full_input = (
            f"{history_text}"
            f"<Context>\n{context}\n</Context>\n"
            f"<Question>{safe_question}</Question>"
        )

        # ------------------------------------------------------------------
        # 4. LLM
        # ------------------------------------------------------------------
        steps.append("Calling LLM…")
        answer, reasoning_trace = get_openai_response_with_trace(
            input=full_input,
            instructions=RAG_SYSTEM_MESSAGE,
            client=client,
            model=model,
            provider=state.llm_provider,
            reasoning_effort=getattr(settings, "reasoning_effort", None),
        )
        steps.append("LLM response received")

        # ------------------------------------------------------------------
        # 5. Output guardrails
        # ------------------------------------------------------------------
        if use_guardrails and getattr(state, "guardrail_numeric", True):
            steps.append("Output guardrails: numeric grounding")
            output_guard = run_output_guardrails(answer, context_texts)
            if output_guard.warnings:
                steps.append(f"Output guardrails: warned ({len(output_guard.warnings)})")
            else:
                steps.append("Output guardrails: passed")
        else:
            steps.append("Output guardrails: skipped")

            class _OutPass:
                warnings: List[str] = []
                ungrounded_numbers: List[str] = []
                details: List[str] = []

                def to_dict(self):
                    return {
                        "numeric_grounding_passed": True,
                        "ungrounded_numbers": [],
                        "warnings": [],
                        "details": ["Output guardrails disabled"],
                    }

            output_guard = _OutPass()

        # ------------------------------------------------------------------
        # 6. Memory
        # ------------------------------------------------------------------
        state.memory.add_user_message(question)
        state.memory.add_ai_message(answer)
        max_messages = settings.max_history_turns * 2
        if len(state.memory.messages) > max_messages:
            state.memory.messages = state.memory.messages[-max_messages:]

        combined = {
            "input": input_guard.to_dict(),
            "output": output_guard.to_dict(),
            "status": (
                "blocked"
                if not getattr(input_guard, "overall_passed", True)
                else (
                    "warned"
                    if (input_guard.warnings or output_guard.warnings)
                    else "passed"
                )
            ),
            "warnings": list(input_guard.warnings) + list(output_guard.warnings),
            "ungrounded_numbers": list(
                getattr(output_guard, "ungrounded_numbers", []) or []
            ),
        }

        return {
            "mode": state.mode,
            "answer": answer,
            "sources": sources,
            "guardrails": combined,
            "reasoning_trace": reasoning_trace or None,
            "pipeline_steps": steps,
            "retrieval_type_used": rtype,
            "history_turns": len(state.memory.messages) // 2,
        }
    def get_context(self, session_id: str) -> str:
        state = self.get_session(session_id)
        if state is None:
            raise ValueError("Session not found")
        if state.mode == "direct":
            return state.direct_text or ""
        return ""

    def get_status(self, session_id: str) -> Dict[str, Any]:
        state = self.get_session(session_id)
        if state is None:
            raise ValueError("Session not found")
        return {
            "session_id": session_id,
            "mode": state.mode,
            "filename": state.filename,
            "file_size": state.file_size,
            "chunk_count": state.chunk_count,
            "pages": state.pages,
        }

    def clear_history(self, session_id: str) -> bool:
        state = self.get_session(session_id)
        if state is None:
            return False
        state.memory.clear()
        logger.info("History cleared", context={"session_id": session_id})
        return True


rag_engine = RAGEngine()