require("melee.melee_common")
require("melee.personalities")

---@class AIPersonality
---@field name string
---@field unit_composition table<UnitId, number>
---@field main_build_order fun(melee_ai: MeleeAI, town: Town)[]
---@field expansion_build_order fun(melee_ai: MeleeAI, town: Town)[]
---@field build_anti_ground_defenses fun(melee_ai: MeleeAI, town: Town)
---@field build_anti_air_defenses fun(melee_ai: MeleeAI, town: Town)
---@field expand_rate? number Default is 3 minutes 30 seconds
---@field expand_wait_after_fail? number Default is 30 seconds
---@field attack_wave_budget_scaling? number Default is 1.0
---@field attack_wave_cost_bias? number Default derived from AI:new_attack()

---@class TownState
---@field started_anti_ground_defenses boolean
---@field started_anti_air_defenses boolean

---@class MeleeAI
---@field debugging {enabled: boolean, last_tick: number, last_script_id: number}
---@field player Player
---@field personality AIPersonality
---@field ai AI
---@field main_town Town
---@field towns table<Town, TownState>
---@field started_expanding boolean
---@field started_attacking boolean
---@field seen_air_threats boolean
MeleeAI = {}
MeleeAI.__index = MeleeAI

setmetatable(MeleeAI, {
    __call = function(cls, ...)
        return cls.new(...)
    end,
})

---@param config? table
---@param player? Player
function MeleeAI.new(config, player)
    local self = setmetatable({}, MeleeAI)
    self.debugging = {
        enabled = PlayerVars(player).ai_debug or false,
        last_tick = -1,
        last_script_id = -1,
    }
    self.player = player or Player()
    self.personality = config or Personalities.pick_random(player)
    self.ai = AI(self.player.id)
    self.main_town = self.ai:start_town()
    self.towns = {
        [self.main_town] = {
            started_anti_ground_defenses = false,
            started_anti_air_defenses = false,
        },
    }
    self.started_expanding = false
    self.started_attacking = false
    self.seen_air_threats = true -- TODO: detect dynamically or after timeout
    -- TODO: maybe have a dynamic unit loss threshold check that triggers a stronger defense response

    self:cdebug(TextColor.Yellow3, "Starting AI with personality: " .. self.personality.name)

    return self
end

function MeleeAI:debug_header()
    if self.debugging.last_tick ~= Time.elapsed or self.debugging.last_script_id ~= Script.id then
        self.debugging.last_tick = Time.elapsed
        self.debugging.last_script_id = Script.id
        print()
        cprint(TextColor.Yellow3, "-- " .. tostring(self.player.name) .. " (P" .. tostring(self.player.id) .. "):\n")
    end
end

function MeleeAI:debug(...)
    if self.debugging.enabled then
        self:debug_header()
        print(...)
    end
end

function MeleeAI:cdebug(color, ...)
    if self.debugging.enabled then
        self:debug_header()
        cprint(color, ...)
    end
end

-- How often the banked-resource snapshot is refreshed while a dynamic_wait is
-- in progress, so a long wait reacts to resources changing mid-wait instead
-- of freezing the scale from a single sample taken at the start.
local DYNAMIC_WAIT_RESAMPLE_INTERVAL = secs(2)

---@param duration integer
function MeleeAI:dynamic_wait(duration)
    local elapsed = 0
    while elapsed < duration do
        local chunk = math.min(DYNAMIC_WAIT_RESAMPLE_INTERVAL, duration - elapsed)
        local resources = self.player.minerals + self.player.gas
        local resource_scale = math.max(0, math.min(1, (1100 - resources) / 600))
        local scaled_chunk = math.floor(chunk * resource_scale)
        if scaled_chunk > 0 then
            wait(scaled_chunk)
        end
        elapsed = elapsed + chunk
    end
end

---@protected
---@return table<UnitId, number>
function MeleeAI:get_final_unit_composition()
    local unit_composition = {}
    for unit_id, preference in pairs(self.personality.unit_composition) do
        unit_composition[unit_id] = preference
    end
    return unit_composition
end

function MeleeAI:run()
    for unit_id, preference in pairs(self:get_final_unit_composition()) do
        if preference > 0 then
            multirun(function()
                wait_until(function() return self.player:can_make(unit_id) end, secs(1))
                self:cdebug(TextColor.Teal, "Added " .. UnitId.name(unit_id) .. " to defense")
                self.ai.defense:auto_set(unit_id)
                wait_while(function() return self.player:can_make(unit_id) end, secs(1))
                self:cdebug(TextColor.Teal, "Added " .. UnitId.name(unit_id) .. " to defense")
                self.ai.defense:auto_remove(unit_id)
            end)
        end
    end

    multirun(function()
        self.main_town:auto_kill_script()
        self.main_town:build(1, MeleeCommon.default_town_building(self.player.race), 120)
        self.main_town:build(self.main_town.max_workers, MeleeCommon.default_worker(self.player.race), 100)
        for i, build_order_step in ipairs(self.personality.main_build_order) do
            build_order_step(self, self.main_town)
            if i == 1 then
                -- If opener didn't enable this, we now should
                self:start_attacking()
                self:start_expanding()
            end
        end
    end)
end

function MeleeAI:start_expanding()
    if self.started_expanding then
        return
    end
    self.started_expanding = true
    multirun(function()
        while true do
            local town = self.ai:expand()
            self:cdebug(TextColor.Green, "Expanding:")
            self:debug(town)
            if town.is_alive then
                self.towns[town] = {
                    started_anti_ground_defenses = false,
                    started_anti_air_defenses = false,
                }
                town:multirun(function()
                    for i, build_order_step in ipairs(self.personality.expansion_build_order) do
                        build_order_step(self, town)
                        if i == 1 then
                            self:start_anti_ground_defenses(town)
                            self:start_anti_air_defenses(town)
                        end
                    end
                end)
                wait(self.personality.expand_rate or (mins(3) + secs(30)))
            else
                wait(self.personality.expand_wait_after_fail or secs(30))
            end
        end
    end)
end

function MeleeAI:start_attacking()
    if self.started_attacking then
        return
    end
    self.started_attacking = true
    multirun(function()
        while true do
            local budget = self.player:unit_count(MeleeCommon.default_worker(self.player.race)) * 50
                + (self.player.minerals + self.player.gas) / 2
            budget = budget * (self.personality.attack_wave_budget_scaling or 1.0)

            local attack = self.ai:new_attack(budget, self.personality.attack_wave_cost_bias or nil)
            for unit_id, preference in pairs(self:get_final_unit_composition()) do
                if preference > 0 and self.player:can_make(unit_id) then
                    attack:compose(preference, unit_id)
                end
            end

            self:cdebug(TextColor.Orange, "Sending attack:")
            self:debug(attack)
            attack:send()
            wait(secs(5))
        end
    end)
end

function MeleeAI:start_anti_ground_defenses(town)
    if self.personality.build_anti_ground_defenses == nil then
        return
    end
    if self.towns[town].started_anti_ground_defenses then
        return
    end
    self.towns[town].started_anti_ground_defenses = true

    multirun(function()
        self.personality.build_anti_ground_defenses(self, town)
    end)
end

function MeleeAI:start_anti_air_defenses(town)
    if self.personality.build_anti_air_defenses == nil then
        return
    end
    if self.towns[town].started_anti_air_defenses then
        return
    end
    self.towns[town].started_anti_air_defenses = true

    multirun(function()
        wait_until(function() return self.seen_air_threats end, secs(1))
        self.personality.build_anti_air_defenses(self, town)
    end)
end