import logging
import json
import time
from typing import Optional, Literal
import queue

import numpy as np
import yarp

from wrapyfi.connect.clients import Client, Clients
from wrapyfi.middlewares.yarp import YarpMiddleware
from wrapyfi.encoders import JsonEncoder, JsonDecodeHook


class YarpClient(Client):

    def __init__(
        self,
        name: str,
        in_topic: str,
        carrier: Literal["tcp", "udp", "mcast"] = "tcp",
        persistent: bool = True,
        yarp_kwargs: Optional[dict] = None,
        **kwargs,
    ):
        """
        Initialize the client.

        :param name: str: Name of the client
        :param in_topic: str: Name of the input topic preceded by '/' (e.g. '/topic')
        :param carrier: str: Carrier protocol (e.g. 'tcp'). Default is 'tcp'
        :param persistent: bool: Whether to keep the service connection alive across multiple service calls. Default is True
        :param yarp_kwargs: dict: Additional kwargs for  the Yarp middleware
        :param kwargs: dict: Additional kwargs for the client
        """
        super().__init__(name, in_topic, carrier=carrier, **kwargs)
        YarpMiddleware.activate(**yarp_kwargs or {})

        self.persistent = persistent

    def close(self):
        """
        Close the client.
        """
        if hasattr(self, "_port") and self._port:
            if self._port is not None:
                self._port.close()

    def __del__(self):
        self.close()


@Clients.register("NativeObject", "yarp")
class YarpNativeObjectClient(YarpClient):
    def __init__(
        self,
        name: str,
        in_topic: str,
        carrier: Literal["tcp", "udp", "mcast"] = "tcp",
        persistent: bool = True,
        serializer_kwargs: Optional[dict] = None,
        deserializer_kwargs: Optional[dict] = None,
        **kwargs,
    ):
        """
        The NativeObject listener using the YARP Bottle construct assuming the data is serialized as a JSON string.
        Deserializes the data (including plugins) using the decoder and parses it to a Python object.

        :param name: str: Name of the client
        :param carrier: str: Carrier protocol (e.g. 'tcp'). Default is 'tcp'
        :param in_topic: str: Name of the input topic preceded by '/' (e.g. '/topic')
        :param serializer_kwargs: dict: Additional kwargs for the serializer
        :param deserializer_kwargs: dict: Additional kwargs for the deserializer
        :param persistent: bool: Whether to keep the service connection alive across multiple service calls. Default is True
        :param serializer_kwargs: dict: Additional kwargs for the serializer
        :param deserializer_kwargs: dict: Additional kwargs for the deserializer
        :param kwargs: dict: Additional kwargs for the client
        """
        super().__init__(
            name, in_topic, carrier=carrier, persistent=persistent, **kwargs
        )
        self._port = None
        self._queue = queue.Queue(maxsize=1)

        self._plugin_encoder = JsonEncoder
        self._plugin_kwargs = kwargs
        self._serializer_kwargs = serializer_kwargs or {}
        self._plugin_decoder_hook = JsonDecodeHook(**kwargs).object_hook
        self._deserializer_kwargs = deserializer_kwargs or {}

    def establish(self):
        """
        Establish the client's connection to the YARP service.
        """
        while not yarp.Network.exists(self.in_topic):
            logging.info(f"[YARP] Waiting for input port: {self.in_topic}")
            time.sleep(0.2)
        self._port = yarp.RpcClient()
        rnd_id = str(np.random.randint(100000, size=1)[0])
        self._port.open(self.in_topic + ":in" + rnd_id)
        self._port.addOutput(self.in_topic, self.carrier)
        if self.persistent:
            self.established = True

    def request(self, *args, **kwargs):
        """
        Send a request to the YARP service.

        :param args: tuple: Positional arguments to send in the request
        :param kwargs: dict: Keyword arguments to send in the request
        :return: Any: The response from the YARP service
        """
        if not self.established:
            self.establish()
        try:
            self._request(*args, **kwargs)
        except Exception as e:
            logging.error("[YARP] Service call failed: %s" % e)
        return self._await_reply()

    def _request(self, *args, **kwargs):
        """
        Internal method to send a request to the YARP service.

        :param args: tuple: Positional arguments to send in the request
        :param kwargs: dict: Keyword arguments to send in the request
        """
        args_str = json.dumps(
            [args, kwargs],
            cls=self._plugin_encoder,
            **self._plugin_kwargs,
            serializer_kwrags=self._serializer_kwargs,
        )
        args_msg = yarp.Bottle()
        args_msg.clear()
        args_msg.addString(args_str)

        msg = yarp.Bottle()
        msg.clear()

        self._port.write(args_msg, msg)
        obj = json.loads(
            msg.get(0).asString(),
            object_hook=self._plugin_decoder_hook,
            **self._deserializer_kwargs,
        )
        self._queue.put(obj, block=False)

    def _await_reply(self):
        """
        Wait for and return the reply from the YARP service.

        :return: Any: The response from the YARP service
        """
        try:
            reply = self._queue.get(block=True)
            return reply
        except queue.Full:
            logging.warning(
                f"[YARP] Discarding data because queue is full. "
                f"This happened due to bad synchronization in {self.__name__}"
            )
            return None


