# BMLibrarian Lite - Biomedical Literature Research Tool
# Copyright (C) 2024-2025 Dr Horst Herb
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

"""
Lite search agent for literature queries.

This agent handles converting natural language questions to search queries,
executing searches across PubMed and/or Europe PMC, and caching results
in SQLite with vector embeddings.

Supports multiple search providers:
- PubMed: NCBI E-utilities API (default)
- Europe PMC: REST API with preprint support
- Both: Combined search with automatic deduplication
"""

import logging
from typing import Optional, Callable

from bmlibrarian_lite.pubmed import (
    PubMedSearchClient,
    PubMedQuery,
    ArticleMetadata,
)
from ..storage import LiteStorage
from ..config import LiteConfig
from ..data_models import LiteDocument, DocumentSource, SearchSession, SearchProvider
from ..embeddings import LiteEmbedder
from ..query_converter import LiteQueryConverter
from ..search_service import SearchService, UnifiedSearchResult
from ..query_translator import QueryTranslator
from .base import LiteBaseAgent

logger = logging.getLogger(__name__)


class LiteSearchAgent(LiteBaseAgent):
    """
    Search agent for literature queries with SQLite caching.

    Converts natural language queries to search queries and executes
    searches across PubMed and/or Europe PMC, caching results in
    SQLite with vector embeddings.

    This agent:
    1. Uses LLM to convert research questions to optimized search queries
    2. Executes searches via PubMed, Europe PMC, or both (configurable)
    3. Merges and deduplicates results when using both providers
    4. Caches results in SQLite with embeddings for later retrieval
    5. Provides semantic search over cached documents

    Attributes:
        storage: LiteStorage instance for persistence
        search_service: Unified search service for multi-provider searches
    """

    TASK_ID = "query_conversion"

    def __init__(
        self,
        storage: Optional[LiteStorage] = None,
        config: Optional[LiteConfig] = None,
        **kwargs,
    ) -> None:
        """
        Initialize the search agent.

        Args:
            storage: LiteStorage instance (creates one if not provided)
            config: Lite configuration
            **kwargs: Additional arguments for base agent
        """
        super().__init__(config=config, **kwargs)
        self.storage = storage or LiteStorage(self.config)

        # Initialize unified search service (handles both PubMed and Europe PMC)
        self.search_service = SearchService(self.config)

        # Keep PubMed client for backward compatibility
        self._search_client = self.search_service.pubmed_client

        # Initialize lite query converter (simplified, focused queries)
        # Uses query_conversion task configuration
        self._query_converter = LiteQueryConverter(
            llm_client=self.llm_client,
            model=self._get_model("query_conversion"),
        )

        # Create embedder (lazy initialization)
        self._embedder: Optional[LiteEmbedder] = None

    @property
    def embedder(self) -> LiteEmbedder:
        """Get or create embedder."""
        if self._embedder is None:
            self._embedder = LiteEmbedder(
                model_name=self.config.embeddings.model
            )
        return self._embedder

    def convert_query(self, question: str) -> PubMedQuery:
        """
        Convert a natural language question to a PubMed query.

        Uses LLM to parse the question and generate a focused
        PubMed query with key MeSH terms and keywords.

        Args:
            question: Natural language research question

        Returns:
            PubMedQuery object with query string and metadata
        """
        logger.info(f"Converting question to PubMed query: {question[:100]}...")

        query = self._query_converter.convert(question)

        logger.debug(f"Generated query: {query.query_string}")
        return query

    def search(
        self,
        question: str,
        max_results: Optional[int] = None,
        provider: Optional[SearchProvider] = None,
        include_preprints: Optional[bool] = None,
        progress_callback: Optional[Callable[[str], None]] = None,
    ) -> tuple[SearchSession, list[LiteDocument]]:
        """
        Search for literature and cache results.

        This is the main entry point for searching. It:
        1. Converts the question to a search query
        2. Executes the search via the configured provider(s)
        3. Merges and deduplicates results (if searching both providers)
        4. Caches results in SQLite

        Args:
            question: Natural language research question
            max_results: Maximum results to fetch (uses config default if None)
            provider: Search provider (uses config default if None)
            include_preprints: Whether to include preprints (Europe PMC only)
            progress_callback: Optional callback for progress updates

        Returns:
            Tuple of (search session, list of documents)
        """
        max_results = max_results or self.config.search.max_results
        provider = provider or self.config.search.search_provider

        if progress_callback:
            progress_callback(f"Converting question to search query...")

        # Convert to search query (generates PubMed syntax)
        pubmed_query = self.convert_query(question)
        logger.info(f"Generated query: {pubmed_query.query_string}")

        if progress_callback:
            provider_name = provider.display_name
            progress_callback(f"Searching {provider_name}...")

        # Execute search via unified search service
        search_result = self.search_service.search(
            query=pubmed_query.query_string,
            max_results=max_results,
            provider=provider,
            include_preprints=include_preprints,
            progress_callback=progress_callback,
        )

        logger.info(
            f"Found {search_result.total_count} total results, "
            f"retrieved {search_result.fetched_count} documents"
        )
        if search_result.duplicates_removed > 0:
            logger.info(f"Removed {search_result.duplicates_removed} duplicates")

        if not search_result.documents:
            # No results found
            session = self.storage.create_search_session(
                query=pubmed_query.query_string,
                natural_language_query=question,
                document_count=0,
                metadata={
                    "provider": provider.value,
                    "total_available": search_result.total_count,
                },
            )
            return session, []

        # Use documents directly from search result
        documents = search_result.documents

        if progress_callback:
            progress_callback(f"Caching {len(documents)} documents...")

        # Store documents with embeddings
        if documents:
            self.storage.add_documents(documents, embedding_function=self.embedder)

        # Create search session with provider metadata
        session = self.storage.create_search_session(
            query=pubmed_query.query_string,
            natural_language_query=question,
            document_count=len(documents),
            metadata={
                "provider": provider.value,
                "total_available": search_result.total_count,
                "pubmed_count": search_result.pubmed_count,
                "europepmc_count": search_result.europepmc_count,
                "duplicates_removed": search_result.duplicates_removed,
            },
        )

        # Record document-question associations for later retrieval
        if documents:
            doc_ids = [doc.id for doc in documents]
            self.storage.add_question_documents(
                question=question,
                document_ids=doc_ids,
                search_session_id=session.id,
            )

        logger.info(f"Cached {len(documents)} documents in session {session.id}")
        return session, documents

    def search_with_query(
        self,
        query_string: str,
        natural_language_query: Optional[str] = None,
        max_results: Optional[int] = None,
        provider: Optional[SearchProvider] = None,
        include_preprints: Optional[bool] = None,
        progress_callback: Optional[Callable[[str], None]] = None,
    ) -> tuple[SearchSession, list[LiteDocument]]:
        """
        Search with a pre-formatted query string.

        Use this when you want to bypass the LLM query conversion
        and use a specific query directly. The query can be in either
        PubMed or Europe PMC syntax and will be translated as needed.

        Args:
            query_string: Pre-formatted query string (PubMed or Europe PMC syntax)
            natural_language_query: Original question (for session tracking)
            max_results: Maximum results to fetch
            provider: Search provider (uses config default if None)
            include_preprints: Whether to include preprints (Europe PMC only)
            progress_callback: Optional callback for progress updates

        Returns:
            Tuple of (search session, list of documents)
        """
        max_results = max_results or self.config.search.max_results
        provider = provider or self.config.search.search_provider

        if progress_callback:
            provider_name = provider.display_name
            progress_callback(f"Searching {provider_name}...")

        # Execute search via unified search service
        search_result = self.search_service.search(
            query=query_string,
            max_results=max_results,
            provider=provider,
            include_preprints=include_preprints,
            progress_callback=progress_callback,
        )

        logger.info(f"Found {search_result.total_count} results")

        if not search_result.documents:
            session = self.storage.create_search_session(
                query=query_string,
                natural_language_query=natural_language_query or query_string,
                document_count=0,
                metadata={"provider": provider.value},
            )
            return session, []

        documents = search_result.documents

        if documents:
            if progress_callback:
                progress_callback(f"Caching {len(documents)} documents...")
            self.storage.add_documents(documents, embedding_function=self.embedder)

        session = self.storage.create_search_session(
            query=query_string,
            natural_language_query=natural_language_query or query_string,
            document_count=len(documents),
            metadata={
                "provider": provider.value,
                "total_available": search_result.total_count,
                "pubmed_count": search_result.pubmed_count,
                "europepmc_count": search_result.europepmc_count,
                "duplicates_removed": search_result.duplicates_removed,
            },
        )

        # Record document-question associations for later retrieval
        question = natural_language_query or query_string
        if documents:
            doc_ids = [doc.id for doc in documents]
            self.storage.add_question_documents(
                question=question,
                document_ids=doc_ids,
                search_session_id=session.id,
            )

        return session, documents

    def semantic_search(
        self,
        query: str,
        n_results: Optional[int] = None,
    ) -> list[LiteDocument]:
        """
        Search cached documents by semantic similarity.

        Uses embeddings to find documents similar to the query
        from the local SQLite cache.

        Args:
            query: Search query (natural language)
            n_results: Maximum results to return

        Returns:
            List of matching documents ordered by similarity
        """
        n_results = n_results or self.config.search.max_results

        return self.storage.search_documents(
            query=query,
            n_results=n_results,
            embedding_function=self.embedder,
        )

    def get_document(self, document_id: str) -> Optional[LiteDocument]:
        """
        Get a document by ID from the cache.

        Args:
            document_id: Document ID (e.g., "pmid-12345678")

        Returns:
            LiteDocument if found, None otherwise
        """
        return self.storage.get_document(document_id)

    def get_documents_by_pmids(self, pmids: list[str]) -> list[LiteDocument]:
        """
        Get multiple documents by their PMIDs.

        Args:
            pmids: List of PubMed IDs

        Returns:
            List of found documents (may be fewer than requested)
        """
        documents = []
        for pmid in pmids:
            doc = self.storage.get_document(f"pmid-{pmid}")
            if doc:
                documents.append(doc)
        return documents

    def _articles_to_documents(
        self,
        articles: list[ArticleMetadata],
    ) -> list[LiteDocument]:
        """
        Convert PubMed articles to LiteDocuments.

        Args:
            articles: List of ArticleMetadata from PubMed

        Returns:
            List of LiteDocument objects
        """
        documents = []

        for article in articles:
            # Skip articles without abstracts
            if not article.abstract:
                logger.debug(f"Skipping PMID {article.pmid} - no abstract")
                continue

            # Extract year from publication date
            year = None
            if article.publication_date:
                try:
                    # Try to extract year from date string
                    year = int(article.publication_date[:4])
                except (ValueError, TypeError):
                    pass

            doc = LiteDocument(
                id=f"pmid-{article.pmid}",
                title=article.title,
                abstract=article.abstract,
                authors=article.authors,
                year=year,
                journal=article.publication,
                doi=article.doi,
                pmid=article.pmid,
                pmc_id=article.pmc_id,
                source=DocumentSource.PUBMED,
                url=article.url,
                mesh_terms=article.mesh_terms,
            )
            documents.append(doc)

        return documents
