"""
Library RAG Search Engine

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

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

from ..search_engine_base import BaseSearchEngine
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


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
                    rag_service = LibraryRAGService(
                        username=self.username,
                        embedding_model=embedding_model,
                        embedding_provider=embedding_provider,
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap,
                    )

                    # 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 as e:
                    logger.warning(
                        f"Error searching collection {collection_id}: {e}"
                    )
                    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[:500] + "..."
                    if len(doc.page_content) > 500
                    else doc.page_content
                )

                # Generate URL to document content (PDF or TXT)
                # Use relative URLs - they work in any deployment scenario
                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:
                                # Use PDFStorageManager to check if PDF available
                                from pathlib import Path

                                library_root = get_setting_from_snapshot(
                                    "research_library.storage_path",
                                    self.settings_snapshot,
                                    "~/.local-deep-research/library",
                                )
                                library_root = Path(library_root).expanduser()
                                pdf_manager = PDFStorageManager(
                                    library_root, "auto"
                                )
                                if pdf_manager.has_pdf(document, session):
                                    document_url = (
                                        f"/library/document/{doc_id}/pdf"
                                    )
                                else:
                                    document_url = (
                                        f"/library/document/{doc_id}/txt"
                                    )
                            else:
                                # Fallback to TXT if document not found
                                document_url = f"/library/document/{doc_id}/txt"
                    except Exception as e:
                        logger.warning(
                            f"Error querying document {doc_id}: {e}, using TXT fallback"
                        )
                        document_url = f"/library/document/{doc_id}/txt"
                else:
                    document_url = "#"

                result = {
                    "title": title,
                    "snippet": snippet,
                    "url": document_url,
                    "link": document_url,  # Add "link" for source extraction
                    "source": "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 ... 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[:500] + "..."
                            if len(document.text_content) > 500
                            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
