diff --git a/README.md b/README.md
index 6ac1aa0..7b6af5d 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ implementation("io.nexure:fsm:2.0.0")
 
  ## Usage
  Here is a fictional example of what a state machine could look like that models the process of a
- payment. In summary it has
+ payment. In summary, it has
  - _Initial_ state `CREATED`
  - _Intermediary_ states `PENDING` and `AUTHORIZED`
  - _Terminal_ states `SETTLED` and `REFUSED`
@@ -33,4 +33,4 @@ implementation("io.nexure:fsm:2.0.0")
 
 See [ExampleStateMachineTest.kt](lib/src/test/kotlin/io/nexure/fsm/ExampleStateMachine.kt) for an
 example of how a state machine with the above states and transitions is built, and how it can
-be invoked to execute certain actions on a given state transition.
+be invoked to execute certain actions on state transition.
diff --git a/lib/build.gradle.kts b/lib/build.gradle.kts
index 90b7588..bd2b47b 100644
--- a/lib/build.gradle.kts
+++ b/lib/build.gradle.kts
@@ -29,6 +29,7 @@ dependencies {
 
     // Use the Kotlin test library.
     testImplementation("org.jetbrains.kotlin:kotlin-test")
+    testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.6.4")
 
     // Use the Kotlin JUnit integration.
     testImplementation("org.jetbrains.kotlin:kotlin-test-junit")
diff --git a/lib/src/main/kotlin/io/nexure/fsm/Action.kt b/lib/src/main/kotlin/io/nexure/fsm/Action.kt
deleted file mode 100644
index 8dc9717..0000000
--- a/lib/src/main/kotlin/io/nexure/fsm/Action.kt
+++ /dev/null
@@ -1,5 +0,0 @@
-package io.nexure.fsm
-
-interface Action<N : Any> {
-    fun action(signal: N)
-}
diff --git a/lib/src/main/kotlin/io/nexure/fsm/Edge.kt b/lib/src/main/kotlin/io/nexure/fsm/Edge.kt
index ebdbe15..204d69f 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/Edge.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/Edge.kt
@@ -3,11 +3,10 @@ package io.nexure.fsm
 /**
  * A connection for a transition state from one state to another, with an event.
  */
-internal class Edge<S : Any, E : Any, N : Any>(
+internal class Edge<S : Any, E : Any>(
     val source: S,
     val target: S,
     val event: E,
-    val action: (N) -> Unit
 ) {
     operator fun component1(): S = source
     operator fun component2(): S = target
diff --git a/lib/src/main/kotlin/io/nexure/fsm/StateMachine.kt b/lib/src/main/kotlin/io/nexure/fsm/StateMachine.kt
index 7481a16..0165a7d 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/StateMachine.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/StateMachine.kt
@@ -3,9 +3,8 @@ package io.nexure.fsm
 /**
  * S = State
  * E = Event triggering a transition between two states
- * N = Signal, data associated with a state change
  */
-interface StateMachine<S : Any, E : Any, N : Any> {
+interface StateMachine<S : Any, E : Any> {
     /**
      * Return all the possible states of the state machine
      */
@@ -28,16 +27,21 @@ interface StateMachine<S : Any, E : Any, N : Any> {
 
     /**
      * Execute a transition from [state] to another state depending on [event].
-     * If an action is associated with the state transition, it will then be executed,
-     * with [signal] as input. Returns a [Transition] indicating if the transition was permitted and
-     * successful or not.
+     * Returns a [Transition] indicating if the transition was permitted and
+     * successful or not by the state machine.
      *
      * It is recommended that the return value is checked for the desired outcome, if it is critical
      * that an event for example is accepted and not rejected.
      */
-    fun onEvent(state: S, event: E, signal: N): Transition<S>
+    fun onEvent(state: S, event: E): Transition<S>
+
+    /**
+     * Return a list of events that are accepted by the state machine in the given state. The returned list will be an
+     * empty list if the state is a terminal state.
+     */
+    fun acceptedEvents(state: S): Set<E>
 
     companion object {
-        fun <S : Any, E : Any, N : Any> builder(): StateMachineBuilder<S, E, N> = StateMachineBuilder()
+        fun <S : Any, E : Any> builder(): StateMachineBuilder.Uninitialized<S, E> = StateMachineBuilder()
     }
 }
diff --git a/lib/src/main/kotlin/io/nexure/fsm/StateMachineBuilder.kt b/lib/src/main/kotlin/io/nexure/fsm/StateMachineBuilder.kt
index eb723f9..2987f00 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/StateMachineBuilder.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/StateMachineBuilder.kt
@@ -1,108 +1,66 @@
 package io.nexure.fsm
 
-@Suppress("UNUSED_PARAMETER")
-private fun <N : Any> noOp(signal: N) {}
-
 /**
  * - [S] - the type of state that the state machine handles
  * - [E] - the type of events that the can trigger state changes
- * - [N] - the type of the input used in actions which are executed on state transitions
  */
-class StateMachineBuilder<S : Any, E : Any, N : Any> private constructor(
-    private var initialState: S? = null,
-    private val transitions: List<Edge<S, E, N>> = emptyList(),
-    private val interceptors: List<(S, S, E, N) -> (N)> = emptyList(),
-    private val postInterceptors: List<(S, S, E, N) -> Unit> = emptyList()
-) {
-    constructor() : this(null, emptyList(), emptyList(), emptyList())
-
-    /**
-     * Set the initial state for this state machine. There must be exactly one initial state,
-     * no more or less. Failing to set an initial state for a state machine will cause an
-     * [InvalidStateMachineException] to be thrown when [build()] is invoked.
-     *
-     * Calling this method more than once, with a different initial state will also cause an
-     * [InvalidStateMachineException] to be thrown, but immediately upon the second call to this
-     * method rather when the state machine is built.
-     */
-    @Throws(InvalidStateMachineException::class)
-    fun initial(state: S): StateMachineBuilder<S, E, N> {
-        return if (initialState == null) {
-            StateMachineBuilder(state, transitions, interceptors, postInterceptors)
-        } else if (state === initialState) {
-            StateMachineBuilder(initialState, transitions, interceptors, postInterceptors)
-        } else {
-            throw InvalidStateMachineException("There can only be one initial state")
-        }
+sealed class StateMachineBuilder<S : Any, E : Any> {
+    class Uninitialized<S : Any, E : Any> internal constructor(
+        private val transitions: List<Edge<S, E>> = emptyList(),
+    ) : StateMachineBuilder<S, E>() {
+        /**
+         * Set the initial state for this state machine. There can only one initial state,
+         * no more or less.
+         */
+        fun initial(state: S): Initialized<S, E> = Initialized(state, transitions)
     }
 
-    /**
-     * Create a state transition from [source] state to [target] state that will be triggered by
-     * [event], and execute an optional [action] when doing the state transition. There can be
-     * multiple events that connect [source] and [target], but there must never be any ambiguous
-     * transitions.
-     *
-     * For example, having both of the following transitions, would NOT be permitted
-     * - `(S1, E1) -> S2`
-     * - `(S1, E1) -> S3`
-     *
-     * since it would not be clear if the new state should be `S2` or `S3` when event `E1` is
-     * received.
-     */
-    fun connect(
-        source: S,
-        target: S,
-        event: E,
-        action: (signal: N) -> Unit = ::noOp
-    ): StateMachineBuilder<S, E, N> = connect(Edge(source, target, event, action))
-
-    fun connect(source: S, target: S, event: E, action: Action<N>): StateMachineBuilder<S, E, N> =
-        connect(source, target, event, action::action)
+    class Initialized<S : Any, E : Any> internal constructor(
+        private val initialState: S,
+        private val transitions: List<Edge<S, E>>,
+    ) : StateMachineBuilder<S, E>() {
+        /**
+         * Create a state transition from [source] state to [target] state that will be triggered by
+         * [event]. There can be multiple events that connect [source] and [target],
+         * but there must never be any ambiguous transitions.
+         *
+         * For example, having both of the following transitions, would NOT be permitted
+         * - `(S1, E1) -> S2`
+         * - `(S1, E1) -> S3`
+         *
+         * since it would not be clear if the new state should be `S2` or `S3` when event `E1` is
+         * received.
+         */
+        fun connect(
+            source: S,
+            target: S,
+            event: E,
+        ): Initialized<S, E> = connect(Edge(source, target, event))
 
-    private fun connect(edge: Edge<S, E, N>): StateMachineBuilder<S, E, N> =
-        StateMachineBuilder(initialState, transitions.plus(edge), interceptors, postInterceptors)
+        private fun connect(edge: Edge<S, E>): Initialized<S, E> =
+            Initialized(initialState, transitions.plus(edge))
 
-    /**
-     * Add an interceptor that is run _before_ any state machine action or state transition is done.
-     * The interceptor will only be run if the current state permits the event in question. No
-     * execution of the interceptor will be done if the state is rejected.
-     */
-    fun intercept(
-        interception: (source: S, target: S, event: E, signal: N) -> N
-    ): StateMachineBuilder<S, E, N> =
-        StateMachineBuilder(initialState, transitions, interceptors.plus(interception), postInterceptors)
+        /**
+         * @throws InvalidStateMachineException if the configured state machine is not valid. The main
+         * reasons for a state machine not being valid are:
+         * - No initial state
+         * - More than one initial state
+         * - The state machine is not connected (some states are not possible to reach from the initial
+         * state)
+         * - The same source state and event is defined twice
+         */
+        @Throws(InvalidStateMachineException::class)
+        fun build(): StateMachine<S, E> {
+            StateMachineValidator.validate(initialState, transitions)
 
-    /**
-     * Add an interceptor that is run _after_ a successful processing of an event by the state
-     * machine. This interceptor will not be run if the event was rejected by the state machine, or
-     * if there was an exception thrown while executing the state machine action (if any).
-     */
-    fun postIntercept(
-        interception: (source: S, target: S, event: E, signal: N) -> Unit
-    ): StateMachineBuilder<S, E, N> =
-        StateMachineBuilder(initialState, transitions, interceptors, postInterceptors.plus(interception))
-
-    /**
-     * @throws InvalidStateMachineException if the configured state machine is not valid. The main
-     * reasons for a state machine not being valid are:
-     * - No initial state
-     * - More than one initial state
-     * - The state machine is not connected (some states are not possible to reach from the initial
-     * state)
-     * - The same source state and event is defined twice
-     */
-    @Throws(InvalidStateMachineException::class)
-    fun build(): StateMachine<S, E, N> {
-        val initState: S = initialState
-            ?: throw InvalidStateMachineException("No initial state set for state machine")
-
-        StateMachineValidator.validate(initState, transitions)
+            return StateMachineImpl(
+                initialState,
+                transitions,
+            )
+        }
+    }
 
-        return StateMachineImpl(
-            initState,
-            transitions,
-            interceptors,
-            postInterceptors
-        )
+    companion object {
+        operator fun <S : Any, E : Any> invoke(): Uninitialized<S, E> = Uninitialized()
     }
 }
diff --git a/lib/src/main/kotlin/io/nexure/fsm/StateMachineImpl.kt b/lib/src/main/kotlin/io/nexure/fsm/StateMachineImpl.kt
index 1c45df4..183413b 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/StateMachineImpl.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/StateMachineImpl.kt
@@ -1,20 +1,15 @@
 package io.nexure.fsm
 
-internal class StateMachineImpl<S : Any, E : Any, N : Any>(
+internal class StateMachineImpl<S : Any, E : Any>(
     private val initialState: S,
-    private val transitions: List<Edge<S, E, N>>,
-    private val interceptors: List<(S, S, E, N) -> (N)>,
-    private val postInterceptors: List<(S, S, E, N) -> Unit>
-) : StateMachine<S, E, N> {
-    private val allowedTransitions: Map<S?, Set<Pair<S, E?>>> = transitions
+    private val transitions: List<Edge<S, E>>,
+) : StateMachine<S, E> {
+    private val allowedTransitions: Map<S, Set<Pair<S, E>>> = transitions
         .groupBy { it.source }
         .map { it.key to it.value.map { edge -> edge.target to edge.event }.toSet() }
         .toMap()
 
-    private val nonTerminalStates: Set<S> = allowedTransitions.keys.filterNotNull().toSet()
-
-    private val transitionActions: Map<Triple<S, E, S>, (N) -> Unit> = transitions
-        .associate { Triple(it.source, it.event, it.target) to it.action }
+    private val nonTerminalStates: Set<S> = allowedTransitions.keys
 
     override fun states(): Set<S> = transitions.asSequence()
         .map { listOf(it.source, it.target) }
@@ -28,31 +23,16 @@ internal class StateMachineImpl<S : Any, E : Any, N : Any>(
     override fun reduceState(events: List<E>): S =
         events.fold(initialState) { state, event -> nextState(state, event) ?: state }
 
-    private fun executeTransition(source: S, target: S, event: E, signal: N) {
-        val action: (N) -> Unit = transitionActions[Triple(source, event, target)] ?: return
-        val interceptedSignal: N = runInterception(source, target, event, signal)
-        action.invoke(interceptedSignal)
-        postIntercept(source, target, event, interceptedSignal)
-    }
-
-    private fun runInterception(source: S, target: S, event: E, signal: N): N {
-        return interceptors.fold(signal) { acc, operation ->
-            operation(source, target, event, acc)
-        }
-    }
-
-    private fun postIntercept(source: S, target: S, event: E, signal: N) {
-        postInterceptors.forEach { intercept -> intercept(source, target, event, signal) }
-    }
-
-    override fun onEvent(state: S, event: E, signal: N): Transition<S> {
+    override fun onEvent(state: S, event: E): Transition<S> {
         val next: S = nextState(state, event) ?: return Rejected
-        executeTransition(state, next, event, signal)
-        return Executed(next)
+        return Accepted(next)
     }
 
+    override fun acceptedEvents(state: S): Set<E> =
+        allowedTransitions.getOrDefault(state, emptySet()).map { it.second }.toSet()
+
     private fun nextState(source: S, event: E): S? {
-        val targets: Set<Pair<S, E?>> = allowedTransitions.getOrDefault(source, emptySet())
+        val targets: Set<Pair<S, E>> = allowedTransitions.getOrDefault(source, emptySet())
         return targets.firstOrNull { it.second == event }?.first
     }
 }
diff --git a/lib/src/main/kotlin/io/nexure/fsm/StateMachineValidator.kt b/lib/src/main/kotlin/io/nexure/fsm/StateMachineValidator.kt
index 98921d8..dcc2916 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/StateMachineValidator.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/StateMachineValidator.kt
@@ -4,7 +4,7 @@ import java.util.Deque
 import java.util.LinkedList
 
 internal object StateMachineValidator {
-    fun <S : Any, E : Any> validate(initialState: S, transitions: List<Edge<S, E, *>>) {
+    fun <S : Any, E : Any> validate(initialState: S, transitions: List<Edge<S, E>>) {
         rejectDuplicates(transitions)
         findIllegalCombinations(transitions)
         isConnected(initialState, transitions)
@@ -13,8 +13,8 @@ internal object StateMachineValidator {
     /**
      * Check so a combination of source state, target state and event is not defined more than once
      */
-    private fun <S : Any, E : Any> rejectDuplicates(transitions: List<Edge<S, E, *>>) {
-        val duplicate: Edge<S, E, *>? = transitions
+    private fun <S : Any, E : Any> rejectDuplicates(transitions: List<Edge<S, E>>) {
+        val duplicate: Edge<S, E>? = transitions
             .duplicatesBy { Triple(it.source, it.target, it.event) }
             .firstOrNull()
 
@@ -35,8 +35,8 @@ internal object StateMachineValidator {
      * These two transitions would not be allowed to exist in the same state machine at the same
      * time.
      */
-    private fun <S : Any, E : Any> findIllegalCombinations(transitions: List<Edge<S, E, *>>) {
-        val illegal: Edge<S, E, *>? = transitions
+    private fun <S : Any, E : Any> findIllegalCombinations(transitions: List<Edge<S, E>>) {
+        val illegal: Edge<S, E>? = transitions
             .groupBy { it.source }
             .filter { it.value.size > 1 }
             .map { x -> x.value.map { y -> x.value.map { it to y } } }
@@ -57,7 +57,7 @@ internal object StateMachineValidator {
      * and event but different target, since a source state which is triggered
      * by a specific event should always result in the same target state.
      */
-    private fun <S : Any, E : Any> illegalCombination(e0: Edge<S, E, *>, e1: Edge<S, E, *>): Boolean {
+    private fun <S : Any, E : Any> illegalCombination(e0: Edge<S, E>, e1: Edge<S, E>): Boolean {
         if (e0 === e1) {
             return false
         }
@@ -70,7 +70,7 @@ internal object StateMachineValidator {
     /**
      * Validate the configuration of the state machine, making sure that state machine is connected
      */
-    private fun <S : Any, E : Any> isConnected(initialState: S, transitions: List<Edge<S, E, *>>) {
+    private fun <S : Any, E : Any> isConnected(initialState: S, transitions: List<Edge<S, E>>) {
         val stateTransitions: Map<S?, Set<S>> = transitions
             .groupBy { it.source }
             .mapValues { it.value.map { value -> value.target }.toSet() }
diff --git a/lib/src/main/kotlin/io/nexure/fsm/Transition.kt b/lib/src/main/kotlin/io/nexure/fsm/Transition.kt
index 6c095bb..c6085eb 100644
--- a/lib/src/main/kotlin/io/nexure/fsm/Transition.kt
+++ b/lib/src/main/kotlin/io/nexure/fsm/Transition.kt
@@ -2,19 +2,19 @@ package io.nexure.fsm
 
 /**
  * Outcome of processing of event by state machine. This will indicate if an event was accepted or
- * rejected by the state machine, or if there was an exception while processing the event.
+ * rejected by the state machine.
  */
 sealed class Transition<out S : Any> {
-    fun transitioned(): Boolean = this is Executed<S>
+    fun transitioned(): Boolean = this is Accepted<S>
 
     /**
      * Returns
-     * - The new state if the transition was permitted and successful ([Executed])
+     * - The new state if the transition was permitted and successful ([Accepted])
      * - `null` if the transitioned was rejected ([Rejected])
      */
     fun stateOrNull(): S? {
         return when (this) {
-            is Executed -> this.state
+            is Accepted -> this.state
             Rejected -> null
         }
     }
@@ -22,9 +22,9 @@ sealed class Transition<out S : Any> {
     /**
      * Invoke this lambda if a transition was executed and successful
      */
-    inline fun onExecution(handle: (state: S) -> Unit): Transition<S> {
-        if (this is Executed) {
-            handle(this.state)
+    inline fun onTransition(action: (state: S) -> Unit): Transition<S> {
+        if (this is Accepted) {
+            action(this.state)
         }
         return this
     }
@@ -41,16 +41,15 @@ sealed class Transition<out S : Any> {
 }
 
 /**
- * The event was successfully processed. Any action associated with the transition, including
- * interceptors, will have been executed successfully.
+ * The event was accepted and a state transition occurred. The [state] property reflects the new state.
  */
-data class Executed<S : Any>(val state: S) : Transition<S>() {
+data class Accepted<S : Any>(val state: S) : Transition<S>() {
     override fun toString(): String = "Executed($state)"
 }
 
 /**
  * The event was rejected because the event was not permitted by the state machine in the current
- * state, which means that no part of any state transition action or interceptors were executed.
+ * state.
  */
 object Rejected : Transition<Nothing>() {
     override fun toString(): String = "Rejected"
diff --git a/lib/src/test/kotlin/io/nexure/fsm/ExampleStateMachine.kt b/lib/src/test/kotlin/io/nexure/fsm/ExampleStateMachine.kt
index e96e7e4..b45ef46 100644
--- a/lib/src/test/kotlin/io/nexure/fsm/ExampleStateMachine.kt
+++ b/lib/src/test/kotlin/io/nexure/fsm/ExampleStateMachine.kt
@@ -18,57 +18,38 @@ enum class PaymentEvent {
     FundsMoved
 }
 
-data class PaymentData(
-    val id: String,
-    val amount: Int
-)
-
-fun buildExampleStateMachine(): StateMachine<PaymentState, PaymentEvent, PaymentData> {
-    return StateMachineBuilder<PaymentState, PaymentEvent, PaymentData>()
+fun buildExampleStateMachine(): StateMachine<PaymentState, PaymentEvent> {
+    return StateMachineBuilder<PaymentState, PaymentEvent>()
         //       ┏━ Initial state
         .initial(PaymentState.Created)
         //       ┏━ Source state       ┏━ Target state       ┏━ Event triggering transition
         .connect(PaymentState.Created, PaymentState.Pending, PaymentEvent.PaymentSubmitted)
-        .connect(PaymentState.Pending, PaymentState.Authorized, PaymentEvent.BankAuthorization) {
-            // Invoke some optional action when payment was authorized
-        }
-        .connect(PaymentState.Pending, PaymentState.Refused, PaymentEvent.BankRefusal) {
-            // Invoke some optional action when payment was refused
-        }
+        .connect(PaymentState.Pending, PaymentState.Authorized, PaymentEvent.BankAuthorization)
+        .connect(PaymentState.Pending, PaymentState.Refused, PaymentEvent.BankRefusal)
         .connect(PaymentState.Authorized, PaymentState.Settled, PaymentEvent.FundsMoved)
-        // This will be called *before* every state transition, with the possibility of altering
-        // the input `signal` if needed
-        .intercept { source, target, event, signal ->
-            println("Will execute transition from $source to $target due to event $event")
-            signal
-        }
-        // This will be called *after* every state transition
-        .postIntercept { source, target, event, _ ->
-            println("Transitioned from $source to $target due to event $event")
-        }
         .build()
 }
 
-fun callExampleStateMachine(fsm: StateMachine<PaymentState, PaymentEvent, PaymentData>) {
-    val payment = PaymentData("foo", 42)
-
+fun callExampleStateMachine(fsm: StateMachine<PaymentState, PaymentEvent>) {
     // Transition from state CREATED into state PENDING
-    val state1 = fsm.onEvent(PaymentState.Created, PaymentEvent.PaymentSubmitted, payment)
-    assertEquals(Executed(PaymentState.Pending), state1)
+    val state1 = fsm.onEvent(PaymentState.Created, PaymentEvent.PaymentSubmitted)
+    assertEquals(Accepted(PaymentState.Pending), state1)
 
     // Transition from state PENDING into state AUTHORIZED
-    val state2 = fsm.onEvent(PaymentState.Pending, PaymentEvent.BankAuthorization, payment)
-    assertEquals(Executed(PaymentState.Authorized), state2)
+    val state2 = fsm.onEvent(PaymentState.Pending, PaymentEvent.BankAuthorization).onTransition {
+        // Invoke some optional action when payment was authorized
+    }
+    assertEquals(Accepted(PaymentState.Authorized), state2)
 
     // Transition from state AUTHORIZED into state SETTLED
-    val state3 = fsm.onEvent(PaymentState.Authorized, PaymentEvent.FundsMoved, payment)
-    assertEquals(Executed(PaymentState.Settled), state3)
+    val state3 = fsm.onEvent(PaymentState.Authorized, PaymentEvent.FundsMoved)
+    assertEquals(Accepted(PaymentState.Settled), state3)
 }
 
 class ExampleStateMachineTest {
     @Test
     fun testExampleStateMachine() {
-        val fsm: StateMachine<PaymentState, PaymentEvent, PaymentData> = buildExampleStateMachine()
+        val fsm: StateMachine<PaymentState, PaymentEvent> = buildExampleStateMachine()
         callExampleStateMachine(fsm)
     }
 }
diff --git a/lib/src/test/kotlin/io/nexure/fsm/StateMachineTest.kt b/lib/src/test/kotlin/io/nexure/fsm/StateMachineTest.kt
index bcca016..49c1123 100644
--- a/lib/src/test/kotlin/io/nexure/fsm/StateMachineTest.kt
+++ b/lib/src/test/kotlin/io/nexure/fsm/StateMachineTest.kt
@@ -1,14 +1,18 @@
 package io.nexure.fsm
 
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertTrue
 import org.junit.Test
 import java.util.concurrent.Semaphore
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class StateMachineTest {
     @Test
     fun `test list of all states`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .connect(State.S2, State.S3, Event.E1)
@@ -20,7 +24,7 @@ class StateMachineTest {
 
     @Test
     fun `test single initial state`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .build()
@@ -28,27 +32,9 @@ class StateMachineTest {
         assertEquals(State.S1, fsm.initialState())
     }
 
-    @Test(expected = InvalidStateMachineException::class)
-    fun `test throws on multiple initial states`() {
-        StateMachine.builder<State, Event, Unit>()
-            .initial(State.S1)
-            .initial(State.S2)
-            .connect(State.S1, State.S2, Event.E2)
-            .connect(State.S2, State.S3, Event.E3)
-            .build()
-    }
-
-    @Test(expected = InvalidStateMachineException::class)
-    fun `test throws on no initial states`() {
-        StateMachine.builder<State, Event, Unit>()
-            .connect(State.S1, State.S2, Event.E2)
-            .connect(State.S2, State.S3, Event.E3)
-            .build()
-    }
-
     @Test
     fun `test list of single terminal state`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .build()
@@ -58,7 +44,7 @@ class StateMachineTest {
 
     @Test
     fun `test list of multiple terminal states`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E2)
             .connect(State.S1, State.S3, Event.E3)
@@ -71,25 +57,21 @@ class StateMachineTest {
     fun `test execution on signal on state change`() {
         var n: Int = 0
 
-        fun plus(signal: Int) {
-            n += signal
-        }
-
-        val fsm = StateMachine.builder<Char, Event, Int>()
+        val fsm = StateMachine.builder<Char, Event>()
             .initial('a')
-            .connect('a', 'b', Event.E2, ::plus)
-            .connect('b', 'c', Event.E3, ::plus)
+            .connect('a', 'b', Event.E2)
+            .connect('b', 'c', Event.E3)
             .build()
 
-        assertEquals(Executed('b'), fsm.onEvent('a', Event.E2, 2))
-        assertEquals(Executed('c'), fsm.onEvent('b', Event.E3, 3))
+        assertEquals(Accepted('b'), fsm.onEvent('a', Event.E2).onTransition { n += 2})
+        assertEquals(Accepted('c'), fsm.onEvent('b', Event.E3).onTransition { n += 3})
 
         assertEquals(5, n)
     }
 
     @Test(expected = InvalidStateMachineException::class)
     fun `test throws on duplicate transition`() {
-        StateMachine.builder<Char, Event, Int>()
+        StateMachine.builder<Char, Event>()
             .initial('a')
             .connect('a', 'b', Event.E1)
             .connect('b', 'c', Event.E2)
@@ -99,7 +81,7 @@ class StateMachineTest {
 
     @Test(expected = InvalidStateMachineException::class)
     fun `test throws on state machine that is not connected`() {
-        StateMachine.builder<Char, Event, Int>()
+        StateMachine.builder<Char, Event>()
             .initial('a')
             .connect('a', 'b', Event.E1)
             .connect('c', 'd', Event.E2)
@@ -108,7 +90,7 @@ class StateMachineTest {
 
     @Test(expected = InvalidStateMachineException::class)
     fun `throw on same source state and event twice for different target state`() {
-        StateMachine.builder<State, Event, Int>()
+        StateMachine.builder<State, Event>()
             .initial(State.S1)
             // Same state (S1) and event (E2) twice, but for different target states
             .connect(State.S1, State.S2, Event.E2)
@@ -118,94 +100,119 @@ class StateMachineTest {
 
     @Test(expected = InvalidStateMachineException::class)
     fun `throw on same source state and event twice for same target state`() {
-        StateMachine.builder<State, Event, Int>()
+        StateMachine.builder<State, Event>()
             .initial(State.S1)
             // Same state (S1) and event (E2) twice with same target state (S2),
             // but with different actions
-            .connect(State.S1, State.S2, Event.E2) {
-                // do X
-            }
-            .connect(State.S1, State.S2, Event.E2) {
-                // do Y
-            }
+            .connect(State.S1, State.S2, Event.E2)
+            .connect(State.S1, State.S2, Event.E2)
             .build()
     }
 
     @Test
-    fun `test interceptor changes signal`() {
-        val fsm = StateMachine.builder<State, Event, Int>()
-            .intercept { _, _, _, signal -> signal * 10 }
-            .initial(State.S1)
-            .connect(State.S1, State.S2, Event.E2) { assertEquals(100, it) }
-            .build()
+    fun `test post interceptor can be built and executed`() {
+        var value: Int = 0
 
-        fsm.onEvent(State.S1, Event.E2, 10)
-    }
+        fun StateMachine<State, Event>.executeWithCustomPostInterception(
+            currentState: State,
+            event: Event
+        ): Transition<State> {
+            return this.onEvent(currentState, event).onTransition { newState ->
+                value += 32
+                println("State changed from $currentState to $newState")
+            }
+        }
 
-    @Test
-    fun `test post interceptor is executed`() {
-        var value: Int = 0
-        val fsm = StateMachine.builder<State, Event, Int>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E2)
-            .postIntercept { _, _, _, signal -> value += signal }
             .build()
 
-        fsm.onEvent(State.S1, Event.E2, 32)
+        fsm.executeWithCustomPostInterception(State.S1, Event.E2)
         assertEquals(32, value)
     }
 
     @Test
     fun `test transition with Action`() {
         val semaphore = Semaphore(10)
-        val fsm = StateMachine.builder<State, Event, Int>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
-            .connect(State.S1, State.S2, Event.E1) { i ->
-                semaphore.tryAcquire(i)
-            }
+            .connect(State.S1, State.S2, Event.E1)
             .build()
 
-        fsm.onEvent(State.S1, Event.E1, 4)
+        fsm.onEvent(State.S1, Event.E1).onTransition {
+            semaphore.tryAcquire(4)
+        }
+
         assertEquals(6, semaphore.availablePermits())
     }
 
     @Test
     fun `test on event runs action`() {
         var executed = false
-        val fsm = StateMachine.builder<State, Event, Boolean>()
+        val fsm = StateMachine.builder<State, Event>()
+            .initial(State.S1)
+            .connect(State.S1, State.S2, Event.E1)
+            .build()
+
+        fsm.onEvent(State.S1, Event.E1).onTransition { executed = true }
+        assertTrue(executed)
+    }
+
+    @Test
+    fun `test on event runs action from non suspend function reference`() {
+        var executed = false
+        fun toggle(state: State) { executed = !executed }
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
-            .connect(State.S1, State.S2, Event.E1) { signal -> executed = signal }
+            .connect(State.S1, State.S2, Event.E1)
             .build()
 
-        fsm.onEvent(State.S1, Event.E1, true)
+        fsm.onEvent(State.S1, Event.E1).onTransition(::toggle)
+        assertTrue(executed)
+    }
+
+    @Test
+    fun `test on event runs action calling suspend function`() = runTest {
+        var executed = false
+        suspend fun toggle() {
+            delay(1L)
+            executed = !executed
+        }
+        val fsm = StateMachine.builder<State, Event>()
+            .initial(State.S1)
+            .connect(State.S1, State.S2, Event.E1)
+            .build()
+
+        fsm.onEvent(State.S1, Event.E1).onTransition { toggle() }
         assertTrue(executed)
     }
 
     @Test
     fun `test going back to initial state is permitted`() {
-        val fsm = StateMachine.builder<State, Event, Boolean>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .connect(State.S2, State.S1, Event.E2)
             .build()
 
-        assertEquals(Executed(State.S2), fsm.onEvent(State.S1, Event.E1, false))
-        assertEquals(Executed(State.S1), fsm.onEvent(State.S2, Event.E2, false))
+        assertEquals(Accepted(State.S2), fsm.onEvent(State.S1, Event.E1))
+        assertEquals(Accepted(State.S1), fsm.onEvent(State.S2, Event.E2))
     }
 
     @Test
     fun `test going back to same state is permitted`() {
-        val fsm = StateMachine.builder<State, Event, Boolean>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S1, Event.E1)
             .build()
 
-        assertEquals(Executed(State.S1), fsm.onEvent(State.S1, Event.E1, false))
+        assertEquals(Accepted(State.S1), fsm.onEvent(State.S1, Event.E1))
     }
 
     @Test
     fun `test reduce state`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .connect(State.S2, State.S3, Event.E2)
@@ -218,7 +225,7 @@ class StateMachineTest {
 
     @Test
     fun `test reduce state with repeated event to be ignored`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .connect(State.S2, State.S3, Event.E2)
@@ -231,7 +238,7 @@ class StateMachineTest {
 
     @Test
     fun `test reduce state with empty list to be resolved to initial state`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .build()
@@ -242,7 +249,7 @@ class StateMachineTest {
 
     @Test
     fun `test reduce state with invalid event for state should be ignored`() {
-        val fsm = StateMachine.builder<State, Event, Unit>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .connect(State.S2, State.S3, Event.E2)
@@ -254,23 +261,47 @@ class StateMachineTest {
 
     @Test
     fun `test rejected transition`() {
-        val fsm = StateMachine.builder<State, Event, Boolean>()
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
             .connect(State.S1, State.S2, Event.E1)
             .build()
 
-        assertEquals(Rejected, fsm.onEvent(State.S1, Event.E3, false))
+        assertEquals(Rejected, fsm.onEvent(State.S1, Event.E3))
     }
 
     @Test(expected = StackOverflowError::class)
     fun `test throwable in action is not caught`() {
         val exception = StackOverflowError("foo")
-        val fsm = StateMachine.builder<State, Event, Boolean>()
+        val fsm = StateMachine.builder<State, Event>()
+            .initial(State.S1)
+            .connect(State.S1, State.S2, Event.E1)
+            .build()
+
+        fsm.onEvent(State.S1, Event.E1).onTransition { throw exception }
+    }
+
+    @Test
+    fun `test acceptedEvents on non-terminal state`() {
+        val fsm = StateMachine.builder<State, Event>()
             .initial(State.S1)
-            .connect(State.S1, State.S2, Event.E1) { throw exception }
+            .connect(State.S1, State.S2, Event.E1)
+            .connect(State.S1, State.S3, Event.E2)
+            .connect(State.S3, State.S4, Event.E3)
+            .build()
+
+        assertEquals(setOf(Event.E1, Event.E2), fsm.acceptedEvents(State.S1))
+    }
+
+    @Test
+    fun `test acceptedEvents is empty on terminal state`() {
+        val fsm = StateMachine.builder<State, Event>()
+            .initial(State.S1)
+            .connect(State.S1, State.S2, Event.E1)
+            .connect(State.S1, State.S3, Event.E2)
+            .connect(State.S3, State.S4, Event.E3)
             .build()
 
-        fsm.onEvent(State.S1, Event.E1, false)
+        assertEquals(emptySet<Event>(), fsm.acceptedEvents(State.S4))
     }
 }