"""External tools — web search and documentation retrieval.

Provides ``web_search`` and ``read_target_docs`` tools for
gathering external context about the target agent.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from ziran.infrastructure.llm.base import BaseLLMClient

logger = logging.getLogger(__name__)


async def web_search(
    llm_client: BaseLLMClient,
    *,
    query: str,
    context: str = "",
) -> dict[str, Any]:
    """Search for CVEs, known vulnerabilities, and framework-specific weaknesses.

    Uses the LLM to synthesize knowledge about known attack patterns
    relevant to the query. This is a knowledge-based search, not a
    live web search — the LLM draws on its training data.

    Args:
        llm_client: LLM client for knowledge synthesis.
        query: Search query (e.g. 'LangChain prompt injection CVEs').
        context: Additional context about the target.

    Returns:
        Synthesized information about relevant vulnerabilities.
    """
    prompt = (
        "You are a cybersecurity researcher. Based on your knowledge, provide "
        "information about the following security topic as it relates to AI agents:\n\n"
        f"Query: {query}\n"
    )
    if context:
        prompt += f"Context: {context}\n"
    prompt += (
        "\nProvide:\n"
        "1. Known CVEs or vulnerabilities related to this topic\n"
        "2. Common attack patterns\n"
        "3. Affected frameworks or components\n"
        "4. Severity assessment\n\n"
        "Be factual and cite specific CVE IDs where applicable."
    )

    try:
        response = await llm_client.complete(
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
            max_tokens=1500,
        )
        return {
            "query": query,
            "results": response.content,
            "token_usage": {
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
                "total_tokens": response.total_tokens,
            },
        }
    except Exception as e:
        logger.warning("Web search failed: %s", e)
        return {"query": query, "results": "", "error": str(e)}


async def read_target_docs(
    *,
    url: str,
    llm_client: BaseLLMClient | None = None,
) -> dict[str, Any]:
    """Fetch and summarize the target agent's documentation.

    Downloads the content at the given URL and optionally uses
    the LLM to summarize it for the pentesting agent's context.

    Args:
        url: URL of the target documentation.
        llm_client: Optional LLM client for summarization.

    Returns:
        Document content and optional summary.
    """
    import httpx

    try:
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            content = resp.text[:10000]  # Limit content size
    except Exception as e:
        logger.warning("Failed to fetch target docs from %s: %s", url, e)
        return {"url": url, "content": "", "error": str(e)}

    result: dict[str, Any] = {
        "url": url,
        "content_length": len(content),
        "content_preview": content[:2000],
    }

    if llm_client:
        try:
            response = await llm_client.complete(
                messages=[
                    {
                        "role": "user",
                        "content": (
                            "Summarize this agent/API documentation, focusing on:\n"
                            "1. Agent capabilities and tools\n"
                            "2. Authentication and access controls\n"
                            "3. Input/output handling\n"
                            "4. Potential attack surface\n\n"
                            f"Documentation:\n{content[:5000]}"
                        ),
                    }
                ],
                temperature=0.2,
                max_tokens=1024,
            )
            result["summary"] = response.content
        except Exception as e:
            logger.warning("Doc summarization failed: %s", e)
            result["summary_error"] = str(e)

    return result
