import threading
from typing import Any, Protocol

import httpx2


class ApiRuntimeStats:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._total_received = 0

    def record_received(self) -> None:
        with self._lock:
            self._total_received += 1

    def snapshot(self) -> dict[str, int]:
        with self._lock:
            return {
                "total_received": self._total_received,
            }


class WorkerStatsUnavailable(Exception):
    pass


class WorkerStatsProvider(Protocol):
    async def get_stats(
        self,
    ) -> dict[str, Any]:
        ...

    async def close(self) -> None:
        ...


class HttpWorkerStatsClient:
    def __init__(
        self,
        url: str,
        timeout_seconds: float,
    ) -> None:
        self._url = url

        self._client = httpx2.AsyncClient(
            timeout=timeout_seconds,
        )

    async def get_stats(
        self,
    ) -> dict[str, Any]:
        try:
            response = await self._client.get(
                self._url
            )

            response.raise_for_status()

            return response.json()

        except httpx2.HTTPError as exc:
            raise WorkerStatsUnavailable(
                f"Unable to retrieve worker stats: {exc}"
            ) from exc

    async def close(self) -> None:
        await self._client.aclose()