Finite State Machines

/**
 * Finite State Machine actor trait
 * @tparam S - State type (typically sealed trait or enum)
 * @tparam D - State data type
 */
trait FSM[S, D] extends Actor {
  /**
   * Set initial state and data
   * @param stateName - Initial state
   * @param stateData - Initial data
   */
  def startWith(stateName: S, stateData: D): Unit
  
  /**
   * Define behavior for a state
   * @param stateName - State to handle
   * @param stateTimeout - Optional timeout for this state
   * @param stateFunction - Partial function handling events in this state
   */
  def when(stateName: S, stateTimeout: Duration = null)(
    stateFunction: StateFunction
  ): Unit
  
  /**
   * Define behavior for events regardless of state
   * @param stateFunction - Partial function handling events in any state
   */
  def whenUnhandled(stateFunction: StateFunction): Unit
  
  /**
   * Transition to new state
   * @param nextStateName - Target state
   * @return State transition with new state name
   */
  def goto(nextStateName: S): State
  
  /**
   * Stay in current state
   * @return State transition staying in current state
   */
  def stay(): State
  
  /**
   * Stop FSM with reason
   * @param reason - Stop reason
   * @return State transition to stopped state
   */
  def stop(reason: Reason = Normal): State
  
  /**
   * Current state name
   */
  def stateName: S
  
  /**
   * Current state data
   */
  def stateData: D
  
  /**
   * Next state data (during transition)
   */
  def nextStateData: D
  
  /**
   * Initialize FSM - called automatically
   */
  def initialize(): Unit
  
  /**
   * Transform state function with additional logic
   * @param func - Transformation function
   * @return Transformed StateFunction
   */
  def transform(func: StateFunction): StateFunction
}

/**
 * State function type - maps events to state transitions
 */
type StateFunction = PartialFunction[Event, State]

/**
 * Event wrapper containing message and current state data
 */
case class Event(event: Any, stateData: D)

/**
 * State transition result
 */
case class State(
  stateName: S,
  stateData: D,
  timeout: Option[Duration] = None,
  stopReason: Option[Reason] = None,
  replies: List[Any] = Nil
) {
  /**
   * Set timeout for this state
   * @param timeout - State timeout duration
   * @return Updated State
   */
  def forMax(timeout: Duration): State
  
  /**
   * Add reply to be sent
   * @param replyValue - Reply message
   * @return Updated State with reply
   */
  def replying(replyValue: Any): State
  
  /**
   * Set new state data
   * @param nextStateData - New state data
   * @return Updated State with new data
   */
  def using(nextStateData: D): State
}

import akka.actor.FSM
import scala.concurrent.duration._

// Define states and data
sealed trait State
case object Idle extends State
case object Active extends State
case object Stopped extends State

case class Data(items: List[String] = List.empty, count: Int = 0)

// FSM implementation
class WorkerFSM extends FSM[State, Data] {
  
  // Set initial state
  startWith(Idle, Data())
  
  // Define state behaviors
  when(Idle) {
    case Event("start", data) =>
      println("Starting work")
      goto(Active) using data.copy(count = data.count + 1)
      
    case Event("stop", _) =>
      goto(Stopped)
  }
  
  when(Active, stateTimeout = 30.seconds) {
    case Event("work", data) =>
      val newItems = s"item-${data.count}" :: data.items
      stay() using data.copy(items = newItems) replying "work done"
      
    case Event("pause", data) =>
      println("Pausing work")
      goto(Idle) using data
      
    case Event(StateTimeout, data) =>
      println("Work timeout")
      goto(Idle) using data
      
    case Event("stop", _) =>
      goto(Stopped)
  }
  
  when(Stopped) {
    case Event(_, _) =>
      stay() replying "stopped"
  }
  
  // Handle unhandled events
  whenUnhandled {
    case Event("status", data) =>
      stay() replying s"State: $stateName, Data: $data"
    case Event(e, data) =>
      log.warning(s"Unhandled event $e in state $stateName with data $data")
      stay()
  }
  
  // Lifecycle hooks
  onTransition {
    case Idle -> Active =>
      println("Transitioning from Idle to Active")
    case Active -> Idle =>
      println("Transitioning from Active to Idle")
    case _ -> Stopped =>
      println("Stopping FSM")
  }
  
  initialize()
}

/**
 * Set single-shot timer
 * @param name - Timer name (unique identifier)
 * @param msg - Message to send when timer fires
 * @param timeout - Delay before firing
 */
def setTimer(name: String, msg: Any, timeout: Duration): Unit

/**
 * Set repeating timer
 * @param name - Timer name
 * @param msg - Message to send repeatedly
 * @param timeout - Initial delay and repeat interval
 */
def setTimer(name: String, msg: Any, timeout: Duration, repeat: Boolean): Unit

/**
 * Cancel named timer
 * @param name - Timer name to cancel
 */
def cancelTimer(name: String): Unit

/**
 * Check if timer is active
 * @param name - Timer name to check
 * @return true if timer exists and is active
 */
def isTimerActive(name: String): Boolean

/**
 * Cancel all active timers
 */
def cancelAllTimers(): Unit

class TimerFSM extends FSM[State, Data] {
  
