import tables
, strutils
, options


type
 # entry_action: Option[Callback]
 # exit_action: Option[Callback]
 State_Actions = tuple[ entry_action: Option[Callback]
                      , exit_action: Option[Callback]
                      ]

 Callback = proc(): void
 StateEvent[S,E] = tuple[state: S, event: E]
 Transition[S] = tuple[nextState: S, action: Option[Callback]]

 State_Machine*[S,E] =  ref object of RootObj
  initial_state*: S
  current_state*: Option[S]
  state_actions*: array[S, StateActions]
  transitions*: TableRef[StateEvent[S,E], Transition[S]]
  free_transitions*: TableRef[S, Transition[S]]
  default_transition*: Option[Transition[S]]

  #TransitionNotFoundException = object of Exception

proc current_state_is[S,E]( m: State_Machine[S,E]
                          , nextState: S
                          ) =
  if m.current_state.isSome:
   if m.state_actions[m.current_state.get].exit_action.isSome:
     get(m.state_actions[m.current_state.get].exit_action)()
  
  m.current_state = some(nextState)
  if m.state_actions[m.current_state.get].entry_action.isSome:
   get(m.state_actions[m.current_state.get].entry_action)()

proc reset*[S,E]( m: State_Machine[S,E]
                ) =
 m.setCurrentState m.initial_state

proc initial_state_is*[S,E]( m: State_Machine[S,E]
                          , state: S
                          ) =
 m.initial_state = state

proc add_state_actions*[S,E]( m: State_Machine[S,E]
                            , state: S
                            , entry_action: Callback = nil
                            , exit_action: Callback = nil
                            ) =
 let 
  entry = if entry_action == nil: none(Callback) 
          else: some(entry_action)
  
  exit = if exit_action == nil: none(Callback) 
         else: some(exit_action)

 m.state_actions[state] = (entry, exit)

proc a_state_machine*[S,E]( initial_state: S
                          ): State_Machine[S,E] =
 result = State_Machine[S,E]()
 result.transitions = newTable[ StateEvent[S,E]
                              , Transition[S]
                              ]()
 
 result.free_transitions = newTable[ S
                                   , Transition[S]
                                   ]()
 
 result.initial_state_is initial_state

proc add_free_transition*[S,E]( m: State_Machine[S,E]
                              , state: S, nextState: S
                              ) =
 m.free_transitions[state] = (nextState, none(Callback))

proc addTransitionAny*[S,E](m: State_Machine[S,E], state, nextState: S, action: Callback) =
 m.free_transitions[state] = (nextState, some(action))

proc add_transition*[S,E](m: State_Machine[S,E], state: S, event: E, nextState: S) =
 m.transitions[(state, event)] = (nextState, none(Callback))

proc add_transition*[S,E](m: State_Machine[S,E], state: S, event: E, nextState: S, action: Callback) =
 m.transitions[(state, event)] = (nextState, some(action))

proc setDefaultTransition*[S,E](m: State_Machine[S,E], state: S) =
 m.default_transition = some((state, none(Callback)))

proc setDefaultTransition*[S,E](m: State_Machine[S,E], state: S, action: Callback) =
 m.default_transition = some((state, some(action)))

proc the_transition*[S,E]( m: State_Machine[S,E]
                         , event: E
                         , state: S
                         ): Transition[S] =
 let map = (state, event)
 if m.transitions.hasKey(map): result = m.transitions[map]
 elif m.free_transitions.hasKey(state): result = m.free_transitions[state]
 elif m.default_transition.isSome: result = m.default_transition.get
 else: 
  echo map 
  quit "TransitionNotFoundException"
                         

proc process*[S,E]( m: State_Machine[S,E]
                  , event: E
                  ) =
 let transition = m.the_transition(event, m.current_state.get)
 if transition[1].isSome: get(transition[1])()
 m.current_state_is transition[0]
 #echo event, " ", m.current_state.get

when isMainModule:

  type
    StateName = enum
      SOLID
      LIQUID
      GAS
      PLASMA

    Event = enum
      MELT
      EVAPORATE
      SUBLIMATE
      IONIZE

  var m = newMachine[StateName, Event](LIQUID)

  proc cb() =
    echo "i'm evaporating"

  var condition: bool

  proc enterLiquid() =
    echo "entering liquid state"

  proc exitLiquid() =
    echo "exiting liquid state"

  proc enterGas() =
    if condition:
      echo "entering gas state and ionizing immediately"
      m.process(IONIZE)
    else:
      echo "entering gas state"

  proc exitGas() =
    echo "exiting gas state"

  proc enterPlasma() =
    echo "entering plasma state"


  m.add_transition(SOLID, MELT, LIQUID)
  m.add_transition(LIQUID, EVAPORATE, GAS, cb)
  m.add_transition(SOLID, SUBLIMATE, GAS)
  m.add_transition(GAS, IONIZE, PLASMA)
  m.add_transition(SOLID, MELT, LIQUID)

  m.addStateActions(LIQUID, entry_action=enterLiquid, exit_action=exitLiquid)
  m.addStateActions(PLASMA, enterPlasma)
  m.addStateActions(GAS, enterGas, exitGas)

  # to "start" the fsm
  # this is necessary
  m.reset

  assert m.getCurrentState() == LIQUID
  condition = false
  m.process(EVAPORATE)
  assert m.getCurrentState() == GAS
  m.process(IONIZE)
  assert m.getCurrentState() == PLASMA
  echo "\nreseting\n"

  m.reset

  assert m.getCurrentState() == LIQUID
  condition = true
  m.process(EVAPORATE)
  assert m.getCurrentState() == PLASMA, $m.getCurrentState()