FSM.lua

FSM={}
function FSM.createFSM(name)
    local fsm = {
        name = name,
        states = {},
        currentState = nil,
        previousState = nil,
        transitions = {},
        guards = {},
        enterHandlers = {},
        exitHandlers = {},
        stateHandlers = {},
        owner = nil,
        updateFrequency = 1, -- Default update frequency in seconds
        timerId = nil,       -- Store the timer ID for cleanup
        stateGroups = {},    -- State groups for logical grouping
        globalEvents = {},   -- Global events that apply to all states
        parentStates = {}    -- Parent-child relationships for hierarchical states
    }
    
    -- Add state to the FSM
    function fsm:addState(stateName)
        if type(stateName) ~= "string" or stateName == "" then
            GroundAI.log("Error: State name must be a non-empty string")
            return self
        end
        self.states[stateName] = stateName
        return self
    end
    
    -- Add state with parent (for hierarchical states)
    function fsm:addStateWithParent(stateName, parentState)
        self:addState(stateName)
        if parentState and self.states[parentState] then
            self.parentStates[stateName] = parentState
        elseif parentState then
            GroundAI.log("Error: Parent state " .. parentState .. " does not exist")
        end
        return self
    end
    
    -- Add state to a group
    function fsm:addStateToGroup(stateName, groupName)
        if not self.states[stateName] then
            GroundAI.log("Error: Cannot add non-existent state " .. stateName .. " to group")
            return self
        end
        
        if not self.stateGroups[groupName] then
            self.stateGroups[groupName] = {}
        end
        
        self.stateGroups[groupName][stateName] = true
        return self
    end
    
    -- Check if current state is in a group
    function fsm:isInGroup(groupName)
        if not self.stateGroups[groupName] then
            return false
        end
        
        return self.stateGroups[groupName][self.currentState] == true
    end
    
    -- Get parent state of current state
    function fsm:getParentState(stateName)
        stateName = stateName or self.currentState
        return self.parentStates[stateName]
    end
    
    -- Check if state has a specific ancestor
    function fsm:hasAncestor(stateName, ancestorState)
        local currentState = stateName
        while currentState do
            currentState = self.parentStates[currentState]
            if currentState == ancestorState then
                return true
            end
        end
        return false
    end
    
    -- Add multiple states
    function fsm:addStates(stateNames)
        for _, state in ipairs(stateNames) do
            self:addState(state)
        end
        return self
    end
    
    -- Set initial state
    function fsm:setInitialState(stateName)
        if self.states[stateName] then
            self.currentState = stateName
            -- Execute enter handler if exists
            if self.enterHandlers[stateName] then
                self.enterHandlers[stateName](self.owner)
            end
        else
            GroundAI.log("Error: Cannot set initial state to non-existent state " .. stateName)
        end
        return self
    end
    
    -- Add a transition between states
    function fsm:addTransition(fromState, toState, eventName)
        -- Special case for "ANY" state
        if fromState == "ANY" then
            -- Create an entry for this event for all states
            for stateName, _ in pairs(self.states) do
                if not self.transitions[stateName] then
                    self.transitions[stateName] = {}
                end
                
                if not self.transitions[stateName][eventName] then
                    self.transitions[stateName][eventName] = toState
                    GroundAI.log("Added transition from " .. stateName .. " to " .. toState .. " on event " .. eventName)
                else
                    GroundAI.log("Warning: Overriding existing transition from " .. stateName .. " on event " .. eventName)
                    self.transitions[stateName][eventName] = toState
                end
            end
            
            return self
        end
        
        -- Normal state transition
        if not self.states[fromState] then
            GroundAI.log("Error: Cannot add transition from non-existent state " .. fromState)
            return self
        end
        
        if not self.states[toState] then
            GroundAI.log("Error: Cannot add transition to non-existent state " .. toState)
            return self
        end
        
        if not self.transitions[fromState] then
            self.transitions[fromState] = {}
        end
        
        if not self.transitions[fromState][eventName] then
            self.transitions[fromState][eventName] = toState
        else
            GroundAI.log("Warning: Overriding existing transition from " .. fromState .. " on event " .. eventName)
            self.transitions[fromState][eventName] = toState
        end
        
        return self
    end
    
    -- Add a guard to a transition
    function fsm:addGuard(fromState, eventName, guardFunction)
        if not self.guards[fromState] then
            self.guards[fromState] = {}
        end
        self.guards[fromState][eventName] = guardFunction
        return self
    end
    
    -- Add a state entry handler
    function fsm:onEnter(stateName, handlerFunction)
        self.enterHandlers[stateName] = handlerFunction
        return self
    end
    
    -- Add a state exit handler
    function fsm:onExit(stateName, handlerFunction)
        self.exitHandlers[stateName] = handlerFunction
        return self
    end
    
    -- Add a state handler (executed while in the state)
    function fsm:onState(stateName, handlerFunction)
        self.stateHandlers[stateName] = handlerFunction
        return self
    end
    
    -- Add a global event handler (executed regardless of current state)
    function fsm:addGlobalEvent(eventName, handlerFunction)
        self.globalEvents[eventName] = handlerFunction
        return self
    end
    
    -- Force state change without triggering an event
    function fsm:changeState(newState)
        if not self.states[newState] then
            GroundAI.log("Error: Cannot change to non-existent state " .. newState)
            return false
        end
        
        local currentState = self.currentState
        
        -- Execute exit handler for current state if exists
        if currentState and self.exitHandlers[currentState] then
            self.exitHandlers[currentState](self.owner, newState)
        end
        
        -- Change state
        self.previousState = currentState
        self.currentState = newState
        
        -- Execute enter handler for new state if exists
        if self.enterHandlers[newState] then
            self.enterHandlers[newState](self.owner, currentState)
        end
        
        return true
    end
    
    -- Handle an event
    function fsm:handleEvent(eventName, ...)
        -- Check if this is a global event
        if self.globalEvents[eventName] then
            self.globalEvents[eventName](self.owner, ...)
            -- Note: Global events don't prevent state transitions
        end
        
        local currentState = self.currentState
        
        if not currentState then
            GroundAI.log("Error: No current state set for FSM " .. self.name)
            return false
        end
        
        -- Check current state for transition
        if self.transitions[currentState] and self.transitions[currentState][eventName] then
            -- Check guard if exists
            if self.guards[currentState] and self.guards[currentState][eventName] then
                if not self.guards[currentState][eventName](self.owner, ...) then
                    -- Guard prevented transition
                    return false
                end
            end
            
            local nextState = self.transitions[currentState][eventName]
            
            -- Execute exit handler for current state if exists
            if self.exitHandlers[currentState] then
                self.exitHandlers[currentState](self.owner, nextState, ...)
            end
            
            -- Change state
            self.previousState = currentState
            self.currentState = nextState
            
            -- Execute enter handler for new state if exists
            if self.enterHandlers[nextState] then
                self.enterHandlers[nextState](self.owner, currentState, ...)
            end
            
            return true
        end
        
        -- If no transition found in current state, check parent states (hierarchical FSM)
        local parentState = self.parentStates[currentState]
        while parentState do
            if self.transitions[parentState] and self.transitions[parentState][eventName] then
                -- Check guard if exists
                if self.guards[parentState] and self.guards[parentState][eventName] then
                    if not self.guards[parentState][eventName](self.owner, ...) then
                        -- Guard prevented transition
                        return false
                    end
                end
                
                local nextState = self.transitions[parentState][eventName]
                
                -- Execute exit handler for current state if exists
                if self.exitHandlers[currentState] then
                    self.exitHandlers[currentState](self.owner, nextState, ...)
                end
                
                -- Change state
                self.previousState = currentState
                self.currentState = nextState
                
                -- Execute enter handler for new state if exists
                if self.enterHandlers[nextState] then
                    self.enterHandlers[nextState](self.owner, currentState, ...)
                end
                
                return true
            end
            
            -- Move up the hierarchy
            parentState = self.parentStates[parentState]
        end
        
        -- No transition found in current state or any parent state
        return false
    end
    
    -- Update the FSM (run the current state handler)
    function fsm:update(dt)
        local currentState = self.currentState
        
        if not currentState then
            return
        end
        
        -- Execute current state handler if exists
        if self.stateHandlers[currentState] then
            self.stateHandlers[currentState](self.owner, dt)
        end
        
        -- Execute parent state handlers (hierarchical FSM)
        local parentState = self.parentStates[currentState]
        while parentState do
            if self.stateHandlers[parentState] then
                self.stateHandlers[parentState](self.owner, dt)
            end
            parentState = self.parentStates[parentState]
        end
    end
    
    -- Get current state
    function fsm:getState()
        return self.currentState
    end
    
    -- Debug function to print FSM structure
    function fsm:debugPrint()
        GroundAI.log("FSM: " .. self.name)
        GroundAI.log("Current state: " .. (self.currentState or "none"))
        GroundAI.log("States: ")
        
        for state, _ in pairs(self.states) do
            local parentInfo = ""
            if self.parentStates[state] then
                parentInfo = " (parent: " .. self.parentStates[state] .. ")"
            end
            
            GroundAI.log("  - " .. state .. parentInfo)
            
            -- Print transitions
            if self.transitions[state] then
                for event, nextState in pairs(self.transitions[state]) do
                    local guardInfo = ""
                    if self.guards[state] and self.guards[state][event] then
                        guardInfo = " [GUARDED]"
                    end
                    GroundAI.log("    → " .. event .. " → " .. nextState .. guardInfo)
                end
            end
            
            -- Print group memberships
            for groupName, states in pairs(self.stateGroups) do
                if states[state] then
                    GroundAI.log("    ∈ Group: " .. groupName)
                end
            end
        end
        
        -- Print global events
        GroundAI.log("Global Events: ")
        for event, _ in pairs(self.globalEvents) do
            GroundAI.log("  - " .. event)
        end
    end
    
    -- Set the FSM owner (object that the FSM controls)
    function fsm:setOwner(owner)
        self.owner = owner
        return self
    end

    -- Get the FSM owner (object that the FSM controls)
    function fsm:getOwner()
        return self.owner
    end
    
    -- Set update frequency in seconds
    function fsm:setUpdateFrequency(frequency)
        self.updateFrequency = frequency
        return self
    end
    
    -- Start automatic updates
    function fsm:startUpdates()
        -- Stop any existing update timer
        self:stopUpdates()
        
        -- Create an update function closure
        local function updateFSM()
            self:update(self.updateFrequency)
            return timer.getTime() + self.updateFrequency -- Return the time for next update
        end
        
        -- Schedule first update
        self.timerId = timer.scheduleFunction(updateFSM, nil, timer.getTime() + self.updateFrequency)
        return self
    end
    
    -- Stop automatic updates
    function fsm:stopUpdates()
        if self.timerId then
            timer.removeFunction(self.timerId)
            self.timerId = nil
        end
        return self
    end
    
    return fsm
end