import asyncio
import time
from dataclasses import dataclass
from importlib import import_module
from itertools import count
from typing import Any

from thrift.protocol.TBinaryProtocol import TBinaryProtocol
from thrift.Thrift import TMessageType, TType
from thrift.transport import TTransport

_SEQUENCE_NOS = count(1)
_SIMPLE_TYPES = {TType.STRING, TType.BOOL, TType.I64, TType.I32}  # TODO: more


async def _result_wait(self, seconds):
    # Slightly weird -- a function with a `self` argument outside of a class.
    # This is not for calling directly; rather we will add it as an attribute
    # to the class representing the Thrift call result (which is generated
    # from thrift code, so we can't add the method in a "normal" way)
    if (sent_time := getattr(self, "_sent_time", None)) is None:
        print("sent_time not found")
        # This means we're in a weird racy universe.  We don't want to not
        # sleep at all, as that seems likely to make the races worse.  So
        # we'll do a small wait to try to stabilize.
        to_sleep = 0.1
    else:
        # The next line might look weird until you remember that `sent_time`
        # (like all variables in python) has function scope
        to_sleep = seconds - (time.time() - sent_time)
    if to_sleep > 0:
        print(time.time(), "sleeping for", to_sleep)
        await asyncio.sleep(to_sleep)
        print(time.time(), "done sleeping")
    else:
        # FIXME: log a warning
        pass


class _ThriftError:
    """A wrapper for thrift exceptions returned by Finagle"""

    def __init__(self, wrapped):
        self._wrapped = wrapped


# The official way to use the Client classes generated by Thrift is to
# instantiate them with an input and an output protocol and a sequence id, and
# to use them statefully.  Rather than do that, we'll create some proxy
# objects which have the properties needed to read or to write a message.  And
# then we will call the messages directly off of the class (not an instance),
# passing these proxies as the `self` argument.
@dataclass
class _ReadProxy:
    _iprot: Any


@dataclass
class _WriteProxy:
    _seqid: int
    _oprot: Any


