"""Web search tool for Promethian.

Supports multiple search backends:
- API (primary): Uses Promethian API server with Tavily. Requires PROMETHIAN_AUTH_TOKEN.
- DuckDuckGo (fallback): No API key required, but only returns snippets.

When PROMETHIAN_AUTH_TOKEN is set, searches go through the API server.
Otherwise falls back to local DuckDuckGo search.
"""

from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass, field
from typing import Any, Protocol


@dataclass
class SearchResult:
    """A single search result."""
    title: str
    url: str
    snippet: str
    content: str | None = None  # Full extracted content (API/Tavily only)
    score: float | None = None  # Relevance score (API/Tavily only)


@dataclass
class WebSearchResult:
    """Result from a web search."""
    success: bool
    results: list[SearchResult] | None = None
    error: str | None = None
    query: str = ""
    source: str = ""  # "api", "tavily", or "ddg"


class SearchBackend(Protocol):
    """Protocol for search backends."""

    async def search(
        self,
        query: str,
        max_results: int | None = None,
    ) -> WebSearchResult:
        ...

    async def news_search(
        self,
        query: str,
        max_results: int | None = None,
    ) -> WebSearchResult:
        ...

    def get_tool_schemas(self) -> list[dict]:
        ...


class APISearchTool:
    """
    Web search using Promethian API server (recommended).

    The API server uses Tavily for high-quality search with:
    - Full content extraction from pages (not just snippets)
    - Relevance scoring
    - AI-optimized results

    Requires PROMETHIAN_AUTH_TOKEN environment variable.
    """

    def __init__(self, api_url: str | None = None, auth_token: str | None = None, max_results: int = 5):
        self.api_url = (
            api_url
            or os.environ.get("PROMETHIAN_API_URL")
            or "https://promethian-api.fly.dev"
        )
        self.auth_token = auth_token or os.environ.get("PROMETHIAN_AUTH_TOKEN") or os.environ.get("PROMETHIAN_API_KEY")
        if not self.auth_token:
            raise ValueError("PROMETHIAN_AUTH_TOKEN is required for API search")
        self.max_results = max_results
        self._client = None

    async def _ensure_client(self):
        """Ensure API client is available."""
        if self._client is None:
            from promethian.api.client import PromethianAPIClient
            self._client = PromethianAPIClient(
                api_url=self.api_url,
                api_key=self.auth_token,
            )
        return self._client

    async def close(self):
        """Close the API client."""
        if self._client:
            await self._client.close()
            self._client = None

    async def search(
        self,
        query: str,
        max_results: int | None = None,
        include_content: bool = True,
    ) -> WebSearchResult:
        """
        Search the web using the Promethian API.

        Args:
            query: The search query
            max_results: Maximum number of results (default: 5)
            include_content: Whether to include full page content (default: True)

        Returns:
            WebSearchResult with list of results including content
        """
        try:
            client = await self._ensure_client()
            num_results = max_results or self.max_results

            response = await client.web_search(
                query=query,
                max_results=num_results,
                search_type="general",
            )

            if not response.get("success"):
                return WebSearchResult(
                    success=False,
                    error=response.get("error", "Unknown API error"),
                    query=query,
                    source="api",
                )

            results = []
            for r in response.get("results", []):
                results.append(SearchResult(
                    title=r.get("title", ""),
                    url=r.get("url", ""),
                    snippet=r.get("snippet", ""),
                    content=r.get("content"),
                    score=r.get("score"),
                ))

            return WebSearchResult(
                success=True,
                results=results,
                query=query,
                source="api",
            )

        except Exception as e:
            return WebSearchResult(
                success=False,
                error=str(e),
                query=query,
                source="api",
            )

    async def news_search(
        self,
        query: str,
        max_results: int | None = None,
    ) -> WebSearchResult:
        """
        Search for recent news using the Promethian API.

        Args:
            query: The search query
            max_results: Maximum number of results (default: 5)

        Returns:
            WebSearchResult with list of news results
        """
        try:
            client = await self._ensure_client()
            num_results = max_results or self.max_results

            response = await client.web_search(
                query=query,
                max_results=num_results,
                search_type="news",
            )

            if not response.get("success"):
                return WebSearchResult(
                    success=False,
                    error=response.get("error", "Unknown API error"),
                    query=query,
                    source="api",
                )

            results = []
            for r in response.get("results", []):
                results.append(SearchResult(
                    title=r.get("title", ""),
                    url=r.get("url", ""),
                    snippet=r.get("snippet", ""),
                    content=r.get("content"),
                    score=r.get("score"),
                ))

            return WebSearchResult(
                success=True,
                results=results,
                query=query,
                source="api",
            )

        except Exception as e:
            return WebSearchResult(
                success=False,
                error=str(e),
                query=query,
                source="api",
            )

    def get_tool_schemas(self) -> list[dict]:
        """Get Claude tool schemas for API search actions."""
        return [
            {
                "name": "web_search",
                "description": (
                    "Search the web using AI-optimized search. Returns titles, URLs, snippets, "
                    "AND full extracted page content. Use this to find and read information, "
                    "documentation, articles, or any web content. The content field contains "
                    "the full text extracted from each page."
                ),
                "input_schema": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query"
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum number of results (default 5, max 10)",
                            "default": 5
                        }
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "web_news_search",
                "description": (
                    "Search for recent news articles using AI-optimized search. Returns titles, "
                    "URLs, snippets, and full article content. Use this for current events "
                    "or recent developments."
                ),
                "input_schema": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The news search query"
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum number of results (default 5, max 10)",
                            "default": 5
                        }
                    },
                    "required": ["query"]
                }
            },
        ]

    async def execute_tool(self, name: str, params: dict) -> WebSearchResult:
        """Execute a web search tool by name."""
        max_results = min(params.get("max_results", 5), 10)  # Cap at 10

        if name == "web_search":
            return await self.search(
                query=params["query"],
                max_results=max_results,
            )
        elif name == "web_news_search":
            return await self.news_search(
                query=params["query"],
                max_results=max_results,
            )
        else:
            return WebSearchResult(
                success=False,
                error=f"Unknown web search tool: {name}",
                query=params.get("query", ""),
                source="api",
            )


