ground_ai.lua

-- ground_ai.lua - Main entry point for GroundAI system
env.info('[GroundAI] Script loaded', false)

-- Main GroundAI module
GroundAI = {}
GroundAI.groups = {}       -- Store all managed groups
GroundAI.units = {}        -- Store all managed units

-- Debug categories and levels
GroundAI.DEBUG = {
    GENERAL = 1,
    FORMATION = 1,
    ENGAGEMENT = 4,
    TACTIC = 4,
    EVENT = 1,
    FSM = 4
    -- Add more categories as needed
}

-- Debug log function with category support
function GroundAI.log(message, level, category)
    category = category or "GENERAL"
    level = level or 3  -- Default to info level
    local catLevel = GroundAI.DEBUG[category] or 0
    if level <= catLevel then
        local levelStr = "INFO"
        if level == 1 then levelStr = "ERROR"
        elseif level == 2 then levelStr = "WARN"
        elseif level == 4 then levelStr = "DEBUG"
        end
        env.info(string.format("[GroundAI][%s] %s: %s", category, levelStr, message), false)
    end
end



-- Add support for subgroups in group FSM owner
-- Subgroups: table of { name = <string>, units = {unitName=unit}, leader = unit, leaderName = string, fsm = FSM, waypoints = {}, formation = string, spacing = number }

