"""Sync and async mixin classes for SearchResource (backed by research API)."""

from __future__ import annotations

import logging
from typing import Any, Literal, overload

from ..._api.generated.research.research__api__research.models import (
    AsyncResearchResponse,
    JobStatus,
    ResearchRequestList,
    ResearchRequestRequest,
    ResearchResponse,
    SearchRequestRequest,
)
from ...exceptions import api_error_handler, async_api_error_handler
from .._polling import AsyncPollingMixin, PollingMixin

logger = logging.getLogger(__name__)


class _SearchMixin(PollingMixin):
    """Sync mixin for SearchResource methods (research API)."""

    @api_error_handler
    def search(
        self,
        query: str,
        *,
        num_results: int | None = None,
        gl: str | None = None,
        hl: str | None = None,
        allowed_domains: list[str] | None = None,
        blocked_domains: list[str] | None = None,
        wait: bool = True,
        poll_interval: float = 1.0,
        timeout: float = 60.0,
    ) -> AsyncResearchResponse | ResearchResponse:
        """
        Execute a Serper web search.

        Args:
            query: Search query
            num_results: Number of results (1-100, default: 10)
            gl: Country code (e.g. "us")
            hl: Language code (e.g. "en")
            allowed_domains: Only search these domains
            blocked_domains: Exclude these domains
            wait: If True, poll until completion and return results.
            poll_interval: Seconds between status checks (only if wait=True)
            timeout: Maximum wait time in seconds (only if wait=True)

        Returns:
            If wait=False: AsyncResearchResponse with job_id for polling
            If wait=True: ResearchResponse with search results
        """
        logger.debug("Search: %s (num_results=%s)", query, num_results)
        kwargs: dict[str, Any] = {"query": query}
        if num_results is not None:
            kwargs["num_results"] = num_results
        if gl is not None:
            kwargs["gl"] = gl
        if hl is not None:
            kwargs["hl"] = hl
        if allowed_domains is not None:
            kwargs["allowed_domains"] = allowed_domains
        if blocked_domains is not None:
            kwargs["blocked_domains"] = blocked_domains

        request = SearchRequestRequest(**kwargs)  # type: ignore[arg-type]
        job = self._api.search_create(request)
        logger.info("Search job: %s (status=%s)", job.request_uuid, job.status)

        if not wait or job.status != "queued":
            return job

        return self._poll_until_complete(
            job_uuid=job.request_uuid,
            get_status=lambda: self.job_status(job.request_uuid),
            get_result=lambda: self.results(job.request_uuid),
            error_message="Search job failed",
            poll_interval=poll_interval,
            timeout=timeout,
        )

    @api_error_handler
    def research(
        self,
        query: str,
        *,
        num_results: int | None = None,
        crawl_top_n: int | None = None,
        model: str | None = None,
        task_prompt: str | None = None,
        gl: str | None = None,
        hl: str | None = None,
        allowed_domains: list[str] | None = None,
        blocked_domains: list[str] | None = None,
        wait: bool = True,
        poll_interval: float = 2.0,
        timeout: float = 300.0,
    ) -> AsyncResearchResponse | ResearchResponse:
        """
        Execute full research pipeline: search → crawl → LLM summarize.

        Args:
            query: Research query
            num_results: Number of search results (1-100, default: 10)
            crawl_top_n: Number of top URLs to crawl (0-20, default: 3)
            model: LLM model for summarization (default: @cheap)
            task_prompt: Custom instructions for LLM
            gl: Country code
            hl: Language code
            allowed_domains: Only search these domains
            blocked_domains: Exclude these domains
            wait: If True, poll until completion and return results.
            poll_interval: Seconds between status checks (only if wait=True)
            timeout: Maximum wait time in seconds (only if wait=True)

        Returns:
            If wait=False: AsyncResearchResponse with job_id for polling
            If wait=True: ResearchResponse with full results
        """
        logger.debug("Research: %s (crawl_top_n=%s, model=%s)", query, crawl_top_n, model)
        kwargs: dict[str, Any] = {"query": query}
        if num_results is not None:
            kwargs["num_results"] = num_results
        if crawl_top_n is not None:
            kwargs["crawl_top_n"] = crawl_top_n
        if model is not None:
            kwargs["model"] = model
        if task_prompt is not None:
            kwargs["task_prompt"] = task_prompt
        if gl is not None:
            kwargs["gl"] = gl
        if hl is not None:
            kwargs["hl"] = hl
        if allowed_domains is not None:
            kwargs["allowed_domains"] = allowed_domains
        if blocked_domains is not None:
            kwargs["blocked_domains"] = blocked_domains

        request = ResearchRequestRequest(**kwargs)  # type: ignore[arg-type]
        job = self._api.research_create(request)
        logger.info("Research job: %s (status=%s)", job.request_uuid, job.status)

        if not wait or job.status != "queued":
            return job

        return self._poll_until_complete(
            job_uuid=job.request_uuid,
            get_status=lambda: self.job_status(job.request_uuid),
            get_result=lambda: self.results(job.request_uuid),
            error_message="Research job failed",
            poll_interval=poll_interval,
            timeout=timeout,
        )

    @api_error_handler
    def list(self) -> list[ResearchRequestList]:
        """List research history (last 100 records)."""
        return self._api.list()

    @api_error_handler
    def job_status(self, uuid: str) -> JobStatus:
        """Get status of async job."""
        return self._api.status_retrieve(uuid)

    @api_error_handler
    def results(self, uuid: str) -> ResearchResponse:
        """Get full results from completed job."""
        return self._api.results_retrieve(uuid)

    @api_error_handler
    def cancel(self, uuid: str) -> None:
        """Cancel an in-progress job."""
        self._api.cancel_create(uuid)


