"""
Remote search tools that call HTTP APIs.

These tools are designed to work with AWS Lambda (or any compatible) search endpoints.
They use aiohttp to make async HTTP requests.
"""

import json
from typing import Annotated, Any

from aiohttp import ClientSession, ClientTimeout

from lm_deluge.tool import Tool


class RemoteSearchManager:
    """
    Search and fetch tools that call a remote HTTP API.

    Provides three tools:
    - search: Search documents and get previews
    - fetch: Get full document contents by ID
    - batch: Execute multiple search/fetch operations in parallel

    This is simpler than running an MCP server locally - just make HTTP calls
    to your deployed Lambda (or any compatible API).

    Args:
        base_url: Base URL of the search API (e.g., "https://xxx.lambda-url.us-west-2.on.aws")
        indices: List of available index names (for documentation)
        index_descriptions: Optional dict mapping index names to descriptions
        search_tool_name: Name for the search tool (default: "search")
        fetch_tool_name: Name for the fetch tool (default: "fetch")
        batch_tool_name: Name for the batch tool (default: "batch")
        timeout: Request timeout in seconds (default: 30)
        max_results: Maximum results to return (default: 10)

    Example:
        ```python
        manager = RemoteSearchManager(
            base_url="https://s4wc2dsepgfnnwopoxe2unh7pu0xdyds.lambda-url.us-west-2.on.aws",
            indices=["municipal_code", "website", "housing_element"],
            index_descriptions={
                "municipal_code": "City municipal code and ordinances",
                "website": "City website content",
                "housing_element": "Housing element documents",
            },
        )
        tools = manager.get_tools()
        ```
    """

    def __init__(
        self,
        base_url: str,
        indices: list[str],
        *,
        index_descriptions: dict[str, str] | None = None,
        search_tool_name: str = "search",
        fetch_tool_name: str = "fetch",
        batch_tool_name: str = "batch",
        timeout: int = 30,
        max_results: int = 10,
    ):
        # Remove trailing slash from base_url
        self.base_url = base_url.rstrip("/")
        self.indices = indices
        self.index_descriptions = index_descriptions or {}
        self.search_tool_name = search_tool_name
        self.fetch_tool_name = fetch_tool_name
        self.batch_tool_name = batch_tool_name
        self.timeout = ClientTimeout(total=timeout)
        self.max_results = max_results
        self._tools: list[Tool] | None = None

    def _build_index_docs(self) -> str:
        """Build documentation string for available indices."""
        lines = ["Available indices:"]
        for idx in self.indices:
            desc = self.index_descriptions.get(idx, "")
            if desc:
                lines.append(f"  - {idx}: {desc}")
            else:
                lines.append(f"  - {idx}")
        return "\n".join(lines)

    async def _search(
        self,
        index: Annotated[str, "Name of the index to search"],
        query: Annotated[str, "Search query to find relevant documents"],
        limit: Annotated[int, "Maximum number of results to return"] = 5,
        include_content: Annotated[
            bool, "If true, include full document content in results"
        ] = False,
    ) -> str:
        """
        Search for documents matching the query.

        Returns document IDs, relevance scores, and previews.
        Use the fetch tool to get full document contents.
        """
        try:
            if index not in self.indices:
                return json.dumps(
                    {
                        "status": "error",
                        "error": f"Unknown index: {index}. Available: {self.indices}",
                    }
                )

            data = {
                "index": index,
                "queries": [query],
                "limit": min(limit, self.max_results),
                "include_content": include_content,
            }

            async with ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/search",
                    json=data,
                    timeout=self.timeout,
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        return json.dumps(
                            {
                                "status": "error",
                                "error": f"API error: {response.status} - {error_text}",
                            }
                        )
                    result = await response.json()

            # Format results nicely
            results = result.get("results", [])
            formatted = []
            for r in results:
                formatted_result = {
                    "id": r.get("id"),
                    "score": round(r.get("score", 0), 4),
                    "title": r.get("title"),
                    "description": r.get("description"),
                    "preview": r.get("preview"),
                }
                # Include full content if requested and available
                if include_content and "content" in r:
                    formatted_result["content"] = r["content"]
                formatted.append(formatted_result)

            return json.dumps(
                {
                    "status": "success",
                    "index": index,
                    "query": query,
                    "num_results": len(formatted),
                    "results": formatted,
                },
                indent=2,
            )

        except Exception as e:
            return json.dumps({"status": "error", "error": str(e)})

    async def _fetch(
        self,
        index: Annotated[str, "Name of the index to fetch from"],
        document_ids: Annotated[
            list[str], "List of document IDs to fetch full contents for"
        ],
        filters: Annotated[
            list[dict[str, Any]] | None,
            "Optional list of filter specs (one per document). Each filter can have: "
            "start_line/end_line for line range, or grep_pattern/context_lines for regex search",
        ] = None,
    ) -> str:
        """
        Fetch the full contents of documents by their IDs.

        Use search first to find relevant document IDs.
        """
        try:
            if index not in self.indices:
                return json.dumps(
                    {
                        "status": "error",
                        "error": f"Unknown index: {index}. Available: {self.indices}",
                    }
                )

            data: dict[str, Any] = {
                "index": index,
                "document_ids": document_ids,
            }
            if filters is not None:
                data["filters"] = filters

            async with ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/fetch",
                    json=data,
                    timeout=self.timeout,
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        return json.dumps(
                            {
                                "status": "error",
                                "error": f"API error: {response.status} - {error_text}",
                            }
                        )
                    result = await response.json()

            return json.dumps(
                {
                    "status": "success",
                    "index": index,
                    "found": len(result.get("found_ids", [])),
                    "missing": len(result.get("missing_ids", [])),
                    "documents": result.get("documents", []),
                    "missing_ids": result.get("missing_ids", []),
                },
                indent=2,
            )

        except Exception as e:
            return json.dumps({"status": "error", "error": str(e)})

    async def _batch(
        self,
        commands: Annotated[
            list[dict[str, Any]],
            "List of commands to execute. Each command has 'name' ('search' or 'fetch'), "
            "'index', and either 'queries'+'limit' (for search) or 'document_ids'+'filters' (for fetch)",
        ],
    ) -> str:
        """
        Execute multiple search/fetch operations in a single request.

        This is more efficient than making multiple separate requests.
        """
        try:
            # Validate commands reference known indices
            for cmd in commands:
                if cmd.get("index") not in self.indices:
                    return json.dumps(
                        {
                            "status": "error",
                            "error": f"Unknown index: {cmd.get('index')}. Available: {self.indices}",
                        }
                    )

            data = {"commands": commands}

            async with ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/batch",
                    json=data,
                    timeout=self.timeout,
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        return json.dumps(
                            {
                                "status": "error",
                                "error": f"API error: {response.status} - {error_text}",
                            }
                        )
                    result = await response.json()

            return json.dumps(
                {
                    "status": "success",
                    "num_commands": len(commands),
                    "results": result.get("results", []),
                },
                indent=2,
            )

        except Exception as e:
            return json.dumps({"status": "error", "error": str(e)})

    def get_tools(self) -> list[Tool]:
        """Return the search, fetch, and batch tools."""
        if self._tools is not None:
            return self._tools

        index_docs = self._build_index_docs()

        search_tool = Tool.from_function(self._search, name=self.search_tool_name)
        fetch_tool = Tool.from_function(self._fetch, name=self.fetch_tool_name)
        batch_tool = Tool.from_function(self._batch, name=self.batch_tool_name)

        # Update descriptions with index info
        search_tool = search_tool.model_copy(
            update={
                "description": (
                    f"Search for documents matching a query. "
                    f"Returns document IDs, relevance scores, and previews. "
                    f"Use the fetch tool to get full document contents.\n\n"
                    f"{index_docs}"
                )
            }
        )

        fetch_tool = fetch_tool.model_copy(
            update={
                "description": (
                    f"Fetch the full contents of documents by their IDs. "
                    f"Use after searching to get complete document text. "
                    f"Supports optional filters for line ranges or grep patterns.\n\n"
                    f"{index_docs}"
                )
            }
        )

        batch_tool = batch_tool.model_copy(
            update={
                "description": (
                    f"Execute multiple search/fetch operations in parallel. "
                    f"More efficient than making separate requests. "
                    f"Each command has 'name' ('search' or 'fetch'), 'index', and operation-specific fields.\n\n"
                    f"{index_docs}"
                )
            }
        )

        self._tools = [search_tool, fetch_tool, batch_tool]
        return self._tools

    async def search(
        self,
        index: str,
        query: str,
        limit: int = 5,
        include_content: bool = False,
    ) -> dict[str, Any]:
        """
        Direct search method for programmatic use.

        Returns parsed JSON response instead of string.
        """
        result = await self._search(index, query, limit, include_content)
        return json.loads(result)

    async def fetch(
        self,
        index: str,
        document_ids: list[str],
        filters: list[dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        """
        Direct fetch method for programmatic use.

        Returns parsed JSON response instead of string.
        """
        result = await self._fetch(index, document_ids, filters)
        return json.loads(result)

    async def batch(self, commands: list[dict[str, Any]]) -> dict[str, Any]:
        """
        Direct batch method for programmatic use.

        Returns parsed JSON response instead of string.
        """
        result = await self._batch(commands)
        return json.loads(result)

    async def warm_up(self) -> bool:
        """
        Hit the health endpoint to warm up the Lambda.

        Call this when starting up your agent to ensure
        subsequent search/fetch calls are fast.

        Returns:
            True if the API is healthy, False otherwise.
        """
        try:
            async with ClientSession() as session:
                async with session.get(
                    f"{self.base_url}/health",
                    timeout=self.timeout,
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        return result.get("status") == "healthy"
                    return False
        except Exception:
            return False


__all__ = ["RemoteSearchManager"]
