from __future__ import annotations
from typing import Iterator, Optional, Union
from collections.abc import Iterator
from convomeld.state import Action, State, Trigger, TriggerPlaceholder
from uuid import uuid4


class ConvoThread:
    def __init__(self) -> None:
        self._states: dict[str, State] = {}
        self._uuid = uuid4().hex[-6:]
        self._first_state = None
        self._last_state = None
        self._state_count = 0

    def __repr__(self) -> str:
        return f'ConvoThread(states={repr(self._states)}, first_state_name={repr(self._first_state and self._first_state.name)})'
    
    def __str__(self) -> str:
        res = ''

        for state in self.iter_states():
            trigger_str =  f'--{state.triggers[0].value}-->' if len(state.triggers) == 1 else ''
            res += f'{str(state)} {trigger_str} '

        return res

    def iter_states(self, with_prev_trigger: bool = False) -> ConvoThreadIterator:
        self.validate()
        return ConvoThreadIterator(self._states, self._first_state and self._first_state.name, with_prev_trigger=with_prev_trigger)

    def copy(self):
        self.validate()
        copy = ConvoThread()
        
        for state, prev_trigger in self.iter_states(with_prev_trigger=True):
            if prev_trigger is None:
                tp = TriggerPlaceholder.empty()
            else:
                tp = TriggerPlaceholder.from_trigger(prev_trigger)
                
            copy = copy.append_state(state.actions, tp, **state.attrs)

        return copy
        
    # Utils section

    def num_states(self) -> int:
        self.validate()
        return len(self._states)
    
    def num_triggers(self) -> int:
        self.validate()
        return max(len(self._states) - 1, 0)
    
    def get_first_state(self) -> Optional[State]:
        self.validate()
        return self._first_state and self._first_state.copy()

    def get_last_state(self) -> Optional[State]:
        self.validate()
        return self._last_state and self._last_state.copy()
    
    def append_state(self, actions: Union[Action, list[Action]], tp: TriggerPlaceholder, **state_attrs) -> ConvoThread:
        self.validate()
        if isinstance(actions, Action):
            actions = [actions]

        self._state_count += 1
        new_state = State(f'path_{self._uuid}/state_{self._state_count}', actions, **state_attrs)
        self._states[new_state.name] = new_state

        if len(self._states) == 1:
            self._first_state = new_state
            self._last_state = new_state
            self.validate()
            return self
        
        prev_trigger = tp.create_trigger(new_state.name)
        self._last_state.triggers.append(prev_trigger)
        self._last_state = new_state
        self.validate()
        return self
    
    def prepend_state(self, actions: Union[Action, list[Action]], tp: TriggerPlaceholder, **state_attrs) -> ConvoThread:
        self.validate()
        if isinstance(actions, Action):
            actions = [actions]

        self._state_count += 1
        new_state = State(f'path_{self._uuid}/state_{self._state_count}', actions, **state_attrs)
        self._states[new_state.name] = new_state

        if len(self._states) == 1:
            self._first_state = new_state
            self._last_state = new_state
            self.validate()
            return self

        new_trigger = tp.create_trigger(self._first_state.name)
        new_state.triggers.append(new_trigger)
        self._first_state = new_state
        self.validate()
        return self
    
    def pop_first_state(self) -> Optional[State]:
        self.validate()

        if len(self._states) == 0:
            return None
        elif len(self._states) == 1:
            state = self._states.pop(self._first_state.name)
            self._first_state = None
            self._last_state = None
            return state.copy()
        else:
            state = self._states.pop(self._first_state.name)
            self._first_state = self._states[state.triggers[0].state_name]
            return state.copy()
    
    # Validation section

    def validate(self) -> None:
        self._validate_state_names()
        self._validate_linear()

    def _validate_state_names(self) -> None:
        for state_name, state in self._states.items():
            if state_name != state.name:
                raise RuntimeError('State names constraint of ConvoThread violated')
            
            for trigger in state.triggers:
                if trigger.state_name not in self._states:
                    raise RuntimeError('State names constraint of ConvoThread violated')
    
    def _validate_linear(self) -> None:
        if len(self._states) == 0:
            return

        current_state = self._first_state
        num_states = 0

        while current_state is not None and num_states <= len(self._states) + 1:
            if len(current_state.triggers) == 0:
                # Reached end
                current_state = None
            elif len(current_state.triggers) == 1:
                next_state_name = current_state.triggers[0].state_name
                current_state = self._states[next_state_name]
            else:
                raise RuntimeError('Linearity constraint of ConvoThread violated')

            num_states += 1

        if num_states != len(self._states):
            # print(self)
            raise RuntimeError('Linearity constraint of ConvoThread violated')


class ConvoThreadIterator(Iterator[Union[State, tuple[State, Trigger]]]):
        def __init__(self, states: dict[str, State], first_state_name: Optional[str], with_prev_trigger: bool) -> None:
            self._states = {state_name: state.copy() for state_name, state in states.items()}
            self._first_state = self._states.get(first_state_name, None)
            self._with_prev_trigger = with_prev_trigger

            self._current_state = self._first_state
            self._current_trigger = TriggerPlaceholder.none().create_trigger(self._current_state.name)

        def __next__(self) -> State:
            if self._current_state is None:
                raise StopIteration()
            
            current_state = self._current_state
            current_trigger = self._current_trigger

            if len(self._current_state.triggers) == 0:
                # Reached end
                next_state = None
                next_trigger = None
            else:
                next_trigger = self._current_state.triggers[0]
                next_state = self._states[next_trigger.state_name]

            # Prepare next iterator state
            self._current_state = next_state
            self._current_trigger = next_trigger

            # Collect result
            if self._with_prev_trigger:
                result = (current_state, current_trigger)
            else:
                result = current_state

            return result