client.py

"""Querit API client — wraps the official querit Python SDK."""

from __future__ import annotations

from typing import Any

from querit import QueritClient as _SDKClient
from querit.models.request import SearchRequest, SearchFilters, SiteFilter, GeoFilter
from querit.errors import QueritError

from querit_cli.common import QueritAPIError


class QueritClient:
    """Thin wrapper around the official Querit Python SDK."""

    def __init__(self, api_key: str, timeout: int = 30) -> None:
        # SDK internally sets: Authorization: Bearer {api_key}
        # so pass the raw key only — do NOT prepend "Bearer " here
        self._sdk = _SDKClient(api_key=api_key, timeout=timeout)

    def search(
        self,
        query: str,
        *,
        count: int | None = None,
        sites_include: list[str] | None = None,
        sites_exclude: list[str] | None = None,
        time_range: str | None = None,
        countries: list[str] | None = None,
        languages: list[str] | None = None,
    ) -> dict[str, Any]:
        """Call the Querit search API via the official SDK."""
        search_filters = None
        if any([sites_include, sites_exclude, time_range, countries, languages]):
            search_filters = SearchFilters(
                sites=SiteFilter(
                    include=sites_include or [],
                    exclude=sites_exclude or [],
                ) if (sites_include or sites_exclude) else None,
                geo=GeoFilter(
                    countries=countries,
                ) if countries else None,
                time_range=time_range,
                languages=languages,
            )

        request = SearchRequest(
            query=query,
            **({} if count is None else {"count": count}),
            **({"filters": search_filters} if search_filters else {}),
        )

        try:
            response = self._sdk.search(request)
        except QueritError as e:
            raise QueritAPIError(str(e)) from e
        except Exception as e:
            raise QueritAPIError(f"Unexpected error: {e}") from e

        # Normalise SDK response object → dict so output.py stays unchanged.
        # SDK item has .raw (the original API dict) for fields not exposed as properties.
        results = []
        for item in (response.results or []):
            raw = item.raw if isinstance(item.raw, dict) else {}
            entry: dict = {
                "url":       item.url or "",
                "title":     item.title or "",
                "snippet":   item.snippet or "",
                "page_age":  item.page_age or "",
                "site_name": raw.get("site_name", ""),
            }
            site_icon = raw.get("site_icon")
            if site_icon:
                entry["site_icon"] = site_icon
            results.append(entry)

        return {
            "took":          getattr(response, "took", ""),
            "search_id":     getattr(response, "search_id", None),
            "query_context": {"query": query},
            "results":       {"result": results},
        }