from __future__ import annotations

from typing import TypeVar

import httpx
from google.protobuf.json_format import MessageToJson, Parse
from google.protobuf.message import Message

from ._errors import ConnectErrorMapper

__all__ = ["_AsyncJsonTransport", "_JsonTransport"]


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


def _serialise(message: Message) -> str:
    return MessageToJson(
        message,
        indent=None,
        preserving_proto_field_name=False,
        use_integers_for_enums=False,
    )


def _parse(response_text: str, response_cls: type[_T]) -> _T:
    return Parse(response_text, response_cls(), ignore_unknown_fields=True)


def _merge_headers(defaults: dict[str, str], extra: dict[str, str] | None) -> dict[str, str]:
    if not extra:
        return dict(defaults)
    merged = dict(defaults)
    merged.update(extra)
    return merged


class _JsonTransport:
    """Sync Connect-RPC JSON transport over `httpx.Client`."""

    def __init__(
        self,
        client: httpx.Client,
        base_url: str,
        default_headers: dict[str, str],
    ) -> None:
        self._client = client
        self._base_url = base_url.rstrip("/")
        self._default_headers = dict(default_headers)

    def call(
        self,
        service_fqn: str,
        method: str,
        request: Message,
        response_cls: type[_T],
        extra_headers: dict[str, str] | None = None,
    ) -> _T:
        url = f"{self._base_url}/rpc/{service_fqn}/{method}"
        headers = _merge_headers(self._default_headers, extra_headers)
        body = _serialise(request)
        response = self._client.post(url, headers=headers, content=body)
        if response.status_code >= 400:
            ConnectErrorMapper.rethrow(response)
        return _parse(response.text, response_cls)


class _AsyncJsonTransport:
    """Async Connect-RPC JSON transport over `httpx.AsyncClient`."""

    def __init__(
        self,
        client: httpx.AsyncClient,
        base_url: str,
        default_headers: dict[str, str],
    ) -> None:
        self._client = client
        self._base_url = base_url.rstrip("/")
        self._default_headers = dict(default_headers)

    async def call(
        self,
        service_fqn: str,
        method: str,
        request: Message,
        response_cls: type[_T],
        extra_headers: dict[str, str] | None = None,
    ) -> _T:
        url = f"{self._base_url}/rpc/{service_fqn}/{method}"
        headers = _merge_headers(self._default_headers, extra_headers)
        body = _serialise(request)
        response = await self._client.post(url, headers=headers, content=body)
        if response.status_code >= 400:
            ConnectErrorMapper.rethrow(response)
        return _parse(response.text, response_cls)