FusionAds-iOS/FusionAds/Classes/fusion/utils/state/StateMachine.swift

// StateMachine.swift
// Core state machine implementation
// Ported from Android implementation

import Foundation

private struct PendingWorkItem {
    let id: UUID
    weak var work: DispatchWorkItem?
    let what: Int
    let obj: Any?
}

/// A hierarchical state machine for managing states and transitions
public class StateMachine {
    // MARK: - Constants
    
    /// Message.what value when quitting
    private static let SM_QUIT_CMD = -1
    
    /// Message.what value when initializing
    private static let SM_INIT_CMD = -2
    
    // MARK: - Properties
    
    /// Name of the state machine, used for logging
    let name: String
    
    /// Current state of the state machine
    private(set) var currentState: IState?
    
    /// Initial state for the state machine
    private var initialState: IState?
    
    /// Destination state for transitions
    private var destState: IState?
    
    private var params: Any?
    
    /// Parent-child state relationships
    private var stateInfos = [String: StateInfo]()
    
    /// Monitors state changes
    private var monitors = [IStateMonitor]()
    
    /// Processing queue
    private let queue = DispatchQueue.main
    
    /// Queue for deferred messages
    private var deferredMessages = [Message]()
    
    /// Flag to indicate if the state machine has quit
    private var hasQuit = false
    
    /// The current message being processed
    private var currentMessage: Message?
    
    /// Debug flag
    private let isDebugEnabled = false
    
    private var pendingWorks: [PendingWorkItem] = []
    
    // MARK: - Nested Types
    
    /// Halting state entered when transitionToHaltingState is called
    private class HaltingState: State {
        private weak var stateMachine: StateMachine?
        
        init(stateMachine: StateMachine) {
            self.stateMachine = stateMachine
            super.init()
        }
        
        override func processMessage(_ msg: Message) -> Bool {
            stateMachine?.haltedProcessMessage(msg)
            return true
        }
    }
    
    /// Quitting state entered when a valid quit message is handled
    private class QuittingState: State {
        override func processMessage(_ msg: Message) -> Bool {
            return false
        }
    }
    
    public func isCurrent(_ state: IState) -> Bool {
        return (currentState as? State) === (state as? State)
    }
    
    /// Information about a state and its relationships
    private class StateInfo {
        let state: IState
        var parentStateInfo: StateInfo?
        var active = false
        
        init(state: IState) {
            self.state = state
        }
    }
    
    // MARK: - Lifecycle
    
    /// Initialize a state machine with a name
    /// - Parameter name: Name of the state machine
    public init(name: String) {
        self.name = name
    }
    
    // MARK: - Public Methods
    
    /// Add a new state to the state machine
    /// - Parameters:
    ///   - state: The state to add
    ///   - parent: Optional parent state
    public func addState(_ state: IState, parent: IState? = nil) {
        let stateInfo = StateInfo(state: state)
        stateInfo.active = false
        
        if let parent = parent {
            // Look up the parent state info
            if let parentStateInfo = getStateInfo(parent) {
                stateInfo.parentStateInfo = parentStateInfo
            } else {
                log("StateMachine: parent state not found: \(parent.name)")
            }
        }
        
        stateInfo.active = false
        stateInfo.parentStateInfo = parent != nil ? getStateInfo(parent!) : nil
        
        stateInfos[state.name] = stateInfo
    }
    
    /// Set the initial state of the state machine
    /// - Parameter state: The initial state
    public func setInitialState(_ state: IState) {
        initialState = state
    }
    
    /// Add a monitor for state changes
    /// - Parameter monitor: The monitor to add
    public func addStateMonitor(_ monitor: IStateMonitor) {
        monitors.append(monitor)
    }
    
    /// Remove a monitor for state changes
    /// - Parameter monitor: The monitor to remove
    public func removeStateMonitor(_ monitor: IStateMonitor) {
        if let index = monitors.firstIndex(where: { $0 as AnyObject === monitor as AnyObject }) {
            monitors.remove(at: index)
        }
    }
    
