"""Built-in tool: Perplexity web search.

Requires environment variable:
- ``PERPLEXITY_API_KEY``: API key from Perplexity.

Uses the Perplexity Chat Completions API with an online model
to perform grounded web search. Returns the answer with citations.

See https://docs.perplexity.ai/
"""

from __future__ import annotations

import logging
import os

# Any: the OpenAI tool schema is a heterogeneous dict with string
# keys and mixed value types (str, dict, list).
from typing import Any

import httpx

_logger = logging.getLogger(__name__)

_DEFAULT_PERPLEXITY_URL = "https://api.perplexity.ai/chat/completions"

# Perplexity model optimized for web search with citations.
_PERPLEXITY_MODEL = "sonar"


def _perplexity_url() -> str:
    """Resolve the Perplexity URL; ``OMNIGENT_PERPLEXITY_BASE_URL`` overrides for tests."""
    return os.environ.get("OMNIGENT_PERPLEXITY_BASE_URL", _DEFAULT_PERPLEXITY_URL)


def _search_perplexity(
    query: str,
    config: dict[str, str],
) -> str:
    """
    Call the Perplexity Chat Completions API with an online model.

    :param query: The search query string.
    :param config: Spec-level config; checked for ``api_key``
        before falling back to the env var.
    :returns: The answer text with citations, or an error message.
    """
    api_key = config.get("api_key")
    if not api_key:
        return "Error: api_key must be provided in the web_search config in config.yaml."
    try:
        resp = httpx.post(
            _perplexity_url(),
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": _PERPLEXITY_MODEL,
                "messages": [
                    {"role": "user", "content": query},
                ],
            },
            timeout=30.0,
        )
        resp.raise_for_status()
    except httpx.HTTPStatusError as exc:
        return f"Perplexity search error: HTTP {exc.response.status_code}"
    except (httpx.ConnectError, httpx.TimeoutException) as exc:
        return f"Perplexity search error: {exc}"

    return _format_response(resp.json())


def _format_response(data: dict[str, Any]) -> str:
    """
    Extract the answer and citations from Perplexity's response.

    :param data: The parsed JSON response from Perplexity.
    :returns: The answer text followed by numbered citations.
    """
    choices = data.get("choices", [])
    if not choices:
        return "No answer returned."
    message = choices[0].get("message", {})
    content: str = str(message.get("content", ""))
    # Perplexity includes citation URLs in the response metadata.
    citations = data.get("citations", [])
    if citations:
        citation_lines = [f"[{i + 1}] {url}" for i, url in enumerate(citations)]
        content += "\n\nSources:\n" + "\n".join(citation_lines)
    return content