# rag_widget.py
# ------------------------------------------------
# Python: MongoDBRAGPlayground (updated with top_k and score_threshold)
# ------------------------------------------------
import pathlib
import anywidget
import traitlets
from langchain_text_splitters import (
    CharacterTextSplitter,
    RecursiveCharacterTextSplitter,
    MarkdownTextSplitter,
)
from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub")

current_directory = pathlib.Path().resolve()

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

    # Which step is selected in the UI (1=Chunking, 2=Embedding, 3=RAG)
    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 (for backward compatibility)
    embedding_model = traitlets.Unicode("").tag(sync=True)

    # List of available embedding model names for dropdown
    embedding_models = traitlets.List(traitlets.Unicode(), default_value=[]).tag(sync=True)

    # Index of the currently selected embedding model (for indexing)
    selected_embedding_index = traitlets.Int(0).tag(sync=True)

    # Index of the currently selected embedding model for query time
    query_embedding_index = traitlets.Int(0).tag(sync=True)

    # Display name for the chat/generative model
    chat_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)

    # RAG controls
    rag_query = traitlets.Unicode("").tag(sync=True)  # The user's question
    rag_results = traitlets.List(traitlets.Dict()).tag(sync=True)  # RAG search results
    rag_prompt = traitlets.Unicode("").tag(sync=True)  # Final expanded prompt
    rag_answer = traitlets.Unicode("").tag(sync=True)  # Final generated answer
    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

    # RAG TRAITLETS for k and threshold
    top_k = traitlets.Int(5).tag(sync=True)
    score_threshold = traitlets.Union(
        [traitlets.Float(), traitlets.Float(allow_none=True), traitlets.Unicode(allow_none=True)],
        default_value=None,
        allow_none=True
    ).tag(sync=True)

    # The user-editable template
    rag_prompt_template = traitlets.Unicode("").tag(sync=True)

    # 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 test_error(self, message="This is a test error message"):
        """Function to test the error display system"""
        self.set_error(message)

    def __init__(
        self,
        loader=None,
        embedding_model=None,
        llm=None,
        mongo_collection=None,
        index_name=None,
        **kwargs,
    ):
        """
        loader: optional doc loader
        embedding_model: single embedding model or list of embedding models for vectorstore
        llm: language model for generating final RAG answer
        mongo_collection: pymongo collection for storing embedded docs
        index_name: name of the MongoDB Atlas Search index to use
        """
        # Keep the actual Python embedding object separate
        self.loader = loader
        self.llm = llm
        self.mongo_collection = mongo_collection

        # Handle both single embedding model and list of models
        if isinstance(embedding_model, list):
            self._embedding_models = embedding_model  # internal Python objects
            self._embedding_model = embedding_model[0] if embedding_model else None
        else:
            self._embedding_models = [embedding_model] if embedding_model else []
            self._embedding_model = embedding_model  # internal Python object

        # Store the index name for later use
        self._index_name = 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 = []
            # Don't set error yet, wait until super().__init__ is called

        self.chunks_by_page = []

        # Initialize parent class first so traitlets are properly set up
        super().__init__(**kwargs)

        # Now set traitlet values after initialization
        # 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:
                # Leave them as empty strings if something goes wrong
                self.mongo_db_name = ""
                self.mongo_collection_name = ""

        # Human-readable model names for the embedding model UI
        model_display_names = []
        for model in self._embedding_models:
            display_name = ""
            try:
                if isinstance(model, str):
                    display_name = model
                elif hasattr(model, "model"):
                    display_name = str(getattr(model, "model"))
                elif hasattr(model, "model_name"):
                    display_name = str(getattr(model, "model_name"))
                elif model is not None:
                    display_name = model.__class__.__name__
            except Exception as e:
                display_name = f"Unknown Model {len(model_display_names) + 1}"

            # Ensure we have a valid string
            if not display_name or not isinstance(display_name, str):
                display_name = f"Model {len(model_display_names) + 1}"

            model_display_names.append(display_name)

        # Always ensure we have at least an empty list, never None
        self.embedding_models = model_display_names if model_display_names else []

        # Set the current embedding_model for backward compatibility
        # This should be a string, not the list
        if model_display_names:
            self.embedding_model = str(model_display_names[0])
        else:
            self.embedding_model = ""

        # Human-readable model name for the chat/generative model UI
        chat_model_name = ""
        try:
            if isinstance(llm, str):
                chat_model_name = llm
            elif hasattr(llm, "model"):
                chat_model_name = str(getattr(llm, "model"))
            elif hasattr(llm, "model_name"):
                chat_model_name = str(getattr(llm, "model_name"))
            elif llm is not None:
                chat_model_name = llm.__class__.__name__
        except Exception:
            chat_model_name = ""
        self.chat_model = chat_model_name

        if self._index_name:
            self.selected_index = self._index_name

        self.rag_prompt_template = (
            "<context>\n{context}\n</context>\n"
            "<question>{question}</question>\n"
            "<instructions>Answer the user QUESTION using the CONTEXT text above.\n"
            "Keep your answer grounded in the facts of the CONTEXT.\n"
            "If the CONTEXT doesn’t contain the facts to answer the QUESTION, "
            "respond that you don't know.</instructions>"
        )

        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)

        # New palette: brighter, lighter fills so they pop on dark background
        # but still look fine in light mode.
        colors = [
            "rgba(0, 237, 100, 0.40)",   # bright green
            "rgba(0, 214, 255, 0.40)",   # cyan
            "rgba(255, 214, 10, 0.40)",  # yellow
            "rgba(255, 159, 243, 0.40)", # pink
            "rgba(156, 220, 254, 0.40)", # light blue
        ]
        # Overlaps: slightly stronger accent
        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()

    @traitlets.observe("selected_embedding_index")
    def _on_embedding_model_change(self, change):
        """Update the active embedding model when user selects a different one"""
        new_index = change["new"]
        if 0 <= new_index < len(self._embedding_models):
            self._embedding_model = self._embedding_models[new_index]
            self.embedding_model = self.embedding_models[new_index]
            # Reset the loaded state since we're using a different model
            self.loaded_in_mongo = False

    # -----------------
    # 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 == "rag_ask":
            self.run_rag()
        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: RAG
    # -----------------
    def run_rag(self):
        # Clear old state for this single run
        self.rag_results = []
        self.rag_answer = ""
        self.query_embedding = []
        self.clear_error()

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

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

        try:
            # Select the query-time embedding model (may be different from indexing model)
            query_emb_index = self.query_embedding_index
            if 0 <= query_emb_index < len(self._embedding_models):
                query_embedding_model = self._embedding_models[query_emb_index]
                query_model_name = self.embedding_models[query_emb_index] if query_emb_index < len(self.embedding_models) else "Unknown"
            else:
                # Fallback to the indexing model
                query_embedding_model = self._embedding_model
                query_model_name = self.embedding_model

            print(f"🔍 Query embedding model used: {query_model_name}")
            print(f"   (Index: {query_emb_index}, Total models available: {len(self._embedding_models)})")

            # Build vector store with the query embedding model
            vector_store = MongoDBAtlasVectorSearch(
                collection=self.mongo_collection,
                embedding=query_embedding_model,
                index_name=self.selected_index,
                relevance_score_fn="cosine",
            )

            # Generate query embedding using the selected query embedding model
            query_emb = query_embedding_model.embed_query(query)
            self.query_embedding = [float(x) for x in query_emb]
            print(f"✅ Query embedding generated with {query_model_name} (dimension: {len(self.query_embedding)})")

            # k (fallback to 5 if bad)
            k = self.top_k if isinstance(self.top_k, int) and self.top_k > 0 else 5
            search_args = {"k": k}

            # Optional threshold
            threshold_val = None
            if self.score_threshold not in (None, ""):
                try:
                    threshold_val = float(self.score_threshold)
                except (ValueError, TypeError):
                    threshold_val = None

            # One retrieval - get more than k to allow threshold filtering
            # Get a larger set first, then filter by threshold
            retrieval_k = max(k * 3, 20)  # Get 3x k or at least 20 documents
            large_search_args = {"k": retrieval_k}

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

            # Filter by threshold if provided (this gives us ALL docs above threshold)
            if threshold_val is not None:
                docs_above_threshold = [
                    (doc, score)
                    for (doc, score) in all_scored_docs
                    if score >= threshold_val
                ]
            else:
                docs_above_threshold = all_scored_docs

            # Now take top K from the threshold-filtered docs for RAG context
            docs_with_scores = docs_above_threshold[:k]

            # Create a mapping from text to embedding for all docs above threshold
            text_to_embedding = {}
            all_embeddings_for_viz = []

            # Process ALL documents above threshold for visualization
            for idx, (doc, score) in enumerate(docs_above_threshold):
                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 for retrieved docs
            results_with_embeddings = []
            for (doc, score) in docs_with_scores:
                # Try to find the embedding by matching the text content
                embedding = text_to_embedding.get(doc.page_content, [])

                # Format embedding with sample (first 5 values) like in Indexing section
                embedding_sample = (
                    embedding[:5] + ["..."]
                    if len(embedding) > 5
                    else embedding
                )

                result = {
                    "score": float(score),
                    "text": doc.page_content,  # Changed from "content" to "text"
                    "metadata": doc.metadata,
                    "embedding": embedding_sample  # Use sampled version
                }
                results_with_embeddings.append(result)

            self.rag_results = results_with_embeddings

            # Build context
            context_str = "\n\n".join(
                [doc.page_content for (doc, _) in docs_with_scores]
            )

            # Use user template or default
            template = self.rag_prompt_template or (
                "<context>\n{context}\n</context>\n"
                "<question>{question}</question>\n"
                "<instructions>Answer the user QUESTION using the CONTEXT text above.\n"
                "Keep your answer grounded in the facts of the CONTEXT.\n"
                "If the CONTEXT doesn’t contain the facts to answer the QUESTION, "
                "respond that you don't know.</instructions>"
            )

            prompt = ChatPromptTemplate.from_template(template)
            parse_output = StrOutputParser()

            # One generation
            rag_chain = (
                {
                    "context": (lambda x: context_str),
                    "question": RunnablePassthrough(),
                }
                | prompt
                | self.llm
                | parse_output
            )

            answer = rag_chain.invoke(query)
            self.rag_answer = answer

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