class DDGSearchTool:
    """
    Web search tool using DuckDuckGo (fallback).

    Provides:
    - search(query) - Search the web and get results
    - news_search(query) - Search for recent news

    No API key required, but only returns snippets (not full content).
    """

    def __init__(self, max_results: int = 5):
        self.max_results = max_results
        self._ddgs = None

    def _ensure_ddgs(self):
        """Ensure DuckDuckGo search is available."""
        if self._ddgs is None:
            try:
                # Try new package name first
                from ddgs import DDGS
                self._ddgs = DDGS()
            except ImportError:
                try:
                    # Fall back to old package name
                    from duckduckgo_search import DDGS
                    self._ddgs = DDGS()
                except ImportError:
                    raise RuntimeError(
                        "ddgs not installed. Run: pip install ddgs"
                    )
        return self._ddgs

    async def search(
        self,
        query: str,
        max_results: int | None = None,
        region: str = "wt-wt",
    ) -> WebSearchResult:
        """
        Search the web using DuckDuckGo.

        Args:
            query: The search query
            max_results: Maximum number of results (default: 5)
            region: Region for results (default: worldwide)

        Returns:
            WebSearchResult with list of results (snippets only, no full content)
        """
        try:
            ddgs = self._ensure_ddgs()
            num_results = max_results or self.max_results

            # Run in thread pool since ddgs is synchronous
            loop = asyncio.get_event_loop()
            raw_results = await loop.run_in_executor(
                None,
                lambda: list(ddgs.text(query, max_results=num_results, region=region))
            )

            results = [
                SearchResult(
                    title=r.get("title", ""),
                    url=r.get("href", r.get("link", "")),
                    snippet=r.get("body", r.get("snippet", "")),
                    content=None,  # DDG doesn't provide full content
                    score=None,
                )
                for r in raw_results
            ]

            return WebSearchResult(
                success=True,
                results=results,
                query=query,
                source="ddg",
            )

        except Exception as e:
            return WebSearchResult(
                success=False,
                error=str(e),
                query=query,
                source="ddg",
            )

    async def news_search(
        self,
        query: str,
        max_results: int | None = None,
        region: str = "wt-wt",
    ) -> WebSearchResult:
        """
        Search for recent news using DuckDuckGo.

        Args:
            query: The search query
            max_results: Maximum number of results (default: 5)
            region: Region for results (default: worldwide)

        Returns:
            WebSearchResult with list of news results
        """
        try:
            ddgs = self._ensure_ddgs()
            num_results = max_results or self.max_results

            # Run in thread pool since ddgs is synchronous
            loop = asyncio.get_event_loop()
            raw_results = await loop.run_in_executor(
                None,
                lambda: list(ddgs.news(query, max_results=num_results, region=region))
            )

            results = [
                SearchResult(
                    title=r.get("title", ""),
                    url=r.get("url", r.get("link", "")),
                    snippet=r.get("body", r.get("excerpt", "")),
                    content=None,
                    score=None,
                )
                for r in raw_results
            ]

            return WebSearchResult(
                success=True,
                results=results,
                query=query,
                source="ddg",
            )

        except Exception as e:
            return WebSearchResult(
                success=False,
                error=str(e),
                query=query,
                source="ddg",
            )

    def get_tool_schemas(self) -> list[dict]:
        """Get Claude tool schemas for web search actions."""
        return [
            {
                "name": "web_search",
                "description": (
                    "Search the web using DuckDuckGo. Returns titles, URLs, and snippets "
                    "for the top results. Use this to find information, documentation, "
                    "tutorials, or any web content. Note: Only returns snippets, not full "
                    "page content."
                ),
                "input_schema": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query"
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum number of results (default 5, max 10)",
                            "default": 5
                        }
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "web_news_search",
                "description": (
                    "Search for recent news articles using DuckDuckGo. Returns titles, "
                    "URLs, and snippets for recent news. Use this for current events "
                    "or recent developments."
                ),
                "input_schema": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The news search query"
                        },
                        "max_results": {
                            "type": "integer",
                            "description": "Maximum number of results (default 5, max 10)",
                            "default": 5
                        }
                    },
                    "required": ["query"]
                }
            },
        ]

    async def execute_tool(self, name: str, params: dict) -> WebSearchResult:
        """Execute a web search tool by name."""
        max_results = min(params.get("max_results", 5), 10)  # Cap at 10

        if name == "web_search":
            return await self.search(
                query=params["query"],
                max_results=max_results,
            )
        elif name == "web_news_search":
            return await self.news_search(
                query=params["query"],
                max_results=max_results,
            )
        else:
            return WebSearchResult(
                success=False,
                error=f"Unknown web search tool: {name}",
                query=params.get("query", ""),
                source="ddg",
            )


# Type alias for the unified search tool
WebSearchTool = APISearchTool | DDGSearchTool

# Singleton web search tool instance
_web_search: WebSearchTool | None = None


def get_web_search(
    max_results: int = 5,
    auth_token: str | None = None,
) -> WebSearchTool:
    """
    Get the global web search tool instance.

    Uses Promethian API if PROMETHIAN_AUTH_TOKEN is set (server-side Tavily),
    otherwise falls back to DuckDuckGo.

    Args:
        max_results: Default number of results to return
        auth_token: Optional auth token (uses env var if not provided)

    Returns:
        APISearchTool if auth token available, else DDGSearchTool
    """
    global _web_search
    if _web_search is None:
        token = auth_token or os.environ.get("PROMETHIAN_AUTH_TOKEN") or os.environ.get("PROMETHIAN_API_KEY")
        if token:
            try:
                _web_search = APISearchTool(auth_token=token, max_results=max_results)
            except Exception:
                # Fall back to DDG if API init fails
                _web_search = DDGSearchTool(max_results=max_results)
        else:
            _web_search = DDGSearchTool(max_results=max_results)
    return _web_search


def reset_web_search() -> None:
    """Reset the global web search tool instance."""
    global _web_search
    _web_search = None
