# two_stage_retrieval_widget.py
# ------------------------------------------------
# Python: MongoDBTwoStageRetrievalWidget - demonstrates two-stage retrieval with reranking
# ------------------------------------------------
import pathlib
import anywidget
import traitlets
from langchain_text_splitters import (
    CharacterTextSplitter,
    RecursiveCharacterTextSplitter,
    MarkdownTextSplitter,
)
from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
from langchain_classic.retrievers.contextual_compression import ContextualCompressionRetriever
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub")

current_directory = pathlib.Path().resolve()

class MongoDBTwoStageRetrievalWidget(anywidget.AnyWidget):
    _esm = str(pathlib.Path(__file__).parent / "two_stage_retrieval.js")
    _css = str(pathlib.Path(__file__).parent / "two_stage_retrieval.css")

    # Which step is selected in the UI (1=Chunking, 2=Indexing, 3=Two-Stage Retrieval)
    current_step = traitlets.Int(1).tag(sync=True)

    # Chunking controls
    split_strategy = traitlets.Unicode("Fixed").tag(sync=True)
    chunk_size = traitlets.Int(512).tag(sync=True)
    overlap_size = traitlets.Int(0).tag(sync=True)
    current_doc_index = traitlets.Int(0).tag(sync=True)
    document_preview = traitlets.Unicode("").tag(sync=True)
    chunks_table = traitlets.List(traitlets.Dict()).tag(sync=True)

    # Step 2 controls
    selected_index = traitlets.Unicode("").tag(sync=True)
    embeddings_table = traitlets.List(traitlets.Dict()).tag(sync=True)
    mongo_docs_table = traitlets.List(traitlets.Dict()).tag(sync=True)
    embedding_ready = traitlets.Bool(False).tag(sync=True)
    loaded_in_mongo = traitlets.Bool(False).tag(sync=True)
    vector_indexes = traitlets.List(traitlets.Unicode()).tag(sync=True)

    # Display name for the embedding model
    embedding_model = traitlets.Unicode("").tag(sync=True)

    # Display name for the query embedding model (can be different from indexing model)
    query_embedding_model = traitlets.Unicode("").tag(sync=True)

    # Display name for the reranker model
    reranker_model = traitlets.Unicode("").tag(sync=True)

    # DB and collection names, inferred from mongo_collection
    mongo_db_name = traitlets.Unicode("").tag(sync=True)
    mongo_collection_name = traitlets.Unicode("").tag(sync=True)

    # Retrieval controls
    retrieval_query = traitlets.Unicode("").tag(sync=True)  # The user's query
    initial_results = traitlets.List(traitlets.Dict()).tag(sync=True)  # Initial vector search results
    reranked_results = traitlets.List(traitlets.Dict()).tag(sync=True)  # Reranked results
    query_embedding = traitlets.List(traitlets.Float()).tag(sync=True)  # Query embedding vector
    all_doc_embeddings = traitlets.List(traitlets.Dict()).tag(sync=True)  # All document embeddings for visualization

    # Retrieval parameters
    initial_k = traitlets.Int(20).tag(sync=True)  # Number of documents to retrieve initially
    rerank_top_k = traitlets.Int(3).tag(sync=True)  # Number of documents after reranking

    # Command and error
    command = traitlets.Unicode("").tag(sync=True)
    error = traitlets.Unicode("").tag(sync=True)

    def set_error(self, message):
        """Set an error message and ensure it's sent to the frontend"""
        if message:
            self.error = message
            self.send({"type": "update_error", "error": message})
            print(f"Error: {message}")

    def clear_error(self):
        """Clear any error messages"""
        self.error = ""
        self.send({"type": "update_error", "error": ""})

    def __init__(
        self,
        loader=None,
        embedding_model=None,
        query_embedding_model=None,
        reranker=None,
        mongo_collection=None,
        index_name=None,
        **kwargs,
    ):
        """
        loader: optional doc loader
        embedding_model: embedding model for vectorstore indexing
        query_embedding_model: optional different embedding model for query (defaults to embedding_model)
        reranker: reranker model for two-stage retrieval (e.g., VoyageAIRerank)
        mongo_collection: pymongo collection for storing embedded docs
        index_name: name of the MongoDB Atlas Search index to use
        """
        # Keep the actual Python objects separate
        self.loader = loader
        self.mongo_collection = mongo_collection
        self._embedding_model = embedding_model  # internal Python object for indexing
        self._query_embedding_model = query_embedding_model if query_embedding_model is not None else embedding_model
        self._reranker = reranker  # internal reranker object

        # Infer DB and collection name from the provided pymongo collection
        if self.mongo_collection is not None:
            try:
                self.mongo_db_name = self.mongo_collection.database.name
                self.mongo_collection_name = self.mongo_collection.name
            except Exception:
                self.mongo_db_name = ""
                self.mongo_collection_name = ""

        # Human-readable model name for the indexing embedding model UI
        display_name = ""
        try:
            if isinstance(embedding_model, str):
                display_name = embedding_model
            elif hasattr(embedding_model, "model"):
                display_name = str(getattr(embedding_model, "model"))
            elif hasattr(embedding_model, "model_name"):
                display_name = str(getattr(embedding_model, "model_name"))
            elif embedding_model is not None:
                display_name = embedding_model.__class__.__name__
        except Exception:
            display_name = ""
        self.embedding_model = display_name

        # Human-readable model name for the query embedding model UI
        query_display_name = ""
        try:
            if isinstance(self._query_embedding_model, str):
                query_display_name = self._query_embedding_model
            elif hasattr(self._query_embedding_model, "model"):
                query_display_name = str(getattr(self._query_embedding_model, "model"))
            elif hasattr(self._query_embedding_model, "model_name"):
                query_display_name = str(getattr(self._query_embedding_model, "model_name"))
            elif self._query_embedding_model is not None:
                query_display_name = self._query_embedding_model.__class__.__name__
        except Exception:
            query_display_name = display_name  # Fallback to indexing model name
        self.query_embedding_model = query_display_name

        # Human-readable model name for the reranker UI
        reranker_display_name = ""
        try:
            if isinstance(reranker, str):
                reranker_display_name = reranker
            elif hasattr(reranker, "model"):
                reranker_display_name = str(getattr(reranker, "model"))
            elif hasattr(reranker, "model_name"):
                reranker_display_name = str(getattr(reranker, "model_name"))
            elif reranker is not None:
                reranker_display_name = reranker.__class__.__name__
        except Exception:
            reranker_display_name = ""
        self.reranker_model = reranker_display_name

        if index_name:
            self.selected_index = index_name

        try:
            if hasattr(loader, "load"):
                docs = loader.load()
            elif isinstance(loader, list):
                docs = loader
            else:
                docs = []
            self.loaded_pages = [doc.page_content for doc in docs]
        except Exception as e:
            self.loaded_pages = []
            self.set_error(f"Error loading document: {e}")

        self.chunks_by_page = []

        super().__init__(**kwargs)

        self._create_chunks_for_all_pages()
        self._update_highlighted_preview()
        self._build_embeddings_table()
        self.mongo_docs_table = []
        self._load_vector_indexes()

    # -----------------
    # Step 1: CHUNKING
    # -----------------
    def _create_chunks_for_all_pages(self):
        self.chunks_by_page = []
        table_accumulator = []

        if not self.loaded_pages:
            self.chunks_table = []
            return

        splitter = self._get_text_splitter(add_start_index=True)
        for p_idx, page_text in enumerate(self.loaded_pages):
            docs = splitter.create_documents([page_text])
            for c_idx, doc in enumerate(docs):
                table_accumulator.append(
                    {
                        "page_index": p_idx,
                        "chunk_index": c_idx,
                        "chunk_text": doc.page_content,
                        "start_offset": doc.metadata["start_index"],
                        "end_offset": doc.metadata["start_index"]
                        + len(doc.page_content),
                    }
                )
        self.chunks_table = table_accumulator

    def _get_text_splitter(self, add_start_index=False):
        if self.split_strategy == "Recursive":
            return RecursiveCharacterTextSplitter(
                chunk_size=self.chunk_size,
                chunk_overlap=self.overlap_size,
                add_start_index=add_start_index,
            )
        elif self.split_strategy == "Markdown":
            return MarkdownTextSplitter(
                chunk_size=self.chunk_size,
                chunk_overlap=self.overlap_size,
                add_start_index=add_start_index,
            )
        else:
            return CharacterTextSplitter(
                separator="",
                chunk_size=self.chunk_size,
                chunk_overlap=self.overlap_size,
                add_start_index=add_start_index,
            )

    def _build_highlighted_html(self, page_text, chunks_info):
        if not page_text:
            return page_text

        coverage = [[] for _ in range(len(page_text))]
        for info in chunks_info:
            start = info.get("start_offset", 0)
            end = info.get("end_offset", 0)
            c_idx = info["chunk_index"]
            for i in range(start, min(end, len(page_text))):
                coverage[i].append(c_idx)

        colors = [
            "rgba(0, 237, 100, 0.40)",
            "rgba(0, 214, 255, 0.40)",
            "rgba(255, 214, 10, 0.40)",
            "rgba(255, 159, 243, 0.40)",
            "rgba(156, 220, 254, 0.40)",
        ]
        overlap_color = "rgba(255, 99, 132, 0.55)"

        def html_escape(txt):
            return txt.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

        def style_for(indices):
            if len(indices) > 1:
                return f'style="background-color:{overlap_color}"'
            elif len(indices) == 1:
                idx_mod = indices[0] % len(colors)
                return f'style="background-color:{colors[idx_mod]}"'
            else:
                return ""

        output = []
        if len(page_text) == 0:
            return ""

        span_buffer = page_text[0]
        last_cov = coverage[0]

        for i in range(1, len(page_text)):
            current_cov = coverage[i]
            if len(current_cov) != len(last_cov) or set(current_cov) != set(last_cov):
                output.append(f'<span {style_for(last_cov)}>{html_escape(span_buffer)}</span>')
                span_buffer = page_text[i]
                last_cov = current_cov
            else:
                span_buffer += page_text[i]

        if span_buffer:
            output.append(f'<span {style_for(last_cov)}>{html_escape(span_buffer)}</span>')

        return "".join(output)

    def _update_highlighted_preview(self):
        if 0 <= self.current_doc_index < len(self.loaded_pages):
            page_text = self.loaded_pages[self.current_doc_index]
            page_chunks_info = [
                row
                for row in self.chunks_table
                if row["page_index"] == self.current_doc_index
            ]
            self.document_preview = self._build_highlighted_html(
                page_text, page_chunks_info
            )
        else:
            self.document_preview = (
                f"No page found at index {self.current_doc_index}"
            )

    @traitlets.observe("current_doc_index")
    def _on_page_change(self, change):
        self._update_highlighted_preview()

    @traitlets.observe("chunk_size", "split_strategy", "overlap_size")
    def _on_chunk_settings_change(self, change):
        self._create_chunks_for_all_pages()
        self._update_highlighted_preview()
        self._build_embeddings_table()

    # -----------------
    # Step 2: EMBEDDING
    # -----------------
    def _load_vector_indexes(self):
        if self.mongo_collection is None:
            self.vector_indexes = []
            return
        try:
            all_indexes = list(self.mongo_collection.list_search_indexes())
            self.vector_indexes = [idx.get("name", "") for idx in all_indexes]
            if not self.vector_indexes:
                self.vector_indexes = []
        except Exception as e:
            if "command not found" in str(e).lower():
                idx_info = self.mongo_collection.index_information()
                self.vector_indexes = list(idx_info.keys())
            else:
                self.set_error(f"Error fetching search indexes: {e}")
                self.vector_indexes = []

        if self.selected_index and self.selected_index not in self.vector_indexes:
            self.set_error(
                f"Error: Provided vector index '{self.selected_index}' does not exist."
            )

    def _build_embeddings_table(self):
        data = []
        for row in self.chunks_table:
            data.append(
                {
                    "chunk_text": row["chunk_text"],
                    "page_index": row["page_index"],
                    "chunk_index": row["chunk_index"],
                    "start_offset": row.get("start_offset", 0),
                    "end_offset": row.get("end_offset", 0),
                }
            )
        self.embeddings_table = data
        self.embedding_ready = False
        self.loaded_in_mongo = False

    @traitlets.observe("command")
    def _on_command(self, change):
        cmd = change["new"]
        if cmd == "load_into_mongo":
            self.load_into_mongo()
        elif cmd == "retrieve":
            self.run_retrieval()
        elif cmd == "rerank":
            self.run_reranking()
        self.command = ""

    def load_into_mongo(self):
        if self.mongo_collection is None:
            self.set_error("No MongoDB collection provided.")
            return
        if self._embedding_model is None:
            self.set_error("No embedding model provided.")
            return

        self.mongo_collection.delete_many({})
        vector_store = MongoDBAtlasVectorSearch(
            collection=self.mongo_collection,
            embedding=self._embedding_model,
            index_name=self.selected_index,
            relevance_score_fn="cosine",
        )
        texts = [r["chunk_text"] for r in self.chunks_table]
        vector_store.add_texts(texts, metadatas=[{} for _ in texts])

        preview_cursor = self.mongo_collection.find(
            {}, {"_id": 1, "text": 1, "embedding": 1}
        )

        self.mongo_docs_table = [
            {
                "_id": str(doc["_id"]),
                "text": (
                    doc.get("text", "")[:60] + "..."
                    if len(doc.get("text", "")) > 60
                    else doc.get("text", "")
                ),
                "embedding": (
                    doc.get("embedding", [])[:5] + ["..."]
                    if len(doc.get("embedding", [])) > 5
                    else doc.get("embedding", [])
                ),
            }
            for doc in preview_cursor
        ]
        self.loaded_in_mongo = True

    # -----------------
    # Step 3: TWO-STAGE RETRIEVAL
    # -----------------
    def run_retrieval(self):
        """Run initial vector search retrieval"""
        # Clear old state
        self.initial_results = []
        self.reranked_results = []
        self.query_embedding = []
        self.clear_error()

        # Basic validation
        if not self.selected_index:
            self.set_error("No vector index was specified.")
            return
        if self._query_embedding_model is None:
            self.set_error("No query embedding model provided.")
            return

        query = (self.retrieval_query or "").strip()
        if not query:
            self.set_error("No query provided.")
            return

        try:
            # Build vector store
            vector_store = MongoDBAtlasVectorSearch(
                collection=self.mongo_collection,
                embedding=self._query_embedding_model,
                index_name=self.selected_index,
                relevance_score_fn="cosine",
            )

            # Generate query embedding
            query_emb = self._query_embedding_model.embed_query(query)
            self.query_embedding = [float(x) for x in query_emb]

            # Initial retrieval with specified k
            k = self.initial_k if isinstance(self.initial_k, int) and self.initial_k > 0 else 20
            search_args = {"k": k}

            all_scored_docs = vector_store.similarity_search_with_score(
                query, **search_args
            )

            # Store embeddings for visualization
            text_to_embedding = {}
            all_embeddings_for_viz = []

            for idx, (doc, score) in enumerate(all_scored_docs):
                text = doc.page_content
                # Fetch the embedding from MongoDB
                mongo_doc = self.mongo_collection.find_one(
                    {"text": text},
                    {"_id": 1, "text": 1, "embedding": 1}
                )
                if mongo_doc:
                    embedding = mongo_doc.get("embedding", [])
                    if embedding:
                        text_to_embedding[text] = embedding
                        display_text = text[:50] + "..." if len(text) > 50 else text
                        all_embeddings_for_viz.append({
                            "id": str(mongo_doc.get("_id", idx)),
                            "text": display_text,
                            "embedding": embedding,
                            "score": float(score)
                        })

            self.all_doc_embeddings = all_embeddings_for_viz

            # Update UI results with embeddings
            results_with_embeddings = []
            for (doc, score) in all_scored_docs:
                embedding = text_to_embedding.get(doc.page_content, [])
                embedding_sample = (
                    embedding[:5] + ["..."]
                    if len(embedding) > 5
                    else embedding
                )

                result = {
                    "score": float(score),
                    "text": doc.page_content,
                    "metadata": doc.metadata,
                    "embedding": embedding_sample
                }
                results_with_embeddings.append(result)

            self.initial_results = results_with_embeddings

        except Exception as e:
            self.set_error(f"Retrieval Error: {e}")

    def run_reranking(self):
        """Run reranking on the initial results"""
        if not self.initial_results:
            self.set_error("No initial results to rerank. Please run retrieval first.")
            return

        if self._reranker is None:
            self.set_error("No reranker model provided.")
            return

        query = (self.retrieval_query or "").strip()
        if not query:
            self.set_error("No query provided.")
            return

        try:
            # Build vector store
            vector_store = MongoDBAtlasVectorSearch(
                collection=self.mongo_collection,
                embedding=self._query_embedding_model,
                index_name=self.selected_index,
                relevance_score_fn="cosine",
            )

            # Create base retriever
            base_retriever = vector_store.as_retriever(
                search_kwargs={"k": self.initial_k}
            )

            # Create compression retriever with reranker
            compression_retriever = ContextualCompressionRetriever(
                base_compressor=self._reranker,
                base_retriever=base_retriever
            )

            # Execute reranking
            compressed_docs = compression_retriever.invoke(query)

            # Store embeddings for reranked docs
            text_to_embedding = {}
            for result in self.initial_results:
                # Reconstruct full text from initial results
                text = result["text"]
                # Get full embedding from MongoDB
                mongo_doc = self.mongo_collection.find_one(
                    {"text": text},
                    {"_id": 1, "text": 1, "embedding": 1}
                )
                if mongo_doc:
                    embedding = mongo_doc.get("embedding", [])
                    if embedding:
                        text_to_embedding[text] = embedding

            # Process reranked results
            reranked_results = []
            for doc in compressed_docs:
                # Get relevance score from metadata
                relevance_score = doc.metadata.get("relevance_score", 0.0)

                embedding = text_to_embedding.get(doc.page_content, [])
                embedding_sample = (
                    embedding[:5] + ["..."]
                    if len(embedding) > 5
                    else embedding
                )

                result = {
                    "score": float(relevance_score),
                    "text": doc.page_content,
                    "metadata": doc.metadata,
                    "embedding": embedding_sample
                }
                reranked_results.append(result)

            self.reranked_results = reranked_results

        except Exception as e:
            self.set_error(f"Reranking Error: {e}")
