"""Client for interacting with a :class:`NatsRPCService`."""

from __future__ import annotations

import logging
import uuid
from typing import Any

from nats.aio.client import Client as NatsClient
from pydantic import BaseModel, TypeAdapter

from nats_common.services.models import RpcRequestEnvelope, RpcResponseEnvelope
from nats_common.services.utils import _encode_payload

logger = logging.getLogger(__name__)


class NatsRPCClient[
    RpcReqT: BaseModel | None,
    RpcResT: BaseModel | None,
]:
    """Client-side companion to :class:`~nats_common.services.rpcs.NatsRPCService`.

    Sends a NATS request to the RPC subject and waits for the reply.

    Parameters
    ----------
    nats_client:
        An already-connected :class:`nats.aio.client.Client`.
    subject:
        Base subject matching the service (e.g. ``"myapp.echo"``).
    response_model:
        Optional Pydantic model for deserializing the ``data`` field of
        :class:`RpcResponseEnvelope`.
    timeout_seconds:
        Default timeout for every :meth:`call`.

    Usage::

        client = NatsRPCClient(
            nats_client=nats_client,
            subject="myapp.echo",
            response_model=EchoResponse,
        )
        resp = await client.call(data=EchoRequest(message="hello"))
        print(resp.data)  # EchoResponse(echoed="hello")
    """

    def __init__(
        self,
        *,
        nats_client: NatsClient,
        subject: str,
        response_model: type[RpcResT] | None = None,
        timeout_seconds: float = 30.0,
    ) -> None:
        self._nats_client = nats_client
        self.subject = subject.rstrip(".")
        self.request_subject = f"{self.subject}.rpc"
        self.timeout_seconds = timeout_seconds

        if response_model is None:
            self._resp_adapter: TypeAdapter[Any] = TypeAdapter(
                RpcResponseEnvelope[Any],
            )
        else:
            self._resp_adapter = TypeAdapter(
                RpcResponseEnvelope[response_model],  # type: ignore[valid-type]
            )

    async def call(
        self,
        *,
        data: RpcReqT | None = None,
        rpc_id: str | None = None,
        timeout_seconds: float | None = None,
    ) -> RpcResponseEnvelope[RpcResT]:
        """Send an RPC request and wait for the response.

        Parameters
        ----------
        data:
            Request payload.
        rpc_id:
            Unique call identifier.  Auto-generated (UUID hex) if omitted.
        timeout_seconds:
            Override the instance-level default for this single call.

        Returns
        -------
        RpcResponseEnvelope[RpcResT]:
            The deserialized response.  Check ``resp.error`` for server-side
            errors.

        Raises
        ------
        nats.errors.TimeoutError:
            If no reply is received within the timeout.
        """
        if rpc_id is None:
            rpc_id = uuid.uuid4().hex

        req = RpcRequestEnvelope[Any](rpc_id=rpc_id, data=data)
        timeout = timeout_seconds if timeout_seconds is not None else self.timeout_seconds

        msg = await self._nats_client.request(
            self.request_subject,
            _encode_payload(req),
            timeout=timeout,
        )
        return self._resp_adapter.validate_json(msg.data)