@Clients.register("Image", "yarp")
class YarpImageClient(YarpNativeObjectClient):
    def __init__(
        self,
        name: str,
        in_topic: str,
        carrier: Literal["tcp", "udp", "mcast"] = "tcp",
        width: int = -1,
        height: int = -1,
        rgb: bool = True,
        fp: bool = False,
        persistent: bool = True,
        serializer_kwargs: Optional[dict] = None,
        **kwargs,
    ):
        """
        The Image client using the YARP Bottle construct parsed to a numpy array.

        :param name: str: Name of the client
        :param in_topic: str: Name of the input topic preceded by '/' (e.g. '/topic')
        :param carrier: str: Carrier protocol (e.g. 'tcp'). Default is 'tcp'
        :param width: int: The width of the image. Default is -1
        :param height: int: The height of the image. Default is -1
        :param rgb: bool: Whether the image is RGB. Default is True
        :param fp: bool: Whether to utilize floating-point precision. Default is False
        :param persistent: bool: Whether to keep the service connection alive across multiple service calls. Default is True
        :param serializer_kwargs: dict: Additional kwargs for the serializer
        """
        if "jpg" in kwargs:
            logging.warning(
                "[YARP] YARP currently does not support JPG encoding in REQ/REP. Using raw image."
            )
            kwargs.pop("jpg")
        super().__init__(
            name,
            in_topic,
            carrier=carrier,
            persistent=persistent,
            serializer_kwargs=serializer_kwargs,
            **kwargs,
        )
        self.width = width
        self.height = height
        self.rgb = rgb
        self.fp = fp

    def _request(self, *args, **kwargs):
        """
        Internal method to send a request to the YARP service.

        :param args: tuple: Positional arguments to send in the request
        :param kwargs: dict: Keyword arguments to send in the request
        """
        args_str = json.dumps(
            [args, kwargs],
            cls=self._plugin_encoder,
            **self._plugin_kwargs,
            serializer_kwrags=self._serializer_kwargs,
        )
        args_msg = yarp.Bottle()
        args_msg.clear()
        args_msg.addString(args_str)
        msg = yarp.Bottle()
        msg.clear()
        self._port.write(args_msg, msg)
        img = json.loads(
            msg.get(0).asString(),
            object_hook=self._plugin_decoder_hook,
            **self._deserializer_kwargs,
        )
        height, width, channels = img.shape
        if 0 < self.width != width or 0 < self.height != height:
            raise ValueError("Incorrect image shape for client")
        else:
            self._queue.put(img, block=False)


@Clients.register("AudioChunk", "yarp")
class YarpAudioChunkClient(YarpNativeObjectClient):
    def __init__(
        self,
        name: str,
        in_topic: str,
        carrier: Literal["tcp", "udp", "mcast"] = "tcp",
        channels: int = 1,
        rate: int = 44100,
        chunk: int = -1,
        persistent: bool = True,
        serializer_kwargs: Optional[dict] = None,
        **kwargs,
    ):
        """
        The AudioChunk client using the YARP Bottle construct parsed to a numpy array.

        :param name: str: Name of the client
        :param in_topic: str: Name of the input topic preceded by '/' (e.g. '/topic')
        :param carrier: str: Carrier protocol (e.g. 'tcp'). Default is 'tcp'
        :param channels: int: Number of audio channels. Default is 1
        :param rate: int: Sampling rate of the audio. Default is 44100
        :param chunk: int: The size of audio chunks. Default is -1
        :param persistent: bool: Whether to keep the service connection alive across multiple service calls. Default is True
        :param serializer_kwargs: dict: Additional kwargs for the serializer
        """
        super().__init__(
            name,
            in_topic,
            carrier=carrier,
            persistent=persistent,
            serializer_kwargs=serializer_kwargs,
            **kwargs,
        )
        self.channels = channels
        self.rate = rate
        self.chunk = chunk

    def _request(self, *args, **kwargs):
        """
        Internal method to send a request to the YARP service.

        :param args: tuple: Positional arguments to send in the request
        :param kwargs: dict: Keyword arguments to send in the request
        """
        args_str = json.dumps(
            [args, kwargs],
            cls=self._plugin_encoder,
            **self._plugin_kwargs,
            serializer_kwrags=self._serializer_kwargs,
        )
        args_msg = yarp.Bottle()
        args_msg.clear()
        args_msg.addString(args_str)
        msg = yarp.Bottle()
        msg.clear()
        self._port.write(args_msg, msg)
        chunk, channels, rate, aud = json.loads(
            msg.get(0).asString(),
            object_hook=self._plugin_decoder_hook,
            **self._deserializer_kwargs,
        )
        if 0 < self.rate != rate:
            raise ValueError("Incorrect audio rate for client")
        if (
            0 < self.chunk != chunk
            or self.channels != channels
            or aud.size != chunk * channels
        ):
            raise ValueError("Incorrect audio shape for client")
        else:
            self._queue.put((aud, rate), block=False)
