State Machines with Enums and Sealed Classes

enum class OrderStatus {
    PENDING, CONFIRMED, SHIPPED, DELIVERED, CANCELLED;

    fun canTransitionTo(next: OrderStatus): Boolean = when (this) {
        PENDING -> next in setOf(CONFIRMED, CANCELLED)
        CONFIRMED -> next in setOf(SHIPPED, CANCELLED)
        SHIPPED -> next == DELIVERED
        DELIVERED -> false
        CANCELLED -> false
    }
}

sealed class ConnectionState {
    object Disconnected : ConnectionState()
    data class Connecting(val attempt: Int) : ConnectionState()
    data class Connected(val sessionId: String, val connectedAt: Long) : ConnectionState()
    data class Error(val message: String, val retryAfter: Long) : ConnectionState()

    fun display(): String = when (this) {
        is Disconnected -> "Disconnected"
        is Connecting -> "Connecting (attempt $attempt)..."
        is Connected -> "Connected (session: $sessionId)"
        is Error -> "Error: $message (retry in ${retryAfter}ms)"
    }
}

class StateMachine<S, E>(
    private var state: S,
    private val transitions: Map<Pair<S, E>, (S) -> S>
) {
    fun current() = state
    fun transition(event: E): S {
        val key = Pair(state, event)
        val handler = transitions[key]
            ?: throw IllegalStateException("No transition from $state on $event")
        state = handler(state)
        return state
    }

    class Builder<S, E> {
        private val t = mutableMapOf<Pair<S, E>, (S) -> S>()
        fun on(state: S, event: E, action: (S) -> S) { t[Pair(state, event)] = action }
        fun build(initial: S) = StateMachine(initial, t)
    }
}

fun <S, E> stateMachine(initial: S, block: StateMachine.Builder<S, E>.() -> Unit) =
    StateMachine.Builder<S, E>().apply(block).build(initial)

enum class Light { RED, YELLOW, GREEN }
enum class Timer { TICK }

fun main() {
    var status = OrderStatus.PENDING
    if (status.canTransitionTo(OrderStatus.CONFIRMED)) {
        status = OrderStatus.CONFIRMED
        println("Order: $status")
    }

    var conn: ConnectionState = ConnectionState.Disconnected
    println(conn.display())
    conn = ConnectionState.Connecting(1)
    println(conn.display())
    conn = ConnectionState.Connected("sess-abc", System.currentTimeMillis())
    println(conn.display())

    val trafficLight = stateMachine<Light, Timer>(Light.RED) {
        on(Light.RED, Timer.TICK) { Light.GREEN }
        on(Light.GREEN, Timer.TICK) { Light.YELLOW }
        on(Light.YELLOW, Timer.TICK) { Light.RED }
    }
    repeat(6) {
        println("Light: ${trafficLight.current()}")
        trafficLight.transition(Timer.TICK)
    }
}