# embedding_comparator.py
# ------------------------------------------------
# Python: EmbeddingModelComparator with optional rerankers
# and auto-inferred model names
# ------------------------------------------------
import pathlib
from typing import Any, Dict, Tuple, List, Union

import anywidget
import traitlets
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_classic.retrievers.contextual_compression import (
    ContextualCompressionRetriever,
)

current_directory = pathlib.Path(__file__).parent


class MongoDBEmbeddingComparator(anywidget.AnyWidget):
    _esm = str(current_directory / "comparator.js")
    _css = str(current_directory / "comparator.css")

    # Inputs
    query = traitlets.Unicode("").tag(sync=True)
    top_k = traitlets.Int(4).tag(sync=True)
    score_threshold = traitlets.Float(0.0).tag(sync=True)  # 0.0 means ignore

    # Trigger from frontend
    run_id = traitlets.Int(0).tag(sync=True)

    # Outputs
    # Structure: [ { "model": "OpenAI", "results": [ {content, score, metadata} ] }, ... ]
    comparison_data = traitlets.List(traitlets.Dict()).tag(sync=True)
    is_loading = traitlets.Bool(False).tag(sync=True)
    error = traitlets.Unicode("").tag(sync=True)

    def __init__(
        self,
        models_dict: Union[Dict[str, Any], List[Any]],
        loader=None,
        chunk_size: int = 500,
        chunk_overlap: int = 0,
        **kwargs,
    ):
        """
        models_dict can be:
        - dict: {"ignored_key": embedding_model, ...}
        - dict: {"ignored_key": (embedding_model, compressor), ...}
        - dict: {"ignored_key": {"embedding": embedding_model, "compressor": compressor}}
        - list: [embedding_model, (embedding_model, compressor), {...}, ...]

        Display names are inferred automatically from the objects and:
        - If only embedding: <embedding_name>
        - If embedding + compressor: <embedding_name>+<compressor_name>
        """
        super().__init__(**kwargs)

        self.vector_stores: Dict[str, FAISS] = {}
        self.compressors: Dict[str, Any] = {}

        try:
            # Load docs
            if hasattr(loader, "load"):
                raw_docs = loader.load()
            elif isinstance(loader, list):
                raw_docs = loader
            else:
                raw_docs = []

            if not raw_docs:
                self.error = "No documents found in loader."
                return

            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            )
            self.chunks = text_splitter.split_documents(raw_docs)

            # Normalize model specs into a list
            if isinstance(models_dict, dict):
                model_specs: List[Any] = list(models_dict.values())
            elif isinstance(models_dict, list):
                model_specs = models_dict
            else:
                self.error = "models_dict must be a dict or a list."
                return

            print(
                f"Building indices for {len(model_specs)} models. "
                f"This may take a moment..."
            )

            for idx, model_spec in enumerate(model_specs):
                embedding_model, compressor = self._parse_model_spec(model_spec)

                if embedding_model is None:
                    print(
                        f"Model #{idx} has no embedding model configured. "
                        f"Skipping."
                    )
                    continue

                display_name = self._build_display_name(
                    embedding_model, compressor, idx
                )

                # Ensure uniqueness if needed
                base_name = display_name
                suffix = 1
                while display_name in self.vector_stores:
                    display_name = f"{base_name}_{suffix}"
                    suffix += 1

                try:
                    vstore = FAISS.from_documents(self.chunks, embedding_model)
                    self.vector_stores[display_name] = vstore
                    if compressor is not None:
                        self.compressors[display_name] = compressor
                except Exception as e:
                    print(f"Failed to index model '{display_name}': {e}")
                    self.error = f"Failed to index model '{display_name}': {e}"

            print("Indexing complete.")

        except Exception as e:
            self.error = f"Initialization Error: {e}"

    def _parse_model_spec(self, model_spec: Any) -> Tuple[Any, Any]:
        """
        Normalize different ways of specifying models:
        - embedding_model
        - (embedding_model, compressor)
        - {"embedding": embedding_model, "compressor": compressor}
        - {"embedding_model": embedding_model, "reranker": compressor}
        """
        embedding_model = None
        compressor = None

        # Tuple: (embedding, compressor)
        if isinstance(model_spec, tuple) and len(model_spec) == 2:
            embedding_model, compressor = model_spec
            return embedding_model, compressor

        # Dict specification
        if isinstance(model_spec, dict):
            embedding_model = (
                model_spec.get("embedding")
                or model_spec.get("embedding_model")
                or model_spec.get("embedder")
            )
            compressor = model_spec.get("compressor") or model_spec.get("reranker")
            return embedding_model, compressor

        # Plain embedding model
        embedding_model = model_spec
        return embedding_model, None

    def _infer_obj_name(self, obj: Any, prefix: str, index: int) -> str:
        if obj is None:
            return f"{prefix}_{index}"

        # Try common attributes first
        for attr in ("model", "model_name", "deployment_name", "name"):
            if hasattr(obj, attr):
                val = getattr(obj, attr)
                if isinstance(val, str) and val:
                    return val

        # Fallback to class name
        cls_name = obj.__class__.__name__
        if cls_name:
            return cls_name

        return f"{prefix}_{index}"

    def _build_display_name(
        self, embedding_model: Any, compressor: Any, index: int
    ) -> str:
        emb_name = self._infer_obj_name(embedding_model, "embed", index)

        if compressor is None:
            return emb_name

        comp_name = self._infer_obj_name(compressor, "rerank", index)
        # As requested: concatenate with "+"
        return f"{emb_name}+{comp_name}"

    @traitlets.observe("query", "top_k", "score_threshold")
    def _on_input_change(self, change):
        # Optional: auto-run on input changes
        pass

    @traitlets.observe("run_id")
    def _on_run_id(self, change):
        # Called whenever the frontend bumps run_id
        self.run_comparison()

    def run_comparison(self, _=None):
        if not self.query:
            return

        if not self.vector_stores:
            self.error = "No vector stores available. Initialization may have failed."
            return

        self.is_loading = True
        self.error = ""
        self.comparison_data = []

        all_results = []

        try:
            for name, vstore in self.vector_stores.items():
                compressor = self.compressors.get(name)
                if compressor is not None:
                    # Vector search first, then rerank via compressor
                    initial_k = self.top_k
                    compressor_top_k = getattr(compressor, "top_k", None)
                    if isinstance(compressor_top_k, int) and compressor_top_k > 0:
                        initial_k = max(initial_k, compressor_top_k)

                    base_retriever = vstore.as_retriever(
                        search_kwargs={"k": initial_k}
                    )

                    compression_retriever = ContextualCompressionRetriever(
                        base_compressor=compressor,
                        base_retriever=base_retriever,
                    )

                    reranked_docs = compression_retriever.invoke(self.query)

                    filtered = []
                    total = len(reranked_docs) or 1

                    for rank, doc in enumerate(reranked_docs, start=1):
                        raw_score = 0.0
                        if isinstance(doc.metadata, dict):
                            raw_score = float(
                                doc.metadata.get("relevance_score", 0.0)
                            )

                        score = raw_score if raw_score != 0.0 else float(
                            total - rank + 1
                        )

                        if self.score_threshold > 0 and score < self.score_threshold:
                            continue

                        filtered.append(
                            {
                                "content": doc.page_content,
                                "score": float(score),
                                "metadata": doc.metadata,
                            }
                        )

                else:
                    # Plain FAISS similarity_search_with_score
                    docs_and_scores = vstore.similarity_search_with_score(
                        self.query,
                        k=self.top_k,
                    )

                    filtered = []
                    for doc, score in docs_and_scores:
                        if self.score_threshold > 0 and score < self.score_threshold:
                            continue

                        filtered.append(
                            {
                                "content": doc.page_content,
                                "score": float(score),
                                "metadata": doc.metadata,
                            }
                        )

                all_results.append(
                    {
                        "model": name,
                        "results": filtered,
                    }
                )

            self.comparison_data = all_results

        except Exception as e:
            self.error = str(e)
        finally:
            self.is_loading = False
