"""HTTP transport for Omni API calls."""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

import httpx

from omni.auth import SideroSigner
from omni.config.models import CacheConfig, OmniInstanceConfig, RetryConfig
from omni.http.cache import CacheController, is_mutation
from omni.http.methods import RpcMethod
from omni.http.retry import RetryPolicy
from omni.http.stream import parse_gateway_stream
from omni.state import StateStore


class OmniHTTPError(RuntimeError):
    """API request failed."""

    def __init__(self, message: str, *, status_code: int | None = None, body: Any = None) -> None:
        super().__init__(message)
        self.status_code = status_code
        self.body = body


class OmniTransport:
    """Async transport handling retries, caching, and request signing."""

    def __init__(
        self,
        *,
        instance_name: str,
        instance: OmniInstanceConfig,
        state: StateStore,
        signer: SideroSigner | None,
        timeout: float,
        verify_tls: bool | str,
        retry_config: RetryConfig,
        cache_config: CacheConfig,
    ) -> None:
        self.instance_name = instance_name
        self.instance = instance
        self.signer = signer
        self.retry = RetryPolicy(retry_config)
        self.cache = CacheController(cache_config, state)
        self._client = httpx.AsyncClient(
            base_url=instance.url.rstrip("/"),
            timeout=timeout,
            verify=verify_tls,
            headers={"Accept": "application/json"},
        )

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

    def _api_headers(
        self,
        method_path: str,
        metadata: dict[str, str] | None,
    ) -> dict[str, str]:
        headers: dict[str, str] = {}
        metadata = metadata or {}

        # Forward all metadata as grpc gateway headers.
        for key, value in metadata.items():
            headers[f"Grpc-Metadata-{key}"] = value

        if self.signer is not None:
            signed = self.signer.sign_api_request(method_path=method_path, metadata=metadata)
            headers.update(signed.headers)

        return headers

    def _image_headers(self, http_method: str, request_uri: str, body: bytes) -> dict[str, str]:
        if self.signer is None:
            return {}

        return self.signer.sign_image_request(
            http_method=http_method,
            request_uri=request_uri,
            body=body,
        ).headers

    async def call(
        self,
        method: RpcMethod,
        request: dict[str, Any],
        *,
        metadata: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        method_key = method.key
        cache_decision = self.cache.decision(method_key, stream=method.stream)
        cache_key = self.cache.cache_key(method_key, request, metadata, self.instance_name)

        if cache_decision.enabled:
            hit = self.cache.get(cache_key)
            if hit is not None:
                return hit

        attempt = 1
        while True:
            headers = self._api_headers(method.fq_method, metadata)

            try:
                response = await self._client.post(method.path, json=request, headers=headers)
            except Exception as exc:
                if (
                    method.idempotent
                    and self.retry.should_retry_exception(exc)
                    and attempt < self.retry.config.max_attempts
                ):
                    attempt += 1
                    await self.retry.wait(attempt)
                    continue
                raise OmniHTTPError(f"request failed for {method.key}: {exc}") from exc

            if response.status_code >= 400:
                payload: Any
                try:
                    payload = response.json()
                except Exception:
                    payload = response.text

                if (
                    method.idempotent
                    and self.retry.should_retry_status(response.status_code)
                    and attempt < self.retry.config.max_attempts
                ):
                    attempt += 1
                    await self.retry.wait(attempt)
                    continue

                raise OmniHTTPError(
                    f"request failed for {method.key}",
                    status_code=response.status_code,
                    body=payload,
                )

            payload = response.json()
            if not isinstance(payload, dict):
                raise OmniHTTPError(
                    f"invalid JSON response for {method.key}: expected object",
                    status_code=response.status_code,
                    body=payload,
                )

            if cache_decision.enabled:
                self.cache.set(cache_key, payload, cache_decision.ttl)

            if is_mutation(method_key):
                self.cache.invalidate_after_mutation(method_key, self.instance_name)

            return payload

    async def stream(
        self,
        method: RpcMethod,
        request: dict[str, Any],
        *,
        metadata: dict[str, str] | None = None,
    ) -> AsyncIterator[dict[str, Any]]:
        if not method.stream:
            raise ValueError(f"method {method.key} is not streaming")

        headers = self._api_headers(method.fq_method, metadata)
        response = await self._client.post(method.path, json=request, headers=headers)

        if response.status_code >= 400:
            try:
                payload: Any = response.json()
            except Exception:
                payload = response.text
            raise OmniHTTPError(
                f"stream request failed for {method.key}",
                status_code=response.status_code,
                body=payload,
            )

        async for item in parse_gateway_stream(response):
            yield item

    async def image_request(
        self,
        *,
        path: str,
        http_method: str = "GET",
        body: bytes = b"",
        extra_headers: dict[str, str] | None = None,
    ) -> httpx.Response:
        headers = self._image_headers(http_method, path, body)
        if extra_headers:
            headers.update(extra_headers)

        response = await self._client.request(http_method, path, content=body, headers=headers)
        return response

    async def raw_httpx_client(self) -> httpx.AsyncClient:
        return self._client
