#####################################################
# @module    state_machine.py                       #
# @brief     Basic Finite State Machine             #
# @author    Jamie Bird                             #
# @date      Oct 2020                               #
# @copyright Racelogic Ltd., 2020                   #
# @Version   1.0                                    #
#####################################################

class StateMachineError(Exception):
    pass

class StateMachine:
    def __init__(self, states=None):
        self.active_state = None
        self.start_state = None
        self.states = []
        if states:
            for state in states:
                self.add_state(state)

    def add_state(self, state, initial=False):
        if isinstance(state, type(self).State):
            self.states.append(state)
            state._set_machine(self)
            if initial:
                self.start_state = state
            setattr(self, state.name, state)
        else:
            raise StateMachineError("States must be State Objects")

    def current_state(self):
        return self.active_state.name

    def _set_state(self, state):
        if state == self.active_state:
            return False
        elif state in self.states:
            self.active_state = state
            return True
        else:
            self.add_state(state)
            return self._set_state(state)

    def add_transitions(self, t_dict):
        for key, val in t_dict.items():
            if hasattr(self, key):
                getattr(self, key).add_transitions(val)

    def run(self, start_state=None):
        if isinstance(start_state, type(self).State) and (start_state in self.states):
            start_state()
        elif self.start_state:
            self.start_state()
        else:
            raise StateMachineError("Can't start State Machine without a starting point")



    class State:
        def __init__(self, name, value=None, on_enter=None, on_exit=None, next_state=None, transitions=None):
            self.name = name
            self.value = value if value else name
            self.on_enter = []
            self.on_exit = []
            self.add_entry_action(on_enter)
            self.add_exit_action(on_exit)
            self.machine = None
            self.transitions = transitions if transitions else []
            self.next_state = None

        def __call__(self, *args, **kwargs):
            if self.machine:
                if self.machine.active_state:
                    if self.machine.active_state.can_transition(self):
                        self.machine.active_state.__exit__(None, None, None)
                    else:
                        raise StateMachineError("Unable to transition from {} to {}".format(\
                            self.machine.active_state.name, self.name))
                if self.machine._set_state(self):
                    self.__enter__(*args, **kwargs)
            return self.name

        def __eq__(self, other):
            return self.value == other.value

        def __ne__(self, other):
            return self.value != other.value

        def __repr__(self):
            return self

        def __str__(self):
            return self.name

        def check_action(self, action):
            if callable(action):
                return True
            elif action is None:
                return False
            else:
                raise StateMachineError("Action {} not callable".format(action))

        # def add_actions(self, actions, transition):
        #     # if isinstance(actions, (list, tuple, set)):
        #     #     for action in actions:
        #     #         if self.check_action(action):
        #     #             transition.append(action)
        #     if self.check_action(action):
        #         transition.append(action)

        def add_entry_action(self, action):
            if self.check_action(action):
                self.on_enter = action

        def add_exit_action(self, action):
            if self.check_action(action):
                self.on_exit = action

        def set_next_state(self, state):
            if isinstance(state, State):
                self.next_state = state
            else:
                raise StateMachineError("Next State must be a State Object")

        def _set_machine(self, machine):
            if isinstance(machine, StateMachine):
                self.machine = machine
            else:
                raise StateMachineError("Attribute 'machine' must be a StateMachine Object")

        def add_transitions(self, *args):
            # print("Args: {}, type={}".format(args, type(args)))
            self.transitions.append(args[0]) if (len(args) == 1 and isinstance(args[0], str))\
                                                else  self.transitions.extend(*args)
            # print(self.transitions)

        def can_transition(self, state):
            if not self.transitions:
                return True
            return (state.name in self.transitions)

        def __enter__(self, *args, **kwargs):
            if self.on_enter:
                self.on_enter(self, *args, **kwargs)


        def __exit__(self, e_type, value, traceback):
            # print("Exit Func Args: ", e_type, value, traceback)
            if self.on_exit:
                self.on_exit(self)


State = StateMachine.State

if __name__ == "__main__":

    def test_enter(state=None, func=None, t2=None):
        if state:
            print("Started {}".format(state))
        else:
            print("Started Something!")
        if func: func(t2) if t2 else func()

    def start_traffic(x=23):
        print("Traffic Started")

    def stop_traffic(time=10):
        print("Traffic Halted for {} seconds".format(time))

    def test_exit(state=None, func=None, t2=None):
        if state:
            print("Leaving {}".format(state))
        else:
            print("Leaving Something!")
        if func:
            func(t2)

    class TrafficLight(StateMachine):
        def __init__(self):
            super().__init__()
            self.add_state(StateMachine.State('Red', on_enter=test_enter, on_exit=test_exit))
            self.add_state(StateMachine.State('Yellow', on_enter=test_enter, on_exit=test_exit))
            self.add_state(StateMachine.State('Green', on_enter=test_enter, on_exit=test_exit))
            self.add_transitions({'Red': ('Yellow',), 'Green': 'Yellow'})

        def pedestrian_button(self):
            if self.current_state() == 'Green':
                self.Yellow()
                self.Red()
                print("Pedestrian can now go")
                self.Yellow()
                self.Green()
                print("Traffic Restarted")

    tl = TrafficLight()
    tl.run(tl.Red)
    tl.Yellow()
    tl.Green(start_traffic)

    tl.pedestrian_button()

    assert tl.Yellow() == 'Yellow'
    tl.Red(stop_traffic, 20)
    try:
        tl.Red()
    except StateMachineError:
        print("\n\nState Machine Working Correctly\n\n")

#end