"""Connect-JSON transport for the agentrouter Python SDK.

Implements a minimal Connect-RPC client over JSON using :mod:`httpx`. The
transport takes a request protobuf message, serializes it to JSON via
``json_format``, POSTs to ``<base>/<package>.<Service>/<Method>``, and
deserializes the response into the supplied response message type.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Callable, TypeVar

import httpx
from google.protobuf import json_format
from google.protobuf.message import Message

from .errors import AgentrouterError, RETRYABLE, UNAVAILABLE, from_http_response


T = TypeVar("T", bound=Message)


@dataclass
class RetryPolicy:
    """Retry policy controlling transient-failure recovery."""

    max_attempts: int = 3
    initial_backoff_s: float = 0.1
    max_backoff_s: float = 5.0
    total_budget_s: float = 90.0
    jitter_fraction: float = 0.2


@dataclass
class TransportConfig:
    """Resolved transport configuration shared across service handles."""

    base_url: str
    headers: dict[str, str] = field(default_factory=dict)
    timeout_s: float = 30.0
    retry: RetryPolicy = field(default_factory=RetryPolicy)
    httpx_client_factory: Callable[[], httpx.Client] | None = None


class Transport:
    """Connect-JSON over HTTP transport.

    One transport per service endpoint URL is fine; the SDK shares a single
    underlying :class:`httpx.Client` across all transports when possible.
    """

    def __init__(self, cfg: TransportConfig, client: httpx.Client):
        self._cfg = cfg
        self._client = client

    def unary(
        self,
        service_fqn: str,
        method: str,
        request: Message,
        response_cls: type[T],
    ) -> T:
        """Call ``<base>/<service_fqn>/<method>`` and decode into response_cls.

        ``service_fqn`` is the fully-qualified proto service name (e.g.
        ``tars.catalog.v1.CatalogService``).
        """
        path = f"/{service_fqn}/{method}"
        url = f"{self._cfg.base_url}{path}"
        # MessageToJson defaults: camelCase field names, omit defaults, sorted.
        # Use the no-kwargs form so the SDK works across protobuf 5.26+ where
        # individual flag names have been renamed (e.g. including_default_value_fields ->
        # always_print_fields_with_no_presence in newer point releases).
        body = json_format.MessageToJson(request)
        headers = {
            "Content-Type": "application/json",
            "Connect-Protocol-Version": "1",
        }
        headers.update(self._cfg.headers)

        attempts = max(1, self._cfg.retry.max_attempts)
        backoff = self._cfg.retry.initial_backoff_s
        deadline = time.monotonic() + self._cfg.retry.total_budget_s if self._cfg.retry.total_budget_s > 0 else float("inf")

        last_err: AgentrouterError | None = None
        for attempt in range(attempts):
            try:
                resp = self._client.post(url, content=body, headers=headers, timeout=self._cfg.timeout_s)
            except httpx.TimeoutException as e:
                last_err = AgentrouterError("deadline_exceeded", str(e))
            except httpx.TransportError as e:
                last_err = AgentrouterError(UNAVAILABLE, str(e))
            else:
                if resp.status_code == 200:
                    out = response_cls()
                    if resp.content:
                        json_format.Parse(resp.content, out, ignore_unknown_fields=True)
                    return out
                last_err = from_http_response(resp.status_code, resp.content)

            if last_err is not None and last_err.code in RETRYABLE and attempt < attempts - 1:
                wait = min(backoff, self._cfg.retry.max_backoff_s)
                if time.monotonic() + wait < deadline:
                    time.sleep(wait)
                    backoff *= 2
                    continue
            assert last_err is not None
            raise last_err
        assert last_err is not None
        raise last_err

    def close(self) -> None:
        # The SDK owns the httpx.Client and closes it from Client.close().
        pass
