"""
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 ...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
        )

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

        # Get RAG configuration from settings
        self.embedding_model = get_setting_from_snapshot(
            "local_search_embedding_model",
            settings_snapshot,
            "all-MiniLM-L6-v2",
        )
        self.embedding_provider = get_setting_from_snapshot(
            "local_search_embedding_provider",
            settings_snapshot,
            "sentence_transformers",
        )
        self.chunk_size = get_setting_from_snapshot(
            "local_search_chunk_size", settings_snapshot, 1000
        )
        self.chunk_overlap = get_setting_from_snapshot(
            "local_search_chunk_overlap", settings_snapshot, 200
        )

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

    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 []

        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 = []
            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:
                    logger.warning(
                        f"Error searching collection {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:
                logger.info("No results found across any collections")
                return []

            # 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",
                                    self.settings_snapshot,
                                    str(get_library_directory()),
                                )
                                library_root = (
                                    Path(os.path.expandvars(library_root))
                                    .expanduser()
                                    .resolve()
                                )
                                pdf_manager = PDFStorageManager(
                                    library_root, "auto"
                                )
                                if pdf_manager.has_pdf(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:
            logger.exception("Error searching library RAG")
            return []

    def _get_previews(
        self,
        query: str,
        limit: int = 10,
        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.
        """
        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.
        """
        # Check if we should get full content
        from ...config import search_config

        if (
            hasattr(search_config, "SEARCH_SNIPPETS_ONLY")
            and search_config.SEARCH_SNIPPETS_ONLY
        ):
            logger.info("Snippet-only mode, skipping full content retrieval")
            return relevant_items

        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
