"""
Definition for an RPC channel protocol on top of a websocket - enabling bi-directional request/response interactions
"""
import asyncio
from inspect import _empty, getmembers, ismethod, signature
from typing import Any, Coroutine, Dict, List
from asyncio.coroutines import coroutine

from pydantic import ValidationError

from .utils import gen_uid
from .rpc_methods import EXPOSED_BUILT_IN_METHODS, NoResponse, RpcMethodsBase
from .schemas import RpcMessage, RpcRequest, RpcResponse

from .logger import get_logger
logger = get_logger("RPC_CHANNEL")


class DEAFULT_TIMEOUT:
    pass

class RemoteValueError(ValueError):
    pass

class RpcException(Exception):
    pass

class RpcChannelClosedException(Exception):
    """
    Raised when the channel is closed mid-operation
    """
    pass

class UnknownMethodException(RpcException):
    pass


class RpcPromise:
    """
    Simple Event and id wrapper/proxy
    Holds the state of a pending request
    """

    def __init__(self, request: RpcRequest):
        self._request = request
        self._id = request.call_id
        # event used to wait for the completion of the request (upon receiving its matching response)
        self._event = asyncio.Event()

    @property
    def request(self):
        return self._request

    @property
    def call_id(self):
        return self._id

    def set(self):
        """
        Signal compeltion of request with received response
        """
        self._event.set()

    def wait(self):
        """
        Wait on the internal event - triggered on response
        """
        return self._event.wait()


class RpcProxy:
    """
    Helper class
    provide a __call__ interface for an RPC method over a given channel
    """

    def __init__(self, channel, method_name) -> None:
        self.method_name = method_name
        self.channel = channel

    def __call__(self, **kwds: Any) -> Any:
        return self.channel.call(self.method_name, args=kwds)


class RpcCaller:
    """
    Helper class provide an object (aka other) with callable methods for each remote method on the otherside
    """

    def __init__(self, channel, methods=None) -> None:
        self._channel = channel
        self._method_names = [method[0] for method in getmembers(
            methods, lambda i: ismethod(i))] if methods is not None else None


    def __getattribute__(self, name: str):
        if (not name.startswith("_") or name in EXPOSED_BUILT_IN_METHODS) and (self._method_names is None or name in self._method_names):
            return RpcProxy(self._channel, name)
        else:
            return super().__getattribute__(name)


# Callback signatures
async def OnConnectCallback(channel):
    pass

async def OnDisconnectCallback(channel):
    pass

async def OnErrorCallback(channel, err:Exception):
    pass


