"""Diffbot web search API."""

from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from .client import Diffbot, DiffbotAsync

WEB_SEARCH_BASE = "https://llm.diffbot.com/api/v1/web_search"


def web_search(
    client: "Diffbot",
    text: str,
    *,
    num_results: Optional[int] = None,
    max_tokens: Optional[int] = None,
) -> Dict[str, Any]:
    headers = {"Authorization": f"Bearer {client.token}"}
    params: Dict[str, Any] = {"text": text}
    if num_results is not None:
        params["size"] = num_results
    if max_tokens is not None:
        params["maxTokens"] = max_tokens
    response = client._http.get(client.web_search_url, headers=headers, params=params)
    client._raise_for_status(response)
    return response.json()


async def web_search_async(
    client: "DiffbotAsync",
    text: str,
    *,
    num_results: Optional[int] = None,
    max_tokens: Optional[int] = None,
) -> Dict[str, Any]:
    headers = {"Authorization": f"Bearer {client.token}"}
    params: Dict[str, Any] = {"text": text}
    if num_results is not None:
        params["size"] = num_results
    if max_tokens is not None:
        params["maxTokens"] = max_tokens
    response = await client._http.get(client.web_search_url, headers=headers, params=params)
    client._raise_for_status(response)
    return response.json()