import typing
from typing import Any  # noqa
from typing import MutableMapping  # noqa

from mitmproxy.coretypes import serializable
from mitmproxy.utils import typecheck


class StateObject(serializable.Serializable):
    """
    An object with serializable state.

    State attributes can either be serializable types(str, tuple, bool, ...)
    or StateObject instances themselves.
    """

    _stateobject_attributes = None  # type: MutableMapping[str, Any]
    """
    An attribute-name -> class-or-type dict containing all attributes that
    should be serialized. If the attribute is a class, it must implement the
    Serializable protocol.
    """

    def get_state(self):
        """
        Retrieve object state.
        """
        state = {}
        for attr, cls in self._stateobject_attributes.items():
            val = getattr(self, attr)
            state[attr] = get_state(cls, val)
        return state

    def set_state(self, state):
        """
        Load object state from data returned by a get_state call.
        """
        state = state.copy()
        for attr, cls in self._stateobject_attributes.items():
            val = state.pop(attr)
            if val is None:
                setattr(self, attr, val)
            else:
                curr = getattr(self, attr)
                if hasattr(curr, "set_state"):
                    curr.set_state(val)
                else:
                    setattr(self, attr, make_object(cls, val))
        if state:
            raise RuntimeWarning("Unexpected State in __setstate__: {}".format(state))


def _process(typeinfo: typecheck.Type, val: typing.Any, make: bool) -> typing.Any:
    if val is None:
        return None
    elif make and hasattr(typeinfo, "from_state"):
        return typeinfo.from_state(val)
    elif not make and hasattr(val, "get_state"):
        return val.get_state()

    typename = str(typeinfo)

    if typename.startswith("typing.List"):
        T = typecheck.sequence_type(typeinfo)
        return [_process(T, x, make) for x in val]
    elif typename.startswith("typing.Tuple"):
        Ts = typecheck.tuple_types(typeinfo)
        if len(Ts) != len(val):
            raise ValueError("Invalid data. Expected {}, got {}.".format(Ts, val))
        return tuple(
            _process(T, x, make) for T, x in zip(Ts, val)
        )
    elif typename.startswith("typing.Dict"):
        k_cls, v_cls = typecheck.mapping_types(typeinfo)
        return {
            _process(k_cls, k, make): _process(v_cls, v, make)
            for k, v in val.items()
        }
    elif typename.startswith("typing.Any"):
        # FIXME: Remove this when we remove flow.metadata
        assert isinstance(val, (int, str, bool, bytes))
        return val
    else:
        return typeinfo(val)


def make_object(typeinfo: typecheck.Type, val: typing.Any) -> typing.Any:
    """Create an object based on the state given in val."""
    return _process(typeinfo, val, True)


def get_state(typeinfo: typecheck.Type, val: typing.Any) -> typing.Any:
    """Get the state of the object given as val."""
    return _process(typeinfo, val, False)
