"""Reusable RPC service primitives on top of NATS.

Provides :class:`NatsRPCService` for short request/reply calls.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

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

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

logger = logging.getLogger(__name__)

# ===============================================================================
# region RPC Service
# ===============================================================================


class NatsRPCService[
    RpcReqT: BaseModel | None,
    RpcResT: BaseModel | None,
]:
    """Lightweight request/reply RPC service.

    Uses a core-NATS queue subscription for load balancing across replicas.

    Parameters
    ----------
    nats_client:
        An already-connected :class:`nats.aio.client.Client`.
    subject:
        Base subject - the actual subscription is ``<subject>.rpc``.
    rpc_cls:
        Concrete :class:`NatsRPC` subclass that implements
        :meth:`~NatsRPC.handle`.

    Usage::

        class MyRpc(NatsRPC[MyRequest, MyResponse]):
            async def handle(self) -> MyResponse:
                return MyResponse(...)

        svc = NatsRPCService(
            nats_client=nats_client,
            subject="myapp.echo",
            rpc_cls=MyRpc,
            request_model=MyRequest,
            response_model=MyResponse,
        )
        await svc.start()
    """

    def __init__(
        self,
        *,
        nats_client: NatsClient,
        subject: str,
        rpc_cls: type[NatsRPC[RpcReqT, RpcResT]],
        request_model: type[RpcReqT] | None = None,
        response_model: type[RpcResT] | None = None,
        num_max_parallel_requests: int = 32,
        queue_group: str | None = None,
    ) -> None:
        self._nats_client = nats_client
        self.subject = subject.rstrip(".")
        self.rpc_cls = rpc_cls
        self.request_model = request_model
        self.response_model = response_model
        self.num_max_parallel_requests = num_max_parallel_requests

        self.request_subject = f"{self.subject}.rpc"
        self.queue_group = queue_group or _sanitize_name(f"{self.subject}_rpc")

        self._sub: Any | None = None
        self._stop = asyncio.Event()
        self._sema = asyncio.Semaphore(num_max_parallel_requests)

        if request_model is None:
            self._req_adapter: TypeAdapter[Any] = TypeAdapter(
                RpcRequestEnvelope[Any],
            )
        else:
            self._req_adapter = TypeAdapter(
                RpcRequestEnvelope[request_model],  # type: ignore[valid-type]
            )

    # -- Lifecycle ---------------------------------------------------------

    async def start(self) -> None:
        """Subscribe to the RPC subject and start serving requests."""
        self._stop.clear()
        self._sub = await self._nats_client.subscribe(
            self.request_subject,
            queue=self.queue_group,
            cb=self._on_request,
        )

    async def stop(self) -> None:
        """Stop serving RPC requests."""
        self._stop.set()
        if self._sub is not None:
            try:
                await self._sub.unsubscribe()
            except Exception:
                logger.exception("Failed to unsubscribe RPC subscription")

    # -- Client helper -----------------------------------------------------

    async def call(
        self,
        *,
        rpc_id: str,
        data: RpcReqT | None = None,
        timeout_seconds: float = 30.0,
    ) -> RpcResponseEnvelope[Any]:
        """Send an RPC request and wait for the response."""
        req = RpcRequestEnvelope[Any](rpc_id=rpc_id, data=data)
        msg = await self._nats_client.request(
            self.request_subject,
            _encode_payload(req),
            timeout=timeout_seconds,
        )
        return TypeAdapter(RpcResponseEnvelope[Any]).validate_json(msg.data)

    # -- Internal ----------------------------------------------------------

    async def _on_request(self, msg: CoreMsg) -> None:
        if self._stop.is_set():
            return

        await self._sema.acquire()
        try:
            if not msg.reply:
                return

            envelope = self._req_adapter.validate_json(msg.data)
            rpc = self.rpc_cls(service=self, request=envelope, msg=msg)

            try:
                result = await rpc.handle()
                resp = RpcResponseEnvelope[Any](
                    rpc_id=envelope.rpc_id,
                    data=result,
                )
            except Exception as exc:
                resp = RpcResponseEnvelope[Any](
                    rpc_id=envelope.rpc_id,
                    error=str(exc),
                )

            await self._nats_client.publish(msg.reply, _encode_payload(resp))

        except Exception:
            logger.exception("RPC request handling failed")

        finally:
            self._sema.release()


# ===============================================================================
# region NatsRPC
# ===============================================================================


class NatsRPC[
    RpcReqT: BaseModel | None,
    RpcResT: BaseModel | None,
]:
    """Base class for RPC handlers.

    Subclass and override :meth:`handle` to implement your RPC logic.
    """

    def __init__(
        self,
        *,
        service: NatsRPCService[RpcReqT, RpcResT],
        request: RpcRequestEnvelope[Any],
        msg: CoreMsg,
    ) -> None:
        self.service = service
        self.request = request
        self.msg = msg
        self.rpc_id: str = request.rpc_id

    async def handle(self) -> RpcResT | None:
        """Override in a subclass to implement the RPC logic."""
        raise NotImplementedError