-- Utility: Create a subgroup from a subset of units
function GroundAI.createSubgroup(parentGroupName, subgroupName, unitNames, tactic, waypoints, formation, spacing)
    local parentGroup = GroundAI.groups[parentGroupName]
    if not parentGroup then
        GroundAI.log("Cannot create subgroup: Parent group not found: " .. parentGroupName, 1, "GENERAL")
        return nil
    end
    local parentOwner = parentGroup.fsm:getOwner()
    local units = {}
    for _, unitName in ipairs(unitNames) do
        local unit = parentOwner.units[unitName]
        if unit and unit:isExist() then
            table.insert(units, unit)
        end
    end
    if #units == 0 then
        GroundAI.log("Cannot create subgroup: No valid units", 1, "GENERAL")
        return nil
    end
    local leader = AI_UTILS.group.selectLeader(units)
    local unitRefs = {}
    for _, unit in ipairs(units) do
        unitRefs[unit:getName()] = unit
    end
    -- Create FSM for the subgroup
    local fsm = FSM.createFSM("Subgroup_" .. subgroupName)
    fsm:addStates({
        GroundAI.GROUP_STATES.IDLE,
        GroundAI.GROUP_STATES.MOVING,
        GroundAI.GROUP_STATES.HOLD_POSITION,
        GroundAI.GROUP_STATES.ENGAGING,
        GroundAI.GROUP_STATES.FALL_BACK,
        GroundAI.GROUP_STATES.PUSH_THROUGH,
        GroundAI.GROUP_STATES.ATTACK,
        GroundAI.GROUP_STATES.DEAD
    })
    fsm:setOwner({
        groupName = subgroupName,
        parentGroup = parentGroupName,
        units = unitRefs,
        unitList = units,
        leader = leader,
        leaderName = leader:getName(),
        formation = formation or AI_UTILS.formation.FORMATIONS.WEDGE,
        spacing = spacing or 50,
        waypoints = waypoints or {},
        currentWaypoint = 1,
        enemies = {},
        contactTime = 0,
        inContact = false,
        formationPositions = {},
        waypointReached = false,
        alertState = "relaxed",
        tactic = tactic or TACTICS.hold_and_engage,
        isSubgroup = true
    })
    -- Add transitions and handlers as in createGroupFSM (reuse logic or factor out if needed)
    -- ...existing code for transitions and handlers can be reused here...
    -- Store reference
    if not parentOwner.subgroups then parentOwner.subgroups = {} end
    parentOwner.subgroups[subgroupName] = {
        name = subgroupName,
        units = unitRefs,
        leader = leader,
        leaderName = leader:getName(),
        fsm = fsm
    }
    GroundAI.log("Created subgroup '" .. subgroupName .. "' for group '" .. parentGroupName .. "' with " .. tostring(#units) .. " units", 3, "GENERAL")
    return parentOwner.subgroups[subgroupName]
end

-- Utility: Remove a subgroup
function GroundAI.removeSubgroup(parentGroupName, subgroupName)
    local parentGroup = GroundAI.groups[parentGroupName]
    if not parentGroup or not parentGroup.fsm:getOwner().subgroups then return end
    parentGroup.fsm:getOwner().subgroups[subgroupName] = nil
    GroundAI.log("Removed subgroup '" .. subgroupName .. "' from group '" .. parentGroupName .. "'", 3, "GENERAL")
end

-- Utility: Reassign subgroup leader if current leader is dead
function GroundAI.reassignSubgroupLeader(parentGroupName, subgroupName)
    local parentGroup = GroundAI.groups[parentGroupName]
    if not parentGroup or not parentGroup.fsm:getOwner().subgroups then return end
    local subgroup = parentGroup.fsm:getOwner().subgroups[subgroupName]
    if not subgroup then return end
    local aliveUnits = {}
    for unitName, unit in pairs(subgroup.units) do
        if unit:isExist() then table.insert(aliveUnits, unit) end
    end
    if #aliveUnits == 0 then
        GroundAI.log("Subgroup '" .. subgroupName .. "' has no alive units", 2, "GENERAL")
        return
    end
    local newLeader = AI_UTILS.group.selectLeader(aliveUnits)
    subgroup.leader = newLeader
    subgroup.leaderName = newLeader:getName()
    subgroup.fsm:getOwner().leader = newLeader
    subgroup.fsm:getOwner().leaderName = newLeader:getName()
    GroundAI.log("Reassigned leader for subgroup '" .. subgroupName .. "' to " .. newLeader:getName(), 3, "GENERAL")
end


-- Take control of an existing DCS group, splitting it into single unit groups
function GroundAI.takeControl(groupName, tactic)
    GroundAI.log("Taking control of group: " .. groupName, 2, "GENERAL")
    -- Get original group
    local originalGroup = Group.getByName(groupName)
    if not originalGroup then
        GroundAI.log("Could not find group: " .. groupName, 1, "GENERAL")
        return nil
    end
    
    -- Get list of units and their positions
    local units = {}
    local unitPositions = {}
    local originalUnits = originalGroup:getUnits()
    
    for _, unit in ipairs(originalUnits) do
        table.insert(units, unit)
        local pos = unit:getPosition().p
        local orientation = unit:getPosition().x
        local heading = 0
        if orientation and orientation.x and orientation.z then
            heading = math.atan2(orientation.z, orientation.x)
        end
        unitPositions[unit:getName()] = {pos=pos, heading=heading}
    end
    
    -- Store original group info before destroying
    local originalCoalition = originalGroup:getCoalition()
    local originalCategory = originalGroup:getCategory()
    
    -- Find the country ID from env.mission
    local originalCountry = nil
    for coalitionName, coalition in pairs(env.mission.coalition) do
        if type(coalition) == "table" and coalition.country then
            for _, country in pairs(coalition.country) do
                if type(country) == "table" then
                    -- Search ground and ship groups for matching name
                    local found = false
                    
                    -- Check ground vehicle groups
                    if country.vehicle and country.vehicle.group then
                        for _, group in pairs(country.vehicle.group) do
                            if group.name == groupName then
                                originalCountry = country.id
                                found = true
                                break
                            end
                        end
                    end
                    
                    -- Check ship groups if not found yet
                    if not found and country.ship and country.ship.group then
                        for _, group in pairs(country.ship.group) do
                            if group.name == groupName then
                                originalCountry = country.id
                                found = true
                                break
                            end
                        end
                    end
                    
                    if found then break end
                end
            end
        end
    end
    
    if not originalCountry then
        GroundAI.log("Could not determine country for group: " .. groupName, 1, "GENERAL")
        return nil
    end
    
    -- Create single unit groups for each unit
    local singleUnitGroups = {}
    
    for _, unit in ipairs(units) do
        local unitName = unit:getName()
        local unitType = unit:getTypeName()
        local unitData = unitPositions[unitName]
        local unitPos = unitData and unitData.pos or {x=0, y=0, z=0}
        local heading = unitData and unitData.heading or 0
        
        -- Create a new group with just this unit
        local newGroupData = {
            country = originalCountry,
            category = originalCategory,
            name = unitName .. "_Group",
            task = "Ground Nothing",  -- Will only react to direct commands
            units = {
                [1] = {
                    type = unitType,
                    name = unitName,
                    x = unitPos.x,
                    y = unitPos.z,  -- DCS uses y for north-south
                    heading = heading
                }
            }
        }
        
        -- Destroy the old unit (will be replaced)
        Unit.destroy(unit)
        
        -- Create the new group
        coalition.addGroup(originalCountry, originalCategory, newGroupData)
        
        -- Get reference to the new group
        local newGroup = Group.getByName(unitName .. "_Group")
        if newGroup then
            table.insert(singleUnitGroups, newGroup:getUnit(1))
            
            -- Configure group controller to not engage ground targets automatically
            local controller = newGroup:getController()
            controller:setOption(28, 1) --restrict to air units only to prevent automatic engagement of ground units
        end
    end
    
    -- Create FSMs for units and group
    for _, unit in ipairs(singleUnitGroups) do
        GroundAI.units[unit:getName()] = unit
        local unitFSM = UnitFSM.create(unit)
        unitFSM:getOwner().groupId = groupName
    end
    
    -- Create group FSM
    local groupFSM = GroupFSM.create(groupName, singleUnitGroups, tactic)
    GroundAI.groups[groupName] = {
        originalName = groupName,
        units = singleUnitGroups,
        fsm = groupFSM
    }
    
    -- Start the movement if there are waypoints
    local waypoints = AI_UTILS.group.getOriginalWaypoints(groupName)
    if waypoints and #waypoints > 0 then
        groupFSM:handleEvent("FORMATION_MOVE")
    end

    -- Start periodic updates for the group and all unit FSMs
    if groupFSM and groupFSM.startUpdates then
        groupFSM:startUpdates()
    end
    for _, unit in ipairs(singleUnitGroups) do
        local unitFSM = UnitFSM.Units[unit:getName()]
        if unitFSM and unitFSM.startUpdates then
            unitFSM:startUpdates()
        end
    end
    
    return GroundAI.groups[groupName]
end

-- Utility function to set formation for a group
function GroundAI.setFormation(groupName, formationType, spacing)
    local group = GroundAI.groups[groupName]
    if not group then
        GroundAI.log("Cannot set formation: Group not found: " .. groupName, 1, "GENERAL")
        return false
    end
    
    local owner = group.fsm:getOwner()
    owner.formation = formationType
    if spacing then
        owner.spacing = spacing
    end

    GroundAI.log("Set formation for group '" .. groupName .. "' to '" .. tostring(formationType) .. "'", 2, "GENERAL")
    
    -- Update formation immediately
    GroupFSM.updateFormation(group.fsm)
    return true
end

-- Utility function to set tactic for a group
function GroundAI.setTactic(groupName, tacticEnum)
    local group = GroundAI.groups[groupName]
    if not group then
        GroundAI.log("Cannot set tactic: Group not found: " .. groupName, 1, "GENERAL")
        return false
    end
    local owner = group.fsm:getOwner()
    owner.tacticFactory = TACTICS.getTacticFactory(tacticEnum)
    return true
end

-- Start the GroundAI system
function GroundAI.start()
    GroundAI.log("Starting GroundAI system", 2, "GENERAL")
    -- Initialize event tracking
    GroundAI_Events.init(GroundAI)
end

-- Execute startup code
do
    GroundAI.log("GroundAI startup code running", 4, "GENERAL")
    
    -- Register mission event handlers
    local eventHandler = {}
    function eventHandler:onEvent(event)
        if event.id == world.event.S_EVENT_DEAD or event.id == world.event.S_EVENT_CRASH then
            -- Handle unit destruction
            if event.initiator and event.initiator.getName then
                local unitName = event.initiator:getName()
                if UnitFSM.Units[unitName] then
                    UnitFSM.Units[unitName]:handleEvent("DESTROYED")
                end
            end
        end
    end
    world.addEventHandler(eventHandler)
end