# Copyright (c) 2025 Microsoft Corporation.
# Licensed under the MIT License

"""Base handler for keyword-based retrieval from platform search APIs.

KeywordRetrievalHandler extends DefaultAskHandler and overrides only the
retrieval step. The full query pipeline (decontextualization, elicitation,
ranking, summarization) is inherited from DefaultAskHandler.

Subclasses implement _retrieve_items() to query their platform's search API.
"""

import asyncio
import logging
from abc import abstractmethod

from nlweb_core.config import override_scoring_provider
from nlweb_core.handler import DefaultAskHandler, OutputMethod
from nlweb_core.protocol.models import AskRequest
from nlweb_core.retrieved_item import RetrievedItem

from nlweb_platform_handlers.query_rewrite import rewrite_query

logger = logging.getLogger(__name__)


class KeywordRetrievalHandler(DefaultAskHandler):
    """Base handler for platforms with their own search API.

    Instead of querying a pre-indexed vector database, this handler:
    1. Rewrites the natural language query into keyword queries via LLM
    2. Calls the platform's search API with those keywords (via _retrieve_items)
    3. Passes results through the standard ranking pipeline

    Everything else (decontextualization, elicitation, meta, summarization,
    conversation saving) is inherited from DefaultAskHandler.

    Subclasses must implement:
        _retrieve_items(query, site, num_results) -> list[RetrievedItem]
    """

    async def _run_query_body(
        self,
        request: AskRequest,
        output_method: OutputMethod,
        site_config: dict[str, str],
    ) -> list[dict]:
        """Execute the query body using keyword-based retrieval.

        Overrides DefaultAskHandler._run_query_body() to use platform search
        instead of vector DB search.
        """
        from nlweb_core.ranking import Ranking

        query = request.query.effective_query

        # Rewrite query into keyword queries for platform search
        rewritten_queries = await rewrite_query(query)
        logger.info(f"Rewritten queries for platform search: {rewritten_queries}")

        # Retrieve items using rewritten queries
        retrieved_items = await self._retrieve_with_queries(
            queries=rewritten_queries,
            site=request.query.site,
            num_results=request.query.num_results,
        )

        # Use LLM-based scoring for platform-retrieved items.
        # Platform search APIs return pre-filtered results that need LLM evaluation
        # rather than the default Pi Labs scorer which under-scores these items.
        with override_scoring_provider("default", "4.1-mini"):
            final_ranked_answers = await Ranking().rank(
                items=retrieved_items,
                query_text=query,
                item_type=site_config["item_type"],
                max_results=request.query.max_results,
                min_score=request.query.min_score,
                site=request.query.site,
                start_num=self._get_result_offset(request),
            )

        await self._send_results(output_method, final_ranked_answers)
        return final_ranked_answers

    async def _retrieve_with_queries(
        self,
        queries: list[str],
        site: str | None,
        num_results: int,
    ) -> list[RetrievedItem]:
        """Retrieve items using multiple queries in parallel.

        Distributes num_results across queries, deduplicates by URL.
        """
        if len(queries) == 1:
            return await self._retrieve_items(
                query=queries[0],
                site=site,
                num_results=num_results,
            )

        # Distribute results evenly across queries
        results_per_query = max(1, num_results // len(queries))
        remainder = num_results % len(queries)

        tasks = []
        for i, query in enumerate(queries):
            query_results = results_per_query + (1 if i < remainder else 0)
            tasks.append(
                self._retrieve_items(
                    query=query,
                    site=site,
                    num_results=query_results,
                )
            )

        all_results = await asyncio.gather(*tasks, return_exceptions=True)

        # Combine results, filtering out errors and deduplicating by URL
        combined: list[RetrievedItem] = []
        seen_urls: set[str] = set()

        for i, result in enumerate(all_results):
            if isinstance(result, BaseException):
                logger.warning(f"Search failed for query '{queries[i]}': {result}")
                continue
            for item in result:
                if item.url not in seen_urls:
                    seen_urls.add(item.url)
                    combined.append(item)

        return combined[:num_results]

    @abstractmethod
    async def _retrieve_items(
        self,
        query: str,
        site: str | None,
        num_results: int,
    ) -> list[RetrievedItem]:
        """Retrieve items using platform-specific keyword search.

        Subclasses must implement this to query their platform's search API.

        Args:
            query: The search query string (already rewritten to keywords).
            site: The site domain to search.
            num_results: Maximum number of results to return.

        Returns:
            List of RetrievedItem objects with schema.org formatted data.
        """
        ...
