# -*- coding: utf-8 -*-

__author__ = r'wsb310@gmail.com'

import typing
import queue
import asyncio
import grpc
import msgpack


CHANNEL_USABLE_STATE = (grpc.ChannelConnectivity.READY, grpc.ChannelConnectivity.IDLE)


class GRPCClient:

    __slots__ = [
        r'_channel_queue', r'_credentials', r'_options', r'_compression', r'_interceptors',
        r'_request_serializer', r'_response_deserializer',
    ]

    def __init__(
        self, *,
        credentials=None, options=None, compression=None, interceptors=None,
        request_serializer=msgpack.dumps, response_deserializer=msgpack.loads
    ):

        self._channel_queue: queue.SimpleQueue[grpc.aio.Channel] = queue.SimpleQueue()

        self._credentials = credentials
        self._options = options
        self._compression = compression
        self._interceptors = interceptors

        self._request_serializer = request_serializer
        self._response_deserializer = response_deserializer

    async def _make_channel(self, target, timeout) -> grpc.aio.Channel:

        channel = None

        try:

            if self._credentials is None:
                channel = grpc.aio.insecure_channel(
                    target, self._options, self._compression, self._interceptors
                )
            else:
                channel = grpc.aio.secure_channel(
                    target, self._credentials, self._options, self._compression, self._interceptors
                )

            if timeout > 0:
                await asyncio.wait_for(channel.channel_ready(), timeout)
            else:
                await channel.channel_ready()

        except asyncio.TimeoutError as err:

            await channel.close()

            raise err

        setattr(channel, r'name', target)

        return channel

    async def open(self, targets, timeout=0):

        await self.set_channel(targets, timeout)

    async def close(self):

        while not self._channel_queue.empty():
            channel = self._channel_queue.get_nowait()
            await channel.close()

    def get_channel(self) -> grpc.aio.Channel:

        channel = None

        for _ in range(self._channel_queue.qsize()):

            _channel = self._channel_queue.get_nowait()
            self._channel_queue.put_nowait(_channel)

            if _channel.get_state() in CHANNEL_USABLE_STATE:
                channel = _channel
                break

        return channel

    async def set_channel(self, targets, timeout=0):

        channel_dict = {}

        while not self._channel_queue.empty():
            channel = self._channel_queue.get_nowait()
            channel_dict[getattr(channel, r'name')] = channel

        for target in targets:
            if target not in channel_dict:
                channel_dict[target] = await self._make_channel(target, timeout)

        for target, channel in channel_dict.items():
            if target in targets:
                self._channel_queue.put_nowait(channel)
            else:
                await channel.close()

    async def unary_unary(self, method: str, call_params: typing.Dict):

        channel = self.get_channel()

        return await channel.unary_unary(
            method,
            request_serializer=self._request_serializer,
            response_deserializer=self._response_deserializer,
        )(call_params)

    def unary_stream(
            self,
            method: str,
            call_params: typing.Dict
    ) -> typing.AsyncIterable:

        channel = self.get_channel()

        return channel.unary_stream(
            method,
            request_serializer=self._request_serializer,
            response_deserializer=self._response_deserializer,
        )(call_params)

    async def stream_unary(
            self,
            method: str,
            call_params: typing.Union[typing.Iterable[typing.Dict], typing.AsyncIterable[typing.Dict]]
    ):

        channel = self.get_channel()

        return await channel.stream_unary(
            method,
            request_serializer=self._request_serializer,
            response_deserializer=self._response_deserializer,
        )(call_params)

    def stream_stream(
            self,
            method: str,
            call_params: typing.Union[typing.Iterable[typing.Dict], typing.AsyncIterable[typing.Dict]]
    ) -> typing.AsyncIterable:

        channel = self.get_channel()

        return channel.stream_stream(
            method,
            request_serializer=self._request_serializer,
            response_deserializer=self._response_deserializer,
        )(call_params)
