"""
    transitions.extensions.factory
    ------------------------------

    Adds locking to machine methods as well as model functions that trigger events.
    Additionally, the user can inject her/his own context manager into the machine if required.
"""

from collections import defaultdict
from functools import partial
from threading import Lock
import inspect
import warnings
import logging

from transitions_aio.core import Machine, Event, listify

_LOGGER = logging.getLogger(__name__)
_LOGGER.addHandler(logging.NullHandler())


try:
    from contextlib import nested  # Python 2
    from thread import get_ident  # pragma: no cover
    # with nested statements now raise a DeprecationWarning. Should be replaced with ExitStack-like approaches.
    warnings.simplefilter('ignore', DeprecationWarning)  # pragma: no cover

except ImportError:
    from contextlib import ExitStack, contextmanager
    from threading import get_ident

    @contextmanager
    def nested(*contexts):
        """Reimplementation of nested in Python 3."""
        with ExitStack() as stack:
            for ctx in contexts:
                stack.enter_context(ctx)
            yield contexts


class PicklableLock:
    """A wrapper for threading.Lock which discards its state during pickling and
        is reinitialized unlocked when unpickled.
    """

    def __init__(self):
        self.lock = Lock()

    def __getstate__(self):
        return ''

    def __setstate__(self, value):
        return self.__init__()

    def __enter__(self):
        self.lock.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.lock.__exit__(exc_type, exc_val, exc_tb)


class IdentManager:
    """Manages the identity of threads to detect whether the current thread already has a lock."""

    def __init__(self):
        self.current = 0

    def __enter__(self):
        self.current = get_ident()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.current = 0


class LockedEvent(Event):
    """An event type which uses the parent's machine context map when triggered."""

    def trigger(self, model, *args, **kwargs):
        """Extends transitions.core.Event.trigger by using locks/machine contexts."""
        # pylint: disable=protected-access
        # noinspection PyProtectedMember
        # LockedMachine._locked should not be called somewhere else. That's why it should not be exposed
        # to Machine users.
        if self.machine._ident.current != get_ident():
            with nested(*self.machine.model_context_map[id(model)]):
                return super(LockedEvent, self).trigger(model, *args, **kwargs)
        else:
            return super(LockedEvent, self).trigger(model, *args, **kwargs)


class LockedMachine(Machine):
    """Machine class which manages contexts. In it's default version the machine uses a `threading.Lock`
        context to lock access to its methods and event triggers bound to model objects.
    Attributes:
        machine_context (dict): A dict of context managers to be entered whenever a machine method is
            called or an event is triggered. Contexts are managed for each model individually.
    """

    event_cls = LockedEvent

    def __init__(self, model=Machine.self_literal, states=None, initial='initial', transitions=None,
                 send_event=False, auto_transitions=True,
                 ordered_transitions=False, ignore_invalid_triggers=None,
                 before_state_change=None, after_state_change=None, name=None,
                 queued=False, prepare_event=None, finalize_event=None, model_attribute='state',
                 model_override=False, on_exception=None, on_final=None,
                 machine_context=None, **kwargs):

        self._ident = IdentManager()
        self.machine_context = listify(machine_context) or [PicklableLock()]
        self.machine_context.append(self._ident)
        self.model_context_map = defaultdict(list)

        super(LockedMachine, self).__init__(
            model=model, states=states, initial=initial, transitions=transitions,
            send_event=send_event, auto_transitions=auto_transitions,
            ordered_transitions=ordered_transitions, ignore_invalid_triggers=ignore_invalid_triggers,
            before_state_change=before_state_change, after_state_change=after_state_change, name=name,
            queued=queued, prepare_event=prepare_event, finalize_event=finalize_event,
            model_attribute=model_attribute, model_override=model_override, on_exception=on_exception,
            on_final=on_final, **kwargs
        )

    # When we attempt to pickle a locked machine, using IDs wont suffice to unpickle the contexts since
    # IDs have changed. We use a 'reference' store with objects as dictionary keys to resolve the newly created
    # references. This should induce no restrictions compared to transitions 0.8.8 but enable the usage of unhashable
    # objects in locked machine.
    def __getstate__(self):
        state = {k: v for k, v in self.__dict__.items()}
        del state['model_context_map']
        state['_model_context_map_store'] = {mod: self.model_context_map[id(mod)] for mod in self.models}
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.model_context_map = defaultdict(list)
        for model in self.models:
            self.model_context_map[id(model)] = self._model_context_map_store[model]
        del self._model_context_map_store

    def add_model(self, model, initial=None, model_context=None):
        """Extends `transitions.core.Machine.add_model` by `model_context` keyword.
        Args:
            model (list or object): A model (list) to be managed by the machine.
            initial (str, Enum or State): The initial state of the passed model[s].
            model_context (list or object): If passed, assign the context (list) to the machines
                model specific context map.
        """
        models = listify(model)
        model_context = listify(model_context) if model_context is not None else []
        super(LockedMachine, self).add_model(models, initial)

        for mod in models:
            mod = self if mod is self.self_literal else mod
            self.model_context_map[id(mod)].extend(self.machine_context)
            self.model_context_map[id(mod)].extend(model_context)

    def remove_model(self, model):
        """Extends `transitions.core.Machine.remove_model` by removing model specific context maps
            from the machine when the model itself is removed. """
        models = listify(model)

        for mod in models:
            del self.model_context_map[id(mod)]

        return super(LockedMachine, self).remove_model(models)

    def __getattribute__(self, item):
        get_attr = super(LockedMachine, self).__getattribute__
        tmp = get_attr(item)
        if not item.startswith('_') and inspect.ismethod(tmp):
            return partial(get_attr('_locked_method'), tmp)
        return tmp

    def __getattr__(self, item):
        try:
            return super(LockedMachine, self).__getattribute__(item)
        except AttributeError:
            return super(LockedMachine, self).__getattr__(item)

    # Determine if the returned method is a partial and make sure the returned partial has
    # not been created by Machine.__getattr__.
    # https://github.com/tyarkoni/transitions/issues/214
    def _add_model_to_state(self, state, model):
        super(LockedMachine, self)._add_model_to_state(state, model)  # pylint: disable=protected-access
        for prefix in self.state_cls.dynamic_methods:
            callback = "{0}_{1}".format(prefix, self._get_qualified_state_name(state))
            func = getattr(model, callback, None)
            if isinstance(func, partial) and func.func != state.add_callback:
                state.add_callback(prefix[3:], callback)

    # this needs to be overridden by the HSM variant to resolve names correctly
    def _get_qualified_state_name(self, state):
        return state.name

    def _locked_method(self, func, *args, **kwargs):
        if self._ident.current != get_ident():
            with nested(*self.machine_context):
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
