# app/search_engine.py
from __future__ import annotations

import os
from pathlib import Path
from typing import List

import chromadb
from chromadb.config import Settings

from openai import OpenAI
from langchain_core.documents import Document

try:
    from langchain_chroma import Chroma
except Exception:
    from langchain_community.vectorstores import Chroma 

from .vector import DEFAULT_EMBED_MODEL, make_embeddings

from .paths import RAG_STORE_DIR, ensure_dirs


ensure_dirs()


class ECDAStore:
    """
    Persistent ChromaDB store + local embeddings via FastEmbed.

    Provides:
      - load(embedding_model=...)
      - rebuild(docs, embedding_model=...)
      - search(query, k=...)
      - hard_reset()   (call BEFORE deleting rag_store on disk)
      - recreate()     (call AFTER deleting rag_store on disk)
    """

    def __init__(self):
        # LM Studio / OpenAI-compatible chat client (optional usage elsewhere)
        self.client = OpenAI(
            api_key=os.getenv("API_KEY", "lm-studio"),
            base_url=os.getenv("BASE_URL", "http://localhost:1234/v1"),
        )

        # Paths



        self.rag_dir = RAG_STORE_DIR
        self.chroma_dir = self.rag_dir / "chroma_db"
        self.chroma_dir.mkdir(parents=True, exist_ok=True)

        # Runtime state
        self.embedding_model: str = DEFAULT_EMBED_MODEL
        self.embeddings = make_embeddings(embedding_model=self.embedding_model)

        self._client: chromadb.PersistentClient | None = None
        self.vs: Chroma | None = None

        # initialize
        self._ensure_store()

    # -----------------------------
    # Internal helpers
    # -----------------------------
    def _ensure_store(self) -> None:
        """Ensure chroma client and LC vectorstore exist."""
        if self._client is None:
            self.chroma_dir.mkdir(parents=True, exist_ok=True)
            self._client = chromadb.PersistentClient(
                path=str(self.chroma_dir),
                settings=Settings(anonymized_telemetry=False),
            )

        if self.vs is None:
            self.vs = Chroma(
                client=self._client,
                collection_name=self._collection_name(self.embedding_model),
                embedding_function=self.embeddings,
            )

    def _collection_name(self, embedding_model: str) -> str:
        """
        Make a stable collection name per embedding model to avoid mixing vectors.
        """
        safe = "".join(ch for ch in (embedding_model or "") if ch.isalnum() or ch in ("-", "_", "."))
        if not safe:
            safe = "default"
        return f"ecda_data__{safe}"[:63]
    

    def add_documents(self, docs: List[Document], *, embedding_model: str | None = None) -> dict:
        """
        Incrementally add docs to existing collection (fast).
        """
        self.load(embedding_model=embedding_model)
        assert self.vs is not None
        if not docs:
            return {"ok": True, "chunks_added": 0}
        self.vs.add_documents(docs)
        return {"ok": True, "chunks_added": len(docs)}
    

    def delete_by_source_file(self, stored_name: str, *, embedding_model: str | None = None) -> dict:
        self.load(embedding_model=embedding_model)
        if not self.vs:
            return {"ok": True, "deleted": 0}

        col = getattr(self.vs, "_collection", None)
        if col is None:
            return {"ok": True, "deleted": 0}

        try:
            col.delete(where={"source_file": stored_name})
            return {"ok": True, "deleted": "unknown"}
        except Exception as e:
            raise RuntimeError(f"Failed deleting existing chunks for {stored_name}: {e}") from e


    # -----------------------------
    # Public API (used by rag.py)
    # -----------------------------
    def load(self, *, embedding_model: str | None = None) -> None:
        """
        Prepare the store for searching (creates VS if missing).
        """
        self.embedding_model = (embedding_model or DEFAULT_EMBED_MODEL).strip()
        self.embeddings = make_embeddings(embedding_model=self.embedding_model)
        self.vs = None  # force re-open with correct collection/embeddings
        self._ensure_store()

    def rebuild(self, docs: List[Document], *, embedding_model: str | None = None) -> dict:
        """
        Rebuild the collection from scratch (delete collection, then add docs).
        """
        self.load(embedding_model=embedding_model)  # ensures correct collection + embeddings
        assert self._client is not None

        name = self._collection_name(self.embedding_model)

        # Delete existing collection for a clean rebuild
        try:
            self._client.delete_collection(name=name)
        except Exception:
            pass

        # Re-open vectorstore after delete
        self.vs = Chroma(
            client=self._client,
            collection_name=name,
            embedding_function=self.embeddings,
        )

        if docs:
            self.vs.add_documents(docs)

        return {"ok": True, "chunks": len(docs)}

    # def search(self, query: str, k: int = 5) -> List[Document]:
    #     """
    #     Fast KNN search (Chroma HNSW under the hood).
    #     """
    #     self._ensure_store()
    #     if not self.vs:
    #         return []
    #     pairs = self.vs.similarity_search_with_score(query, k=k)
    #     return [doc for (doc, _score) in pairs]

    def search(self, query: str, k: int = 5) -> List[Document]:
        self._ensure_store()
        if not self.vs:
            return []

        # Best default for RAG: MMR (diverse, still fast)
        return self.vs.max_marginal_relevance_search(
            query,
            k=k,
            fetch_k=max(20, k * 4),
            lambda_mult=0.5
        )
    
    def search_with_scores(self, query: str, k: int = 5) -> List[tuple[Document, float]]:
        """Return (doc, score) pairs for relevance filtering."""
        self._ensure_store()
        if not self.vs:
            return []
        
        # Chroma uses similarity_search_with_relevance_scores for normalized scores (0-1, higher=better)
        return self.vs.similarity_search_with_relevance_scores(
            query, 
            k=k
        )


    # -----------------------------
    # Reset helpers (used by /api/reset)
    # -----------------------------
    def hard_reset(self) -> None:
        """
        Clears Chroma safely without restarting the app.
        Call this BEFORE deleting rag_store on disk.
        """
        self.vs = None
        self._client = None

        # Clear Chroma global cache (critical to avoid readonly sqlite after deleting files)
        try:
            from chromadb.api.client import SharedSystemClient
            SharedSystemClient.clear_system_cache()
        except Exception:
            pass

    def recreate(self) -> None:
        """
        Recreate persistent client and vector store after deleting rag_store on disk.
        """
        self.chroma_dir.mkdir(parents=True, exist_ok=True)
        self._client = chromadb.PersistentClient(
            path=str(self.chroma_dir),
            settings=Settings(anonymized_telemetry=False),
        )
        self.vs = None
        self._ensure_store()