class ThriftMessageFactory:
    """A factory class for thrift RPC messages.

    The `fn_name` is the name of the remote function we want to call.
    `client` should be a client class for the service we will be calling,
    generated by thrift.  `stats_name` allows giving a different name to these
    calls in the mite prometheus metrics.  It defaults to `fn_name`, but can
    be overridden in the case we are calling identically-named functions in
    different services.

    """

    @classmethod
    def get_methods(cls, module):
        """TODO"""
        if "Processor" not in module.__dict__:
            raise ValueError(f"{module.__name__} doesn't contain a Processor class")
        for method_name in module.__dict__["Processor"].__dict__.keys():
            if not method_name.startswith("process_"):
                continue
            method_name = method_name[8:]
            yield method_name

    def __init__(
        self,
        fn_name,
        client,
        stats_name=None,
    ):
        self._fn_name = fn_name
        self._client = client
        self._stats_name = stats_name or fn_name
        if getattr(client, fn_name, None) is None:
            raise Exception(
                f"wrong client passed when instantiating message factory for: {fn_name}"
            )
        # Monkeypatch the chained_wait method onto the reply class
        module = import_module(client.__module__)
        reply_wrapper = getattr(module, fn_name + "_result")
        reply_object = self._get_reply_object(reply_wrapper)
        if reply_object is not None:
            reply_object.chained_wait = _result_wait
        # FIXME: this doesn't cover cases where functions have multiple
        # arguments.  Most (all) the cybertron thrift services have a single
        # struct as the first argument, but technically we're not as generic
        # as we need to be...
        self._args_struct = getattr(module, fn_name + "_args")
        self._result_object = getattr(module, fn_name + "_result")

    @staticmethod
    def _get_reply_object(wrapper):
        spec = wrapper.thrift_spec
        if spec == ():
            return
        spec = spec[0]
        if spec is None or spec[1] == TType.LIST:
            # the thrift method returns void or a list FIXME: this means mite
            # needs to check whether the result is a list etc.
            return
        else:
            return spec[3][0]

    def get_request_bytes(self, *args, **kwargs):
        """Get the bytes representing a serialized request to the RPC function.

        `args` and `kwargs` are passed to the constructor for the request
        arguments.

        """
        out_msg = self._args_struct.thrift_spec[1][3][0](*args, **kwargs)
        seq_no = next(_SEQUENCE_NOS)
        trans = TTransport.TMemoryBuffer()
        proto = TBinaryProtocol(trans)
        proxy = _WriteProxy(seq_no, proto)
        getattr(self._client, "send_" + self._fn_name)(proxy, out_msg)
        return trans._buffer.getvalue()

    def get_request_object(self, msg):
        """TODO"""
        trans = TTransport.TMemoryBuffer(msg)
        proto = TBinaryProtocol(trans)
        (name, type, seqid) = proto.readMessageBegin()
        # assert type == TODO
        r = self._args_struct()
        r.read(proto)
        return getattr(r, "request", ())

    def get_reply_object(self, msg):
        """Deserialize a reply from the bytes `msg`.

        If the call resulted in an exception, it will be returned wrapped in a
        `_ThriftError`.

        """
        trans = TTransport.TMemoryBuffer(msg)
        proto = TBinaryProtocol(trans)
        proxy = _ReadProxy(proto)
        try:
            result = getattr(self._client, "recv_" + self._fn_name)(proxy)
        except Exception as e:
            print("finagle error", str(e))
            result = _ThriftError(e)
        return result

    def get_reply_bytes(self, seqid, *args, **kwargs):
        """Get a bytestring representing a reply to this factory's function."""
        result = self._result_object()
        result_struct_class = self._get_reply_object(self._result_object)
        if result_struct_class is not None:
            result_struct = result_struct_class(*args, **kwargs)
            result.success = result_struct
        trans = TTransport.TMemoryBuffer()
        proto = TBinaryProtocol(trans)
        proto.writeMessageBegin(self._fn_name, TMessageType.REPLY, seqid)
        result.write(proto)
        proto.writeMessageEnd()
        return trans._buffer.getvalue()

    def get_reply_args(self, overrides):
        """TODO"""
        obj = self._get_reply_object(self._result_object)
        if obj is None:
            return {}
        kwargs = self._get_args_for_spec(obj.thrift_spec, overrides)
        return kwargs

    def _get_simple_type(self, type):
        if type == TType.STRING:
            return "foo"
        elif type == TType.BOOL:
            return False
        elif type in {TType.I64, TType.I32}:  # TODO: more
            return 1

    def _get_args_for_spec(self, spec, overrides=None):
        kwargs = {}
        if overrides is None:
            overrides = {}
        overrides = overrides.get(self._fn_name, {})
        for t in spec:
            if not isinstance(t, tuple):
                # FIXME: other than None, what non-tuples do we see?
                continue
            name = t[2]
            inner_type = t[1]
            if name in overrides:
                kwargs[name] = overrides[name]
            elif inner_type in _SIMPLE_TYPES:
                kwargs[name] = self._get_simple_type(inner_type)
            elif inner_type == TType.STRUCT:
                kwargs[name] = t[3][0](**self._get_args_for_spec(t[3][0].thrift_spec))
            elif inner_type == TType.MAP:
                key_type = t[3][0]
                if key_type == TType.STRUCT:
                    obj = t[3][1][0]
                    key = obj(**self._get_args_for_spec(obj.thrift_spec))
                elif key_type in _SIMPLE_TYPES:
                    key = self._get_simple_type(key_type)
                else:
                    raise Exception("unk key")
                value_type = t[3][2]
                if value_type == TType.STRUCT:
                    obj = t[3][3][0]
                    value = obj(**self._get_args_for_spec(obj.thrift_spec))
                elif value_type in _SIMPLE_TYPES:
                    value = self._get_simple_type(value_type)
                else:
                    raise Exception("unk val")
                kwargs[name] = {key: value}
            elif inner_type == TType.LIST:
                member_type = t[3][0]
                if member_type == TType.STRUCT:
                    # FIXME: would be nice to provide a value, but getStatus
                    # causes recursion on donwstreamStatuses, and an empty
                    # list is still a list...
                    # obj = t[3][1][0]
                    # value = obj(**self._get_args_for_spec(obj.thrift_spec))
                    # kwargs[name] = [value]
                    kwargs[name] = ()  # use tuple so it's hashable
                elif member_type in _SIMPLE_TYPES:
                    kwargs[name] = (self._get_simple_type(member_type),)
                else:
                    raise Exception("unk member type")
            else:
                breakpoint()
                raise Exception("can't initialize", inner_type)
        return kwargs