    /// Start the state machine
    public func start() {
        guard !hasQuit, let initialState = initialState else {
            log("StateMachine: start called without an initial state or after quitting")
            return
        }
        
        queue.async { [weak self] in
            guard let self = self else { return }
            self.setupInitialState()
        }
    }
    
    /// Send a message to the state machine
    /// - Parameter message: The message to send
    public func sendMessage(_ message: Message) {
        sendMessageDelayed(message, delayed: .milliseconds(0))
    }
    
    public func sendMessageDelayed(_ message: Message, delayed: DispatchTimeInterval) {
        if hasQuit {
            log("StateMachine: sendMessage called after quitting")
            return
        }
        let pendingItemId = UUID()
        var work: DispatchWorkItem?
        work = DispatchWorkItem() {
            [weak self] in
            guard let self = self, !(work?.isCancelled ?? true) else {
                self?.pendingWorks.removeAll {
                    item in
                    return pendingItemId == item.id
                }
                return
            }
            if(message.callback != nil) {
                // 如果message携带callback，则视作由callback承接执行过程，模拟Android的`Handler.post(runnable)`或`Handler.postDelayed(runnable, delay)`
                message.callback?()
            } else {
                self.processMessage(message)
            }
            self.pendingWorks.removeAll {
                item in
                return pendingItemId == item.id
            }
        }
        if delayed == DispatchTimeInterval.milliseconds(0) {
            queue.async(execute: work!)
        } else {
            pendingWorks.append(PendingWorkItem(id: pendingItemId, work: work, what: message.what, obj: message.obj))
            queue.asyncAfter(deadline: .now() + delayed, execute: work!)
        }
    }
    
    public func sendMessage(what: Int, obj: Any? = nil) {
        sendMessage(Message(what: what, obj: obj))
    }
    
    public func sendMessageDelayed(what: Int, delayMillis: Int, obj: Any? = nil) {
        
        sendMessageDelayed(Message(what: what, obj: obj), delayed: .milliseconds(delayMillis))
    }
    
    public func sendMessageDelayed(what: Int, delayed: DispatchTimeInterval, obj: Any? = nil) {
        
        sendMessageDelayed(Message(what: what, obj: obj), delayed: delayed)
    }
    
    public func removeMessages(what: Int) {
        if hasQuit {
            log("StateMachine: removeMessages called after quitting")
            return
        }
        pendingWorks.filter {
            pendingWorkItem in
            return pendingWorkItem.work == nil || pendingWorkItem.what == what
        }.forEach {
            pendingWorkItem in
            if pendingWorkItem.work != nil && !(pendingWorkItem.work?.isCancelled ?? true) {
                pendingWorkItem.work?.cancel()
            }
        }
        
    }
    
    private func clearPendingWorks() {
        pendingWorks.forEach {
            pendingWorkItem in
            if pendingWorkItem.work != nil && !(pendingWorkItem.work?.isCancelled ?? true) {
                pendingWorkItem.work?.cancel()
            }
        }
        pendingWorks.removeAll()
    }
    
    public func post(runnable: @escaping () -> Void, token: Any? = nil) {
        postDelayed(runnable: runnable, token: token, delayed: .microseconds(0))
    }
    
    public func postDelayed(runnable: @escaping () -> Void, token: Any? = nil, delayed: DispatchTimeInterval) {
        sendMessageDelayed(Message(what: 0, obj: token, callback: runnable), delayed: delayed)
    }
    
    private func anyEquals(_ x : Any?, _ y : Any?) -> Bool {
        guard let x = x as? AnyHashable,
              let y = y as? AnyHashable else { return false }
        return x == y
    }
    
    public func removeCallbacksAndMessages(_ token: Any? = nil) {
        if(token == nil) {
            clearPendingWorks()
        }
        pendingWorks.removeAll {
            pendingWorkItem in
            let remove = anyEquals(pendingWorkItem.obj, token)
            if(remove && !(pendingWorkItem.work?.isCancelled ?? true)) {
                pendingWorkItem.work?.cancel()
            }
            return remove
        }
    }
    