class RpcChannel:
    """
    A wire agnostic json-rpc channel protocol for both server and client.
    Enable each side to send RPC-requests (calling exposed methods on other side) and receive rpc-responses with the return value

    provides a .other property for callign remote methods.
    e.g. answer = channel.other.add(a=1,b=1) will (For example) ask the other side to perform 1+1 and will return an RPC-response of 2
    """

    def __init__(self, methods: RpcMethodsBase, socket, channel_id=None, default_response_timeout=None, sync_channel_id=False, **kwargs):
        """

        Args:
            methods (RpcMethodsBase): RPC methods to expose to other side
            socket: socket object providing simple send/recv methods
            channel_id (str, optional): uuid for channel. Defaults to None in which case a random UUID is generated.
            default_response_timeout(int, optional) default timeout for RPC call responses. Defaults to None - i.e. no timeout
            sync_channel_id(bool, optional) should get the other side of the channel id, helps to identify connections, cost a bit networking time.
                Defaults to False - i.e. not getting the other side channel id
        """
        self.methods = methods._copy_()
        # allow methods to access channel (for recursive calls - e.g. call me as a response for me calling you)
        self.methods._set_channel_(self)
        # Pending requests - id-mapped to async-event
        self.requests: Dict[str, asyncio.Event] = {}
        # Received responses
        self.responses = {}
        self.socket = socket
        # timeout
        self.default_response_timeout = default_response_timeout
        # Unique channel id
        self.id = channel_id if channel_id is not None else gen_uid()
        # flag control should we retrieve the channel id of the other side
        self._sync_channel_id = sync_channel_id
        # The channel id of the other side (if sync_channel_id is True)
        self._other_channel_id = None
        # asyncio event to check if we got the channel id of the other side
        self._channel_id_synced = asyncio.Event()
        #
        # convineice caller
        # TODO - pass remote methods object to support validation before call
        self.other = RpcCaller(self)
        # core event callback regsiters
        self._connect_handlers:List[OnConnectCallback] = []
        self._disconnect_handlers:List[OnDisconnectCallback] = []
        self._error_handlers = []
        # internal event
        self._closed = asyncio.Event()

        # any other kwarg goes straight to channel context (Accessible to methods)
        self._context = kwargs or {}

    @property
    def context(self) -> Dict[str, Any]:
        return self._context

    async def get_other_channel_id(self) -> str:
        """
        Method to get the channel id of the other side of the channel
        The _channel_id_synced verify we have it
        Timeout exception can be raised if the value isn't available
        """
        await asyncio.wait_for(self._channel_id_synced.wait(), self.default_response_timeout)
        return self._other_channel_id

    def get_return_type(self, method):
        method_signature = signature(method)
        return method_signature.return_annotation if method_signature.return_annotation is not _empty else str

    async def send(self, data):
        """
        For internal use. wrap calls to underlying socket
        """
        await self.socket.send(data)

    async def receive(self):
        """
        For internal use. wrap calls to underlying socket
        """
        return await self.socket.recv()

    async def close(self):
        res = await self.socket.close()
        # signal closer
        self._closed.set()
        return res

    def isClosed(self):
        return self._closed.is_set()

    async def wait_until_closed(self):
        return await self._closed.wait()

    async def on_message(self, data):
        """
        Handle an incoming RPC message
        This is the main function servers/clients using the channel need to call (upon reading a message on the wire)
        """
        try:
            message = RpcMessage.parse_obj(data)
            if message.request is not None:
                await self.on_request(message.request)
            if message.response is not None:
                await self.on_response(message.response)
        except ValidationError as e:
            logger.error(f"Failed to parse message", {'message':data, 'error':e})
            await self.on_error(e)
        except Exception as e:
            await self.on_error(e)
            raise


    def register_connect_handler(self, coros:List[OnConnectCallback]=None):
        """
        Register a connection handler callback that will be called (As an async task)) with the channel
        Args:
            coros (List[Coroutine]): async callback
        """
        if coros is not None:
            self._connect_handlers.extend(coros)

    def register_disconnect_handler(self, coros:List[OnDisconnectCallback]=None):
        """
        Register a disconnect handler callback that will be called (As an async task)) with the channel id
        Args:
            coros (List[Coroutine]): async callback
        """
        if coros is not None:
            self._disconnect_handlers.extend(coros)

    def register_error_handler(self, coros:List[OnErrorCallback]=None):
        """
        Register an error handler callback that will be called (As an async task)) with the channel and triggered error.
        Args:
            coros (List[Coroutine]): async callback
        """
        if coros is not None:
            self._error_handlers.extend(coros)

    async def on_handler_event(self, handlers, *args, **kwargs):
        await asyncio.gather(*(callback(*args, **kwargs) for callback in handlers))

    async def on_connect(self):
        """
        Run get other channel id if sync_channel_id is True
        Run all callbacks from self._connect_handlers
        """
        if self._sync_channel_id:
            self._get_other_channel_id_task = asyncio.create_task(self._get_other_channel_id())
        await self.on_handler_event(self._connect_handlers, self)

    async def _get_other_channel_id(self):
        """
        Perform call to the other side of the channel to get its channel id
        Each side is generating the channel id by itself so there is no way to identify a connection without this sync
        """
        if self._other_channel_id is None:
            other_channel_id = await self.other._get_channel_id_()
            self._other_channel_id = other_channel_id.result if other_channel_id and other_channel_id.result else None
            if self._other_channel_id is None:
                raise RemoteValueError()
            # update asyncio event that we received remote channel id
            self._channel_id_synced.set()
            return self._other_channel_id
        else:
            return self._other_channel_id

    async def on_disconnect(self):
        # disconnect happend - mark the channel as closed
        self._closed.set()
        await self.on_handler_event(self._disconnect_handlers, self)

    async def on_error(self, error:Exception):
        await self.on_handler_event(self._error_handlers, self, error)

    async def on_request(self, message: RpcRequest):
        """
        Handle incoming RPC requests - calling relevant exposed method
        Note: methods prefixed with "_" are protected and ignored.

        Args:
            message (RpcRequest): the RPC request with the method to call
        """
        # TODO add exception support (catch exceptions and pass to other side as response with errors)
        logger.debug("Handling RPC request - %s", {'request':message, 'channel':self.id})
        method_name = message.method
        # Ignore "_" prefixed methods (except the built in "_ping_")
        if (isinstance(method_name, str) and (not method_name.startswith("_") or method_name in EXPOSED_BUILT_IN_METHODS)):
            method = getattr(self.methods, method_name)
            if callable(method):
                result = await method(**message.arguments)
                if result is not NoResponse:
                    # get indicated type
                    result_type = self.get_return_type(method)
                    # if no type given - try to convert to string
                    if result_type is str and type(result) is not str:
                        result = str(result)
                    response = RpcMessage(response=RpcResponse[result_type](
                        call_id=message.call_id, result=result, result_type=getattr(result_type, "__name__", getattr(result_type, "_name", "unknown-type"))))
                    await self.send(response)

    def get_saved_promise(self, call_id):
        return self.requests[call_id]

    def get_saved_response(self, call_id):
        return self.responses[call_id]

    def clear_saved_call(self, call_id):
        del self.requests[call_id]
        del self.responses[call_id]

    async def on_response(self, response: RpcResponse):
        """
        Handle an incoming response to a previous RPC call

        Args:
            response (RpcResponse): the received response
        """
        logger.debug("Handling RPC response - %s", {'response':response})
        if response.call_id is not None and response.call_id in self.requests:
            self.responses[response.call_id] = response
            promise = self.requests[response.call_id]
            promise.set()

    async def wait_for_response(self, promise, timeout=DEAFULT_TIMEOUT) -> RpcResponse:
        """
        Wait on a previously made call
        Args:
            promise (RpcPromise): the awaitable-wrapper returned from the RPC request call
            timeout (int, None, or DEAFULT_TIMEOUT): the timeout to wait on the response, defaults to DEAFULT_TIMEOUT.
                - DEAFULT_TIMEOUT - use the value passed as 'default_response_timeout' in channel init
                - None - no timeout
                - a number - seconds to wait before timing out
        Raises:
            asyncio.exceptions.TimeoutError - on timeout
            RpcChannelClosedException - if the channel fails before wait could be completed
        """
        if timeout is DEAFULT_TIMEOUT:
            timeout = self.default_response_timeout
        # wait for the promise or until the channel is terminated
        for coro in asyncio.as_completed([promise.wait(), self._closed.wait()], timeout=timeout):
            await coro
            break
        response = self.responses.get(promise.call_id, NoResponse)
        # if the channel was closed before we could finish
        if response is NoResponse:
            raise RpcChannelClosedException(f"Channel Closed before RPC response for {promise.call_id} could be received")
        self.clear_saved_call(promise.call_id)
        return response

    async def async_call(self, name, args={}, call_id=None) -> RpcPromise:
        """
        Call a method and return the event and the sent message (including the chosen call_id)
        use self.wait_for_response on the event and call_id to get the return value of the call
        Args:
            name (str): name of the method to call on the other side
            args (dict): keyword args to pass tot he other side
            call_id (string, optional): a UUID to use to track the call () - override only with true UUIDs
        """
        call_id = call_id or gen_uid()
        msg = RpcMessage(request=RpcRequest(
            method=name, arguments=args, call_id=call_id))
        logger.debug("Calling RPC method - %s", {'message':msg})
        await self.send(msg)
        promise = self.requests[msg.request.call_id] = RpcPromise(msg.request)
        return promise

    async def call(self, name, args={}, timeout=DEAFULT_TIMEOUT):
        """
        Call a method and wait for a response to be received
        """
        promise = await self.async_call(name, args)
        return await self.wait_for_response(promise, timeout=timeout)




