from typing import Any, Callable, Dict, Generic, Optional, Set, TypeVar

from .connection import Connection, ConnectionStatus
from .envelope import Envelope, EnvelopeNotify, EnvelopeRequest
from .util import make_id
from .watchable import Watchable
from . import TransportBase

T = TypeVar('T')
Thunk = Callable[[], None]


class Connection(Generic[T]):
    def __init__(self, transport: TransportBase, device_id: str, description: str, methods: T, send_envelope: Callable[[Connection[T], Envelope], None]) -> None:
        self.status = Watchable(ConnectionStatus.CONNECTING)
        self._close_cbs: Set[Thunk] = set()
        self.description = description
        self._transport = transport
        self._device_id = device_id
        self._other_device_id: Optional[str] = None
        self._methods = methods
        self._send_envelope = send_envelope
        self._deferred_requests: Dict[str, Deferred[Any]] = {}
        self._last_seen = 0

    @property
    def is_closed(self) -> bool:
        return self.status == ConnectionStatus.CLOSED

    def on_close(self, func: Thunk) -> Thunk:
        if self.is_closed:
            raise RpcErrorUseAfterClose('the connection is closed')

        self._close_cbs.add(func)

        def del_cb():
            self._close_cbs.remove(func)

        return del_cb

    def close(self) -> None:
        if self.is_closed:
            return

        self.status = ConnectionStatus.CLOSED

        for func in self._close_cbs:
            func()

        self._close_cbs.clear()

    def handle_incoming_envelope(env: Envelope[T]) -> None:
        # TODO: maybe this function should be in a lock to ensure it only runs one at a time
        if self.is_closed:
            raise RpcErrorUseAfterClose('the connection is closed')

        # TODO: throw error if status is ERROR?
        if env.kind == 'NOTIFY':
            if not hasattr(self._methods, env.method):
                logger.warn(f'> error in NOTIFY handler: no notify method called {env.method}')
            else:
                try:
                    self._methods[env.method](*env.args)
                except BaseException as exc:
                    logger.warn(f'> error when running NOTIFY method: {env} {exc}')

        elif env.kind == 'REQUEST':
            try:
                if not hasattr(self._methods, env.method):
                    raise RpcErrorUnknownMethod(f'unknown method in REQUEST: {env.method}')

                data = self._methods[env.method](*env.args)
                response_env_data = EnvelopeResponseWithData(kind='RESPONSE', from_device_id=self._device_id, envelope_id=env.envelope_id, data=data)
                self._send_envelope(self, response_env_data)
            except BaseException as exc:
                response_env_error = EnvelopeResponseWithError(kind='RESPONSE', from_device_id=self.device_id, envelope_id=env.envelope_id, error=str(exc))

                self._send_envelope(self, response_error)
        elif env.kind == 'RESPONSE':
            deferred = self._deferred_requests.get(env.envelope_id)

            if not deferred:
                return

            if env.data:
                deferred.resolve(env.data)
            elif env.error:
                deferred.reject(RpcErrorFromMethod(env.error))
            else:
                logger.warn('> RESPONSE has neither data nor error. This should never happen')

                deferred.reject(RpcError('> RESPONSE has neither data nor error??'))

            self._deferred_requests.remove(env.envelope_id)

    def notify(self, method: MethodKey, *args: Parameters[BagType[MethodKey]]) -> None:
        if self.is_closed:
            raise RpcErrorUseAfterClose('the connection is closed')

        env = EnvelopeNotify(
            kind='NOTIFY',
            from_device_id=self._device_id,
            envelope_id=f'env:{make_id()}',
            method=method,
            args=args
        )
        self._send_envelope(self, env)

    def request(self, method: str, *args: Parameters[BagType[MethodKey]]) -> ReturnType[BagType[MethodKey]]:
        if self.is_closed:
            raise RpcErrorUseAfterClose('the connection is closed')

        env = EnvelopeRequest(kind='REQUEST', from_device_id=self._device_id, envelope_id=f'env:{make_id}', method=method, args=args)
        deferred = make_deferred()
        self._deferred_requests[env.envelope_id] = deferred
        self._send_envelope(self, env)

        return deferred.promise