"""
Library RAG Search Engine

Provides semantic search over the user's personal research library using RAG.
"""

import os
from typing import List, Dict, Any, Optional
from loguru import logger

from ..search_engine_base import BaseSearchEngine
from ...constants import SNIPPET_LENGTH_LONG
from ...research_library.services.library_rag_service import LibraryRAGService
from ...research_library.services.library_service import LibraryService
from ...config.thread_settings import get_setting_from_snapshot
from ...utilities.llm_utils import get_server_url
from ...utilities.type_utils import to_bool
from ...database.models.library import RAGIndex, Document
from ...research_library.services.pdf_storage_manager import PDFStorageManager
from ...database.session_context import get_user_db_session
from ...config.paths import get_library_directory


class LibraryRAGSearchEngine(BaseSearchEngine):
    """
    Search engine that queries the user's research library using RAG/semantic search.
    """

    # Mark as local RAG engine
    is_local = True

    def __init__(
        self,
        llm: Optional[Any] = None,
        max_filtered_results: Optional[int] = None,
        max_results: int = 10,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Initialize the Library RAG search engine.

        Args:
            llm: Language model for relevance filtering
            max_filtered_results: Maximum number of results to keep after filtering
            max_results: Maximum number of search results
            settings_snapshot: Settings snapshot from thread context
            **kwargs: Additional engine-specific parameters
        """
        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            settings_snapshot=settings_snapshot,
            **kwargs,
        )
        self.username = (
            settings_snapshot.get("_username") if settings_snapshot else None
        )
        # Self-stamp the policy registry name so direct instantiations get
        # the runtime egress backstop without caller cooperation (the
        # factory stamps registry engines, but routes/services may build
        # this engine directly). CollectionSearchEngine overrides this
        # with its collection key after calling super().__init__.
        self._engine_name = "library"

        if not self.username:
            logger.warning(
                "Library RAG search engine initialized without username"
            )

        # Get RAG configuration from settings. NB: keyword arguments — the
        # signature is (key, default, username, settings_snapshot); these
        # calls used to pass (key, settings_snapshot, default) positionally,
        # so the snapshot was never consulted and, outside a thread settings
        # context, each setting silently resolved to the snapshot DICT
        # itself (the misplaced "default").
        self.embedding_model = get_setting_from_snapshot(
            "local_search_embedding_model",
            default="all-MiniLM-L6-v2",
            settings_snapshot=settings_snapshot,
        )
        self.embedding_provider = get_setting_from_snapshot(
            "local_search_embedding_provider",
            default="sentence_transformers",
            settings_snapshot=settings_snapshot,
        )
        self.chunk_size = get_setting_from_snapshot(
            "local_search_chunk_size",
            default=1000,
            settings_snapshot=settings_snapshot,
        )
        self.chunk_overlap = get_setting_from_snapshot(
            "local_search_chunk_overlap",
            default=200,
            settings_snapshot=settings_snapshot,
        )

        # Extract server URL from settings snapshot for link generation
        self.server_url = get_server_url(settings_snapshot)

        # Local RAG retrieval is vector-similarity ranked and can return many
        # near-duplicate chunks per query — enough to flood an agent's context
        # across research iterations. Apply a RAG-specific cap (default 20),
        # independent of the global ``search.max_results`` (which targets web
        # engines). This OVERRIDES the inherited max_results for local search.
        # Default to the inherited max_results so an ABSENT setting preserves
        # the caller's value (programmatic callers, tests); real web/CLI runs
        # carry the shipped default of 20 via default_settings.json, so they ARE
        # capped. A malformed value also falls back to no change.
        rag_max_results = get_setting_from_snapshot(
            "search.rag.max_results",
            default=self.max_results,
            settings_snapshot=settings_snapshot,
        )
        try:
            rag_max_results = int(rag_max_results)
        except (TypeError, ValueError):
            rag_max_results = self.max_results
        if rag_max_results > 0:
            self.max_results = rag_max_results

        # Embeddings already rank by cosine similarity, so the LLM relevance
        # filter is OPT-IN here (unlike keyword engines that hardcode
        # needs_llm_relevance_filter=True). Default on: prune
        # semantically-near-but-irrelevant chunks. The factory only ever SETS
        # enable_llm_relevance_filter=True (for should_filter engines) and never
        # unsets it — and RAG engines are not should_filter (no
        # needs_llm_relevance_filter; dynamic ``collection_<uuid>`` name) — so
        # this assignment is authoritative. The base filter is a no-op without
        # an llm, so it degrades gracefully.
        self.enable_llm_relevance_filter = to_bool(
            get_setting_from_snapshot(
                "search.rag.enable_relevance_filter",
                default=True,
                settings_snapshot=settings_snapshot,
            ),
            default=True,
        )

    def search(
        self,
        query: str,
        limit: int = 10,
        llm_callback=None,
        extra_params: Optional[Dict[str, Any]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Search the library using semantic search.

        Args:
            query: Search query
            limit: Maximum number of results to return
            llm_callback: Optional LLM callback for processing results
            extra_params: Additional search parameters

        Returns:
            List of search results with title, url, snippet, etc.
        """
        if not self.username:
            logger.error("Cannot search library without username")
            return []

        # search() can be called directly (bypassing run()), so apply the
        # runtime egress backstop here too. Memoized per snapshot, so the
        # run() -> _get_previews() -> search() path re-checks for free.
        # Raises PolicyDeniedError on denial — deliberately BEFORE the
        # broad try/except below.
        self._verify_egress_scope()

        try:
            # Initialize services
            library_service = LibraryService(username=self.username)

            # Get all collections for this user
            collections = library_service.get_all_collections()
            if not collections:
                logger.info("No collections found for user")
                return []

            # Search across all collections and merge results
            all_docs_with_scores = []
            failed_collections = []
            for collection in collections:
                collection_id = collection.get("id")
                if not collection_id:
                    continue

                try:
                    # Get the RAG index for this collection to find embedding settings
                    with get_user_db_session(self.username) as session:
                        collection_name = f"collection_{collection_id}"
                        rag_index = (
                            session.query(RAGIndex)
                            .filter_by(
                                collection_name=collection_name,
                                is_current=True,
                            )
                            .first()
                        )

                        if not rag_index:
                            logger.debug(
                                f"No RAG index found for collection {collection_id}"
                            )
                            continue

                        # Get embedding settings from the RAG index
                        embedding_model = rag_index.embedding_model
                        embedding_provider = (
                            rag_index.embedding_model_type.value
                        )
                        chunk_size = rag_index.chunk_size or self.chunk_size
                        chunk_overlap = (
                            rag_index.chunk_overlap or self.chunk_overlap
                        )

                    # Create RAG service with the collection's embedding settings
                    with LibraryRAGService(
                        username=self.username,
                        embedding_model=embedding_model,
                        embedding_provider=embedding_provider,
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap,
                    ) as rag_service:
                        # Get RAG stats to check if there are any indexed documents
                        stats = rag_service.get_rag_stats(collection_id)
                        if stats.get("indexed_documents", 0) == 0:
                            logger.debug(
                                f"No documents indexed in collection {collection_id}"
                            )
                            continue

                        # Load the FAISS index for this collection
                        vector_store = rag_service.load_or_create_faiss_index(
                            collection_id
                        )

                        # Search this collection's index
                        docs_with_scores = (
                            vector_store.similarity_search_with_score(
                                query, k=limit
                            )
                        )

                        # Add collection info to metadata and append to results
                        for doc, score in docs_with_scores:
                            if not doc.metadata:
                                doc.metadata = {}
                            doc.metadata["collection_id"] = collection_id
                            doc.metadata["collection_name"] = collection.get(
                                "name", "Unknown"
                            )
                            all_docs_with_scores.append((doc, score))

                except Exception:
                    # One broken collection must not abort the others, but
                    # record the failure so it is not silently equated with
                    # "no matching documents" below.
                    logger.exception(
                        f"Error searching collection {collection_id}"
                    )
                    failed_collections.append(
                        collection.get("name") or str(collection_id)
                    )
                    continue

            # Sort all results by score (lower is better for distance)
            all_docs_with_scores.sort(key=lambda x: x[1])

            # Take top results across all collections
            docs_with_scores = all_docs_with_scores[:limit]

            if not docs_with_scores:
                if failed_collections:
                    # Zero results AND at least one collection errored:
                    # reporting "no results" here would hide the failure.
                    self._raise_collections_failed(failed_collections)
                logger.info("No results found across any collections")
                return []

            if failed_collections:
                # Partial success: results from the healthy collections are
                # returned, but tell the user they are incomplete. WARNING
                # level streams to the research-log UI via
                # frontend_progress_sink.
                logger.warning(
                    f"Library search results are incomplete: "
                    f"{len(failed_collections)} of {len(collections)} "
                    f"collection(s) failed: {failed_collections}"
                )

            # Convert Document objects to search results format
            results = []
            for doc, score in docs_with_scores:
                # Extract metadata from Document object
                metadata = doc.metadata or {}

                # Try both source_id and document_id for compatibility
                doc_id = metadata.get("source_id") or metadata.get(
                    "document_id"
                )

                # Get title from metadata, with fallbacks
                title = (
                    metadata.get("document_title")
                    or metadata.get("title")
                    or (f"Document {doc_id}" if doc_id else "Untitled")
                )

                # Content is stored in page_content
                snippet = (
                    doc.page_content[:SNIPPET_LENGTH_LONG] + "..."
                    if len(doc.page_content) > SNIPPET_LENGTH_LONG
                    else doc.page_content
                )

                # Generate URL to document content
                # Default to root document page (shows all options: PDF, Text, Chunks, etc.)
                document_url = f"/library/document/{doc_id}" if doc_id else "#"

                if doc_id:
                    try:
                        with get_user_db_session(self.username) as session:
                            document = (
                                session.query(Document)
                                .filter_by(id=doc_id)
                                .first()
                            )
                            if document:
                                from pathlib import Path

                                library_root = get_setting_from_snapshot(
                                    "research_library.storage_path",
                                    default=str(get_library_directory()),
                                    settings_snapshot=self.settings_snapshot,
                                )
                                library_root = (
                                    Path(os.path.expandvars(library_root))
                                    .expanduser()
                                    .resolve()
                                )
                                if PDFStorageManager.pdf_exists(
                                    library_root, document, session
                                ):
                                    document_url = (
                                        f"/library/document/{doc_id}/pdf"
                                    )
                    except Exception:
                        logger.warning(f"Error querying document {doc_id}")

                result = {
                    "title": title,
                    "snippet": snippet,
                    "url": document_url,
                    "link": document_url,  # Add "link" for source extraction
                    "source": "library",
                    "source_type": "library",
                    "relevance_score": float(
                        1 / (1 + score)
                    ),  # Convert distance to similarity
                    "metadata": metadata,
                }

                results.append(result)

            logger.info(
                f"Library RAG search returned {len(results)} results for query: {query}"
            )
            return results

        except Exception:
            # Re-raise so run() records the failure in metrics instead of
            # treating a failed search as "no matching documents".
            logger.exception("Error searching library RAG")
            raise

    @staticmethod
    def _raise_collections_failed(failed_collections: List[str]) -> None:
        """Signal zero results caused (at least partly) by search failures."""
        raise RuntimeError(
            f"Library search returned no results and failed for "
            f"{len(failed_collections)} collection(s): {failed_collections}"
        )

    def _get_previews(
        self,
        query: str,
        limit: Optional[int] = None,
        llm_callback=None,
        extra_params: Optional[Dict[str, Any]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Get preview results for the query.
        Delegates to the search method, using the configured max_results
        when no explicit limit is provided.
        """
        if limit is None:
            limit = self.max_results
        return self.search(query, limit, llm_callback, extra_params)

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for relevant library documents.
        Retrieves complete document text instead of just snippets.
        """
        if not self.username:
            logger.error("Cannot retrieve full content without username")
            return relevant_items

        try:
            from ...database.models.library import Document
            from ...database.session_context import get_user_db_session

            # Retrieve full content for each document
            for item in relevant_items:
                doc_id = item.get("metadata", {}).get("document_id")
                if not doc_id:
                    continue

                # Get full document text from database
                with get_user_db_session(self.username) as db_session:
                    document = (
                        db_session.query(Document).filter_by(id=doc_id).first()
                    )

                    if document and document.text_content:
                        # Replace snippet with full content
                        item["content"] = document.text_content
                        item["snippet"] = (
                            document.text_content[:SNIPPET_LENGTH_LONG] + "..."
                            if len(document.text_content) > SNIPPET_LENGTH_LONG
                            else document.text_content
                        )
                        logger.debug(
                            f"Retrieved full content for document {doc_id}"
                        )

            return relevant_items

        except Exception:
            logger.exception("Error retrieving full content from library")
            return relevant_items

    def close(self):
        """Clean up resources."""
        pass