class _AsyncSearchMixin(AsyncPollingMixin):
    """Async mixin for AsyncSearchResource methods (research API)."""

    @async_api_error_handler
    async def search(
        self,
        query: str,
        *,
        num_results: int | None = None,
        gl: str | None = None,
        hl: str | None = None,
        allowed_domains: list[str] | None = None,
        blocked_domains: list[str] | None = None,
        wait: bool = True,
        poll_interval: float = 1.0,
        timeout: float = 60.0,
    ) -> AsyncResearchResponse | ResearchResponse:
        """
        Execute a Serper web search.

        Args:
            query: Search query
            num_results: Number of results (1-100, default: 10)
            gl: Country code (e.g. "us")
            hl: Language code (e.g. "en")
            allowed_domains: Only search these domains
            blocked_domains: Exclude these domains
            wait: If True, poll until completion and return results.
            poll_interval: Seconds between status checks (only if wait=True)
            timeout: Maximum wait time in seconds (only if wait=True)

        Returns:
            If wait=False: AsyncResearchResponse with job_id for polling
            If wait=True: ResearchResponse with search results
        """
        logger.debug("Async search: %s (num_results=%s)", query, num_results)
        kwargs: dict[str, Any] = {"query": query}
        if num_results is not None:
            kwargs["num_results"] = num_results
        if gl is not None:
            kwargs["gl"] = gl
        if hl is not None:
            kwargs["hl"] = hl
        if allowed_domains is not None:
            kwargs["allowed_domains"] = allowed_domains
        if blocked_domains is not None:
            kwargs["blocked_domains"] = blocked_domains

        request = SearchRequestRequest(**kwargs)  # type: ignore[arg-type]
        job = await self._api.search_create(request)
        logger.info("Search job: %s (status=%s)", job.request_uuid, job.status)

        if not wait or job.status != "queued":
            return job

        return await self._poll_until_complete(
            job_uuid=job.request_uuid,
            get_status=lambda: self.job_status(job.request_uuid),
            get_result=lambda: self.results(job.request_uuid),
            error_message="Search job failed",
            poll_interval=poll_interval,
            timeout=timeout,
        )

    @async_api_error_handler
    async def research(
        self,
        query: str,
        *,
        num_results: int | None = None,
        crawl_top_n: int | None = None,
        model: str | None = None,
        task_prompt: str | None = None,
        gl: str | None = None,
        hl: str | None = None,
        allowed_domains: list[str] | None = None,
        blocked_domains: list[str] | None = None,
        wait: bool = True,
        poll_interval: float = 2.0,
        timeout: float = 300.0,
    ) -> AsyncResearchResponse | ResearchResponse:
        """
        Execute full research pipeline: search → crawl → LLM summarize.

        Args:
            query: Research query
            num_results: Number of search results (1-100, default: 10)
            crawl_top_n: Number of top URLs to crawl (0-20, default: 3)
            model: LLM model for summarization (default: @cheap)
            task_prompt: Custom instructions for LLM
            gl: Country code
            hl: Language code
            allowed_domains: Only search these domains
            blocked_domains: Exclude these domains
            wait: If True, poll until completion and return results.
            poll_interval: Seconds between status checks (only if wait=True)
            timeout: Maximum wait time in seconds (only if wait=True)

        Returns:
            If wait=False: AsyncResearchResponse with job_id for polling
            If wait=True: ResearchResponse with full results
        """
        logger.debug("Async research: %s (crawl_top_n=%s, model=%s)", query, crawl_top_n, model)
        kwargs: dict[str, Any] = {"query": query}
        if num_results is not None:
            kwargs["num_results"] = num_results
        if crawl_top_n is not None:
            kwargs["crawl_top_n"] = crawl_top_n
        if model is not None:
            kwargs["model"] = model
        if task_prompt is not None:
            kwargs["task_prompt"] = task_prompt
        if gl is not None:
            kwargs["gl"] = gl
        if hl is not None:
            kwargs["hl"] = hl
        if allowed_domains is not None:
            kwargs["allowed_domains"] = allowed_domains
        if blocked_domains is not None:
            kwargs["blocked_domains"] = blocked_domains

        request = ResearchRequestRequest(**kwargs)  # type: ignore[arg-type]
        job = await self._api.research_create(request)
        logger.info("Research job: %s (status=%s)", job.request_uuid, job.status)

        if not wait or job.status != "queued":
            return job

        return await self._poll_until_complete(
            job_uuid=job.request_uuid,
            get_status=lambda: self.job_status(job.request_uuid),
            get_result=lambda: self.results(job.request_uuid),
            error_message="Research job failed",
            poll_interval=poll_interval,
            timeout=timeout,
        )

    @async_api_error_handler
    async def list(self) -> list[ResearchRequestList]:
        """List research history (last 100 records)."""
        return await self._api.list()

    @async_api_error_handler
    async def job_status(self, uuid: str) -> JobStatus:
        """Get status of async job."""
        return await self._api.status_retrieve(uuid)

    @async_api_error_handler
    async def results(self, uuid: str) -> ResearchResponse:
        """Get full results from completed job."""
        return await self._api.results_retrieve(uuid)

    @async_api_error_handler
    async def cancel(self, uuid: str) -> None:
        """Cancel an in-progress job."""
        await self._api.cancel_create(uuid)


__all__ = [
    "_SearchMixin",
    "_AsyncSearchMixin",
]
