# Copyright 2025 Acute Technology Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""A wrapper class of ZeroMQ's zmq.Socket class.

This is useful for cases where the application needs to send and receive
messages in a synchronous manner, such as when the application is running in
a single thread.
"""

from typing import Generic, TypeVar

import zmq
import zmq.asyncio

import aqpxlib.message as pxmsg
from aqpxlib.exceptions import BackendError, PxAPIError

RequestType = TypeVar("RequestType")
ResponseType = TypeVar("ResponseType")


# FIXME: This is a temporary solution to create a result object from a bytes object.
class Result(Generic[ResponseType]):
    """A wrapper class for the result of a request."""

    def __init__(
        self,
        topic: str | None = None,
        status: int | None = None,
        result: ResponseType | None = None,
    ) -> None:
        """Create wrapper object for the response data received."""
        self.topic = topic
        self.status = status  # FIXME: Check if this is still needed
        self.result = result


def _create_result(
    frames: list[bytes],
    response_cls: type[ResponseType],
) -> Result[ResponseType]:
    if len(frames) == 2:
        status = int.from_bytes(frames[0], byteorder="little")
        if status != pxmsg.PxResponseStatus.SUCCESS:
            raise PxAPIError(error_code=status)
        return Result(
            topic=None,
            status=status,
            result=response_cls.FromString(frames[1]),
        )
    return Result(
        topic=frames[0].decode("utf-8"),
        status=int.from_bytes(frames[1], byteorder="little"),
        result=response_cls.FromString(frames[2]),
    )


# FIXME: End of temporary solution


class SynchronousSocketWrapper(Generic[RequestType, ResponseType]):
    """A wrapper class for the synchronous ZeroMQsocket."""

    def __init__(self, socket: zmq.Socket | None = None) -> None:
        """Initialize a synchronous socket."""
        self.socket: zmq.Socket | None = socket

    def request(
        self,
        data: RequestType,
        timeout: int | None = None,
    ) -> Result[ResponseType]:
        """Send a request to the server and wait for a response.

        Args:
            data (RequestType): The data to send to the server.
            timeout (Optional[int]): The timeout in milliseconds.

        Returns:
            Optional[ResponseType]: The response from the server.

        """
        assert self.socket is not None
        assert hasattr(data, "__bytes__")

        self.socket.send(bytes(data))
        return self.recv(timeout=timeout)

    def recv(self, timeout: int | None = None) -> Result[ResponseType]:
        """Poll for a response from the server.

        Args:
            timeout (Optional[int]): The timeout in milliseconds.

        Returns:
            Result[T]: The result of the request.

        Raises:
            BackendError: If the request times out or there is a ZeroMQ error.

        """
        assert self.socket is not None

        try:
            status = self.socket.poll(timeout=timeout)
            if status != zmq.POLLIN:
                err_msg = "Timeout waiting for response from server"
                raise BackendError(err_msg)
            resp_frame = self.socket.recv_multipart()

            # response_cls is the second argument of the generic type of the SynchronousSocketWrapper class
            # which is the type of the response message
            return _create_result(
                frames=resp_frame,
                response_cls=self.__orig_class__.__args__[1],
            )
        except zmq.ZMQError:
            err_msg = "ZeroMQ error"
            raise BackendError(err_msg) from None

    def close(self) -> None:
        """Close the socket."""
        if self.socket is not None:
            self.socket.close(linger=0)
            self.socket = None


class AsynchronousSocketWrapper(Generic[RequestType, ResponseType]):
    """A wrapper class for the asynchronous ZeroMQ socket."""

    def __init__(self, socket: zmq.asyncio.Socket | None) -> None:
        """Initialize an asynchronous socket."""
        self.socket: zmq.asyncio.Socket | None = socket

    async def request(self, data: RequestType) -> ResponseType:
        """Send a request to the server and wait for a response.

        Args:
            data (RequestType): The data to send to the server.

        Returns:
            ResponseType: The response from the server.

        """
        assert self.socket is not None
        assert hasattr(data, "__bytes__")

        await self.socket.send(bytes(data))
        return await self.socket.recv()

    async def recv(self, timeout: int | None = None) -> ResponseType:
        """Poll a response from the server.

        Args:
            timeout (Optional[int]): The timeout in milliseconds.

        Returns:
            Result[T]: The result of the request.

        """
        assert self.socket is not None

        try:
            status = self.socket.poll(timeout=timeout)
            if status != zmq.POLLIN:
                err_msg = "Timeout waiting for response from server"
                raise BackendError(err_msg)

            resp_frame = self.socket.recv_multipart()
            return _create_result(
                frames=resp_frame,
                response_cls=self.__orig_class__.__args__[1],
            )

        except zmq.ZMQError:
            err_msg = "ZeroMQ error"
            raise BackendError(err_msg) from None

    def close(self) -> None:
        """Close the socket."""
        if self.socket is not None:
            self.socket.close(linger=0)
            self.socket = None