  when(Idle) {
    case Event("start", data) =>
      // Set periodic health check
      setTimer("healthCheck", "ping", 10.seconds, repeat = true)
      // Set timeout for this session
      setTimer("sessionTimeout", "timeout", 5.minutes)
      goto(Active) using data
  }
  
  when(Active) {
    case Event("ping", data) =>
      // Health check response
      stay() replying "pong"
      
    case Event("timeout", _) =>
      println("Session timed out")
      cancelAllTimers()
      goto(Stopped)
      
    case Event("stop", _) =>
      cancelAllTimers()
      goto(Stopped)
  }
  
  initialize()
}

/**
 * Register transition handler
 * @param transitionHandler - Partial function handling state transitions
 */
def onTransition(transitionHandler: TransitionHandler): Unit

/**
 * Register termination handler
 * @param terminationHandler - Function called when FSM terminates
 */
def onTermination(terminationHandler: PartialFunction[StopEvent, Unit]): Unit

/**
 * Transition handler type
 */
type TransitionHandler = PartialFunction[(S, S), Unit]

/**
 * Stop event containing reason and state data
 */
case class StopEvent(reason: Reason, currentState: S, stateData: D)

class TransitionFSM extends FSM[State, Data] {
  
  // Transition logging
  onTransition {
    case Idle -> Active =>
      log.info("FSM became active")
      // Initialize resources
      
    case Active -> Idle =>  
      log.info("FSM became idle")
      // Clean up temporary resources
      
    case _ -> Stopped =>
      log.info("FSM stopped")
      // Final cleanup
  }
  
  // Termination handling
  onTermination {
    case StopEvent(Normal, state, data) =>
      log.info(s"FSM terminated normally in state $state")
      
    case StopEvent(Shutdown, state, data) =>
      log.info(s"FSM shutdown in state $state")
      
    case StopEvent(Failure(cause), state, data) =>
      log.error(cause, s"FSM failed in state $state with data $data")
  }
  
  initialize()
}

/**
 * Java-friendly FSM base class
 */
abstract class AbstractFSM[S, D] extends FSM[S, D] {
  /**
   * Java-friendly state function builder
   */
  protected def matchEvent(eventType: Class[_]): UnitMatch[Event]
  
  /**
   * Java-friendly transition builder
   */
  protected def matchState(from: S, to: S): UnitMatch[(S, S)]
}

/**
 * FSM stop reasons
 */
sealed trait Reason

/**
 * Normal termination
 */
case object Normal extends Reason

/**
 * Shutdown requested
 */
case object Shutdown extends Reason

/**
 * Failure with cause
 */
final case class Failure(cause: Any) extends Reason

/**
 * State timeout event
 */
case object StateTimeout

/**
 * FSM configuration and logging
 */
trait FSMConfig {
  /**
   * Enable debug logging of FSM events
   */
  def debugEvent: Boolean = false
  
  /**
   * Log FSM state transitions
   */
  def logTransitions: Boolean = true
}

/**
 * FSM actor ref with state access
 */
trait FSMRef[S] {
  /**
   * Current state of the FSM
   */
  def currentState: S
}

/**
 * Timer management interface
 */
trait Timers {
  def startSingleTimer(key: Any, msg: Any, delay: Duration): Unit
  def startTimerWithFixedDelay(key: Any, msg: Any, delay: Duration): Unit
  def startTimerAtFixedRate(key: Any, msg: Any, interval: Duration): Unit
  def cancel(key: Any): Unit
  def cancelAll(): Unit
  def isTimerActive(key: Any): Boolean
}

// Complex FSM with multiple data types
sealed trait ConnectionState
case object Disconnected extends ConnectionState
case object Connecting extends ConnectionState  
case object Connected extends ConnectionState
case object Reconnecting extends ConnectionState

case class ConnectionData(
  retries: Int = 0,
  lastError: Option[Throwable] = None,
  connectedAt: Option[Long] = None
)

class ConnectionFSM extends FSM[ConnectionState, ConnectionData] {
  
  startWith(Disconnected, ConnectionData())
  
  when(Disconnected) {
    case Event("connect", data) =>
      // Attempt connection
      setTimer("connectTimeout", "connectionFailed", 10.seconds)
      goto(Connecting) using data.copy(retries = data.retries + 1)
  }
  
  when(Connecting) {
    case Event("connected", data) =>
      cancelTimer("connectTimeout")
      val now = System.currentTimeMillis()
      goto(Connected) using data.copy(connectedAt = Some(now), lastError = None)
      
    case Event("connectionFailed", data) if data.retries < 3 =>
      val delay = math.pow(2, data.retries).seconds
      setTimer("reconnect", "connect", delay)
      goto(Disconnected) using data
      
    case Event("connectionFailed", data) =>
      goto(Disconnected) using data replying "Max retries exceeded"
      
    case Event(StateTimeout, data) =>
      goto(Reconnecting) using data
  }
  
  when(Connected) {
    case Event("disconnect", _) =>
      goto(Disconnected) using ConnectionData()
      
    case Event("connectionLost", data) =>
      goto(Reconnecting) using data.copy(lastError = Some(new RuntimeException("Connection lost")))
  }
  
  when(Reconnecting, stateTimeout = 5.seconds) {
    case Event("connect", data) =>
      goto(Connecting) using data
      
    case Event(StateTimeout, data) =>
      self ! "connect"
      stay()
  }
  
  initialize()
}