    /// Send a message to the front of the queue
    /// - Parameter message: The message to send
    public func sendMessageAtFrontOfQueue(_ message: Message) {
        if hasQuit {
            log("StateMachine: sendMessageAtFrontOfQueue called after quitting")
            return
        }
        
        queue.async(flags: .barrier) { [weak self] in
            guard let self = self else { return }
            self.processMessage(message)
        }
    }
    
    /// Transition to a new state
    /// - Parameters:
    ///   - state: The state to transition to
    ///   - params: Optional parameters for the transition
    public func transitionTo(_ state: IState, params: Any? = nil) {
        if hasQuit {
            log("StateMachine: transitionTo called after quitting")
            return
        }
        
        guard let stateInfo = getStateInfo(state) else {
            log("StateMachine: transitionTo called with unrecognized state: \(state.name)")
            return
        }
        
        destState = state
        self.params = params
    }
    
    /// Transition to the halting state
    public func transitionToHaltingState() {
        let haltingState = HaltingState(stateMachine: self)
        addState(haltingState)
        transitionTo(haltingState)
    }
    
    /// Quit the state machine
    public func quit() {
        if hasQuit {
            log("StateMachine: quit called after already quitting")
            return
        }
        
        let quittingState = QuittingState()
        addState(quittingState)
        transitionTo(quittingState)
        
        queue.async(flags: .barrier) { [weak self] in
            guard let self = self else { return }
            self.cleanupAfterQuitting()
        }
    }
    
    /// Defer a message to be processed after the next state transition
    /// - Parameter message: The message to defer
    public func deferMessage(_ message: Message) {
        if hasQuit {
            log("StateMachine: deferMessage called after quitting")
            return
        }
        
        deferredMessages.append(message)
    }
    
    // MARK: - public Methods
    
    /// Called for any message received after transitioning to halting state
    /// - Parameter message: The message received
    open func haltedProcessMessage(_ message: Message) {
        // Default implementation does nothing
    }
    
    /// Called once after transitioning to halting state
    open func onHalting() {
        // Default implementation does nothing
    }
    
    /// Called once after quitting
    open func onQuitting() {
        // Default implementation does nothing
    }
    
    /// Called when a message is not handled by any state
    /// - Parameter message: The unhandled message
    open func unhandledMessage(_ message: Message) {
        // Default implementation does nothing
    }
    
    // MARK: - Private Methods
    
    /// Get the StateInfo for a state
    /// - Parameter state: The state to get info for
    /// - Returns: StateInfo object or nil if not found
    private func getStateInfo(_ state: IState) -> StateInfo? {
        return stateInfos[state.name]
    }
    
    /// Setup the initial state
    private func setupInitialState() {
        guard let initialState = initialState else { return }
        
        currentState = initialState
        
        // Find all parent states of the initial state
        var parentStates = [IState]()
        var stateInfo = getStateInfo(initialState)
        
        while let info = stateInfo?.parentStateInfo {
            parentStates.insert(info.state, at: 0)
            stateInfo = info
        }
        
        // Enter parent states from top to bottom
        for state in parentStates {
            state.enter(from: nil, params: nil)
        }
        
        // Enter initial state
        initialState.enter(from: nil, params: nil)
        
        // Process any deferred messages
        processDeferredMessages()
    }
    
    /// Process a message through the state hierarchy
    /// - Parameter message: The message to process
    private func processMessage(_ message: Message) {
        if hasQuit {
            return
        }
        
        currentMessage = message
        
        if message.what == StateMachine.SM_QUIT_CMD {
            transitionTo(QuittingState())
            return
        }
        
        var state: IState? = currentState
        var handled = false
        
        // Try to handle the message starting with the current state and going up the hierarchy
        while let currentState = state, !handled {
            handled = currentState.processMessage(message)
            if !handled {
                state = getStateInfo(currentState)?.parentStateInfo?.state
            }
        }
        
        if !handled {
            unhandledMessage(message)
        }
        
        performStateTransitions(message)
        
        currentMessage = nil
    }
    
    /// Perform any pending state transitions
    /// - Parameter message: The message that triggered potential transitions
    private func performStateTransitions(_ message: Message) {
        if let destState = destState {
            let orgState = currentState
            
            // Find common ancestor state
            let (commonAncestor, statesToExit, statesToEnter) = findCommonAncestor(currentState!, destState)
            
            // Exit current states up to common ancestor
            for state in statesToExit {
                state.exit(to: destState)
            }
            
            // Enter new states from common ancestor down to destination
            for state in statesToEnter {
                state.enter(from: orgState, params: params)
            }
            
            // Update current state
            currentState = destState
            
            // Process deferred messages
            processDeferredMessages()
            
            // Notify monitors
            notifyMonitors(destState, lastState: orgState, params: message.obj)
            
            // Reset destination state
            self.destState = nil
            self.params = nil
            
            // Check if we're entering a special state
            if destState is QuittingState {
                onQuitting()
                cleanupAfterQuitting()
            } else if destState is HaltingState {
                onHalting()
            }
        }
    }
    
    /// Find the common ancestor of two states and determine which states to exit and enter
    /// - Parameters:
    ///   - currentState: The current state
    ///   - destState: The destination state
    /// - Returns: Tuple containing common ancestor, states to exit, and states to enter
    private func findCommonAncestor(_ currentState: IState, _ destState: IState) -> (IState?, [IState], [IState]) {
        var commonAncestor: IState? = nil
        var statesToExit = [IState]()
        var statesToEnter = [IState]()
        
        // Build path from current state to root
        var currentPath = [IState]()
        var currentStateInfo = getStateInfo(currentState)
        while let info = currentStateInfo {
            currentPath.insert(info.state, at: 0)
            currentStateInfo = info.parentStateInfo
        }
        
        // Build path from destination state to root
        var destPath = [IState]()
        var destStateInfo = getStateInfo(destState)
        while let info = destStateInfo {
            destPath.insert(info.state, at: 0)
            destStateInfo = info.parentStateInfo
        }
        
        // Find common ancestor
        var commonIndex = 0
        while commonIndex < min(currentPath.count, destPath.count) && currentPath[commonIndex].name == destPath[commonIndex].name {
            commonAncestor = currentPath[commonIndex]
            commonIndex += 1
        }
        
        // States to exit - from current state up to (but not including) common ancestor
        statesToExit = Array(currentPath.reversed().prefix(currentPath.count - commonIndex))
        
        // States to enter - from one below common ancestor down to destination state
        statesToEnter = Array(destPath.suffix(destPath.count - commonIndex))
        
        return (commonAncestor, statesToExit, statesToEnter)
    }
    
    /// Process all deferred messages
    private func processDeferredMessages() {
        guard !deferredMessages.isEmpty else { return }
        
        let messages = deferredMessages
        deferredMessages.removeAll()
        
        for message in messages {
            sendMessageAtFrontOfQueue(message)
        }
    }
    
    /// Notify all monitors of a state change
    /// - Parameters:
    ///   - newState: The new state
    ///   - lastState: The previous state
    ///   - params: Optional parameters
    private func notifyMonitors(_ newState: IState, lastState: IState?, params: Any?) {
        for monitor in monitors {
            monitor.onStateChanged(stateMachine: self, state: newState, lastState: lastState, params: params)
        }
    }
    
    /// Clean up resources after quitting
    private func cleanupAfterQuitting() {
        hasQuit = true
        currentState = nil
        initialState = nil
        destState = nil
        deferredMessages.removeAll()
        stateInfos.removeAll()
        monitors.removeAll()
        clearPendingWorks()
    }
    
    /// Log a message if debugging is enabled
    /// - Parameter message: The message to log
    private func log(_ message: String) {
        if isDebugEnabled {
            print("\(name): \(message)")
        }
    }
}