from typing import Any, Dict, List, Optional

import requests
from langchain_core.language_models import BaseLLM
from loguru import logger

from ...security.safe_requests import safe_post
from ..rate_limiting import RateLimitError
from ..search_engine_base import BaseSearchEngine


class TavilySearchEngine(BaseSearchEngine):
    """Tavily search engine implementation with two-phase approach"""

    # Mark as public search engine
    is_public = True
    # Mark as generic search engine (general web search)
    is_generic = True

    def __init__(
        self,
        max_results: int = 10,
        region: str = "US",
        time_period: str = "y",
        safe_search: bool = True,
        search_language: str = "English",
        api_key: Optional[str] = None,
        llm: Optional[BaseLLM] = None,
        include_full_content: bool = True,
        max_filtered_results: Optional[int] = None,
        search_depth: str = "basic",
        include_domains: Optional[List[str]] = None,
        exclude_domains: Optional[List[str]] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Initialize the Tavily search engine.

        Args:
            max_results: Maximum number of search results
            region: Region code for search results (not used by Tavily currently)
            time_period: Time period for search results (not used by Tavily currently)
            safe_search: Whether to enable safe search (not used by Tavily currently)
            search_language: Language for search results (not used by Tavily currently)
            api_key: Tavily API key (can also be set via LDR_SEARCH_ENGINE_WEB_TAVILY_API_KEY env var or in UI settings)
            llm: Language model for relevance filtering
            include_full_content: Whether to include full webpage content in results
            max_filtered_results: Maximum number of results to keep after filtering
            search_depth: "basic" or "advanced" - controls search quality vs speed
            include_domains: List of domains to include in search
            exclude_domains: List of domains to exclude from search
            settings_snapshot: Settings snapshot for thread context
            **kwargs: Additional parameters (ignored but accepted for compatibility)
        """
        # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            include_full_content=include_full_content,
            settings_snapshot=settings_snapshot,
        )
        self.search_depth = search_depth
        self.include_domains = include_domains or []
        self.exclude_domains = exclude_domains or []

        # Get API key - check params, settings, or env vars
        tavily_api_key = self._resolve_api_key(
            api_key,
            "search.engine.web.tavily.api_key",
            engine_name="Tavily",
            settings_snapshot=settings_snapshot,
        )

        self.api_key = tavily_api_key
        self.base_url = "https://api.tavily.com"

        # If full content is requested, initialize FullSearchResults
        if include_full_content:
            # Create a simple wrapper for Tavily API calls
            class TavilyWrapper:
                def __init__(self, parent):
                    self.parent = parent

                def run(self, query):
                    return self.parent._get_previews(query)

            self._init_full_search(
                web_search=TavilyWrapper(self),
                language=search_language,
                max_results=max_results,
                region=region,
                time_period=time_period,
                safe_search="moderate" if safe_search else "off",
            )

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information from Tavily Search.

        Args:
            query: The search query

        Returns:
            List of preview dictionaries
        """
        logger.info("Getting search results from Tavily")

        try:
            # Prepare the request payload
            payload = {
                "api_key": self.api_key,
                "query": query[:400],  # Limit query length
                "search_depth": self.search_depth,
                "max_results": min(
                    20, self.max_results
                ),  # Tavily has a max limit
                "include_answer": False,  # We don't need the AI answer
                "include_images": False,  # We don't need images
                "include_raw_content": self.include_full_content,  # Get content if requested
            }

            # Add domain filters if specified
            if self.include_domains:
                payload["include_domains"] = self.include_domains
            if self.exclude_domains:
                payload["exclude_domains"] = self.exclude_domains

            # Apply rate limiting before request
            self._last_wait_time = self.rate_tracker.apply_rate_limit(
                self.engine_type
            )

            # Make the API request
            response = safe_post(
                f"{self.base_url}/search",
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=30,
            )

            # Check for rate limits
            self._raise_if_rate_limit(response.status_code)

            response.raise_for_status()

            # Parse the response
            data = response.json()
            results = data.get("results", [])

            # Format results as previews
            previews = []
            for i, result in enumerate(results):
                url = self._clean_result_url(result.get("url"))
                preview = {
                    "id": url or str(i),  # Use URL as ID
                    "title": result.get("title", ""),
                    "link": url,
                    "snippet": result.get(
                        "content", ""
                    ),  # Tavily calls it "content"
                    "displayed_link": url,
                    "position": i,
                }

                # Store full Tavily result for later
                preview["_full_result"] = result

                previews.append(preview)

            # Store the previews for potential full content retrieval
            self._search_results = previews

            return previews

        except RateLimitError:
            raise  # Re-raise rate limit errors
        except requests.exceptions.RequestException as e:
            safe_msg = self._scrub_error(e)
            logger.warning(f"Error getting Tavily results: {safe_msg}")
            self._raise_if_rate_limit(e)
            return []
        except Exception as e:
            safe_msg = self._scrub_error(e)
            logger.warning(
                f"Unexpected error getting Tavily results: {safe_msg}"
            )
            return []

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant search results.
        Extends base implementation to include Tavily's raw_content.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content if available
        """
        results = super()._get_full_content(relevant_items)

        # If Tavily provided raw_content and full content is requested, use it
        if self.include_full_content:
            for result in results:
                if "raw_content" in result:
                    result["content"] = result.get(
                        "raw_content", result.get("content", "")
                    )

        return results
