-- AmmoCheck
-- Last modified: 2022.08.17
-- https://github.com/RAX-Anomaly/AmmoCheck

local common = animation_common

do_ammo_check_anim = nil
wpn_checking, det_active = nil


-- settings--
local use_clr = false
local hide_counter = true
local hide_ammo_icon = true
local busy_hands_fix = false
local mcm_key = DIK_keys.DIK_T

-- locals --
local clr00_Red = GetARGB(0xff, 0xff, 0x00, 0x00)
local clr01_RedOrange = GetARGB(0xff, 0xff, 0x40, 0x00)
local clr02_Orange = GetARGB(0xff, 0xff, 0x80, 0x00)
local clr03_Amber = GetARGB(0xff, 0xff, 0xc0, 0x00)
local clr04_LemonGlacier = GetARGB(0xff, 0xff, 0xff, 0x00)
local clr05_LaserLemon = GetARGB(0xff, 0xff, 0xff, 0x80)
local clr06_White = GetARGB(0xff, 0xff, 0xff, 0xff)

local clrEE_Purple = utils_xml.get_color("d_purple", true)

-- the misalignment of the colors and the text means 2x the info without 2x the text.
local messages = {
    { m = "st_ac_near_empty", c = clr00_Red }, -- <1/10
    { m = "st_ac_near_empty", c = clr01_RedOrange }, -- <2/10
    { m = "st_ac_less_half", c = clr01_RedOrange }, -- <3/10
    { m = "st_ac_less_half", c = clr02_Orange }, -- <4/10
    { m = "st_ac_about_half", c = clr02_Orange }, -- <5/10
    { m = "st_ac_about_half", c = clr03_Amber }, -- <6/10
    { m = "st_ac_more_half", c = clr03_Amber }, -- <7/10
    { m = "st_ac_more_half", c = clr04_LemonGlacier }, -- <8/10
    { m = "st_ac_nearly_full", c = clr04_LemonGlacier }, -- <9/10
    { m = "st_ac_nearly_full", c = clr05_LaserLemon }, -- <10/10
    { m = "st_ac_full", c = clr06_White } -- full
}

local dbg_log
function print_dbg(msg, ...)
    if not mcm_log then
        printf("![AC] " .. msg, ...)
    elseif not ui_mcm.MCM_DEBUG then -- use MCM's debug log setting
        return
    else
        if not dbg_log then
            dbg_log = mcm_log.new("![AC]")
            dbg_log.enabled = true
        end
        if dbg_log then dbg_log:log(msg, ...) end
    end
end

function null_function()
    return false
end

function get_mag_loaded_shim(id)
    give_news("You are using an older version of MagsRedux. Update to Github version for best experience.")

    local mag_data = get_data(id)
    return (mag_data and mag_data.section ~= "no_mag") and mag_data or nil
end

local mcm_keybinds = ui_mcm and ui_mcm.key_hold
local modifier = 0
local mode = 0

local modes = {
    [0] = { ["call"] = { "on_key_press", "on_key_hold" }, ["function"] = function(key) ui_mcm.simple_press("rax_ammo_check", key, check_Ammo) end },
    [1] = { ["call"] = { "on_key_press", "on_key_hold" }, ["function"] = function(key) if ui_mcm.double_tap("rax_ammo_check", key) then check_Ammo() end end },
    [2] = { ["call"] = { "on_key_hold", "on_key_press" }, ["function"] = function(key) if ui_mcm.key_hold("rax_ammo_check", key) then check_Ammo() end end }
}

local direction_keys = {
    [key_bindings.kFWD] = true,
    [key_bindings.kBACK] = true,
    [key_bindings.kL_STRAFE] = true,
    [key_bindings.kACCEL] = true,
    [key_bindings.kR_STRAFE] = true
}

local weapon_hidden = false

function on_key_press(key)
    -- do not interrupt if directional key pressed
    local bind = dik_to_bind(key)
    if weapon_hidden and not direction_keys[bind] then
        weapon_hidden = false
    end

    if key ~= mcm_key then return end
    if not mcm_keybinds then
        check_Ammo()
        return
    end
    if ui_mcm.get_mod_key(modifier) then
        modes[mode]["function"](key)
    end
end

function on_key_hold(key)
    if key ~= mcm_key then return end
    if ui_mcm.get_mod_key(modifier) then
        modes[mode]["function"](key)
    end
end

function l_round(value)
    local min = math.floor(value + 0.5)
    return min
end

function on_game_start()
    -- aliases--
    gc = game.translate_string

    if magazine_binder then
        get_data = magazine_binder.get_data
        set_data = magazine_binder.set_data
        get_mag_loaded = magazine_binder.get_mag_loaded or get_mag_loaded_shim -- get_mag_loaded does not exist in moddb version. MagsRedux needs update.
        is_supported_weapon = magazine_binder.is_supported_weapon
        is_jammed_weapon = magazines.is_jammed_weapon
        get_sec_chambered = magazines.get_sec_chambered
        print_dbg("MagsRedux installed. Working in integrated mode.")
    else
        get_data = null_function -- Should not be called, just in case.
        set_data = null_function -- Should not be called, just in case.
        is_supported_weapon = null_function -- Always returns false, meaning none of the weapons are supported, fall back on vanilla ammo handling.
        get_mag_loaded = null_function -- There is never a loaded magazine in vanilla.
        is_jammed_weapon = null_function -- Sadly, standalone mode does not report weapon jams for now.
        get_sec_chambered = null_function -- No "one-in-the-chamber" feature in vanilla.
        print_dbg("MagsRedux not found. Working in standalone mode.")
    end

    RegisterScriptCallback("actor_on_first_update", actor_on_first_update)
    RegisterScriptCallback("on_option_change", on_option_change)
    RegisterScriptCallback("on_option_change", on_option_change)
    RegisterScriptCallback("on_key_press", on_key_press)

    common.add_anim_mutator(mag_check_animation, 2)
end

-- Script Callbacks --

function on_mcm_load()
    ch_options = {
        id = "rax_ammo_check",
        sh = true,
        gr = {
            { id = "ammo_check", type = "slide", link = "ui_options_slider_player", text = "ui_mm_title_rax_ammo_check", size = { 512, 50 }, spacing = 20 },
            { id = "usecolor", type = "check", val = 1, def = false },
            { id = "hidecounter", type = "check", val = 1, def = true },
            { id = "hideicon", type = "check", val = 1, def = true },
            { id = "busy_hands_fix", type = "check", val = 1, def = false },
            { id = "keybind", type = "key_bind", val = 2, def = DIK_keys.DIK_F },
            { id = "modifier", type = ui_mcm.kb_mod_radio, val = 2, def = 0, hint = "mcm_kb_modifier", content = { { 0, "mcm_kb_mod_none" }, { 1, "mcm_kb_mod_shift" }, { 3, "mcm_kb_mod_alt" } } }, --I removed control from the list, note the values of the other options were unchanged.
            { id = "mode", type = ui_mcm.kb_mod_radio, val = 2, def = 2, hint = "mcm_kb_mode", content = { { 0, "mcm_kb_mode_press" }, { 1, "mcm_kb_mode_dtap" }, { 2, "mcm_kb_mode_hold" } } },
            { id = "desc_mcm", type = "desc", text = "ui_mcm_rax_ammo_check_update_mcm", clr = { 255, 175, 0, 0 }, precondition = { function() return not mcm_keybinds end } },
        }
    }
    return ch_options
end

function actor_on_first_update()
    on_option_change()
end

function on_option_change()
    if ui_mcm then
        hide_counter = ui_mcm.get("rax_ammo_check/hidecounter")
        use_clr = ui_mcm.get("rax_ammo_check/usecolor")
        hide_ammo_icon = ui_mcm.get("rax_ammo_check/hideicon")
        busy_hands_fix = ui_mcm.get("rax_ammo_check/busy_hands_fix")
        if mcm_keybinds then
            mcm_key = ui_mcm.get("rax_ammo_check/keybind")
            mode = ui_mcm.get("rax_ammo_check/mode")
            modifier = ui_mcm.get("rax_ammo_check/modifier")
            RegisterScriptCallback(modes[mode]["call"][1], this[modes[mode]["call"][1]])
            UnregisterScriptCallback(modes[mode]["call"][2], this[modes[mode]["call"][2]])
        end
    end

    pos = ActorMenu.get_maingame().m_ui_hud_states.m_ui_weapon_cur_ammo:GetWndPos()
    pos.x = ((hide_counter and pos.x > 0) or ((not hide_counter) and pos.x < 0)) and (-1 * pos.x) or pos.x
    ActorMenu.get_maingame().m_ui_hud_states.m_ui_weapon_cur_ammo:SetWndPos(pos)

    pos = ActorMenu.get_maingame().m_ui_hud_states.m_ui_weapon_icon:GetWndPos()
    pos.x = ((hide_ammo_icon and pos.x > 0) or ((not hide_ammo_icon) and pos.x < 0)) and (-1 * pos.x) or pos.x
    ActorMenu.get_maingame().m_ui_hud_states.m_ui_weapon_icon:SetWndPos(pos)
end

-- Main function --
function checkAmmo() end --crash prevention for dirty updates.

function get_ammo_data(weapon, currentAmmo)
    local sec = weapon:section()
    local mag_data = get_mag_loaded(weapon:id())
    local message = ""
    local clr = nil
    local is_simple_weapon = false

    -- empty/gl use cases
    if (is_jammed_weapon(weapon)) then
        message = gc("st_ac_jammed")
        clr = use_clr and clr00_Red or nil
    elseif weapon:weapon_in_grenade_mode() then
        message = currentAmmo == 1 and gc("st_ac_grenade") or gc("st_ac_empty")
        clr = use_clr and (currentAmmo == 1 and clr06_White or clr00_Red) -- grenade or no grenade that is the question
    elseif currentAmmo == 0 then
        if is_supported_weapon(sec) and not mag_data then
            message = gc("st_ac_noMag")
            clr = use_clr and clr00_Red or nil
        else
            message = gc("st_ac_empty")
            clr = use_clr and clr00_Red or nil
        end
    else
        local top_round = nil
        local max_ammo = 0

        if not mag_data and is_supported_weapon(sec) then
            -- some random crap to force it to say oitc
            currentAmmo = 1
            max_ammo = 10
        elseif is_supported_weapon(sec) then
            print_dbg("checking for section %s", mag_data.section)
            max_ammo = SYS_GetParam(2, mag_data.section, "max_mag_size")
            if magazines.retain_round(weapon) then
                top_round = mag_data.loaded[#mag_data.loaded - 1]
            end
            if not top_round then
                top_round = stack.peek(mag_data.loaded)
            end
        else
            top_round = get_sec_chambered(weapon)
            max_ammo = SYS_GetParam(2, sec, "ammo_mag_size")
        end

        local curAmmoPerc = currentAmmo / max_ammo
        local is_double_barrel = SYS_GetParam(2, sec, "ammo_mag_size", 0) == 2

        -- simple mode for weapons without magazines to only say "loaded" or "empty" (except double-barrel shotguns)
        is_simple_weapon = not is_supported_weapon(sec) and not is_double_barrel
        if is_simple_weapon then
            message = "Loaded"
        elseif currentAmmo == max_ammo + 1 then
            message = gc("st_ac_plus1")
            clr = use_clr and clr06_White -- fullest
        elseif (curAmmoPerc > 1) then
            message = gc("st_ac_overfull")
            clr = clrEE_Purple -- always colored this is an error.
        elseif (currentAmmo == 1 and is_double_barrel) then
            message = gc("st_ac_just_one")
            clr = use_clr and clr03_Amber -- last shell in double-barreled shotty
        elseif currentAmmo == 1 and (not mag_data or (mag_data and #mag_data.loaded == 1)) and is_supported_weapon(sec) then
            message = gc("st_ac_oitc")
            clr = use_clr and clr00_Red -- last round in the chamber with no maggie or empty maggie
        else
            idx = l_round((curAmmoPerc * 10) + .5) -- gives an integer value n for the decile of curAmmoPerc and 11 for full.
            print_dbg("IDX is %s %s %s %s", idx, curAmmoPerc, currentAmmo, max_ammo)
            message = messages[idx] and gc(messages[idx].m) or "Error!"
            clr = use_clr and messages[idx] and messages[idx].c or nil
        end

        if curAmmoPerc > 0 and top_round then
            if is_supported_weapon(sec) then
                message = strformat("%s, %srnd %s", message, max_ammo, gc(ui_item.get_sec_short_name(top_round)))
            else
                message = strformat("%s, %s", message, gc(ui_item.get_sec_short_name(top_round)))
            end
        end
    end

    return message, clr, is_simple_weapon
end


function check_ammo_after_reload()
    local weapon = db.actor:active_item()

    -- ends if no weapon is in hand etc.
    if (weapon == nil or (not IsWeapon(weapon)) or IsItem("fake_ammo_wpn", nil, weapon)) then
        return true
    end

    local currentState = weapon:get_state()
    if not (currentState == 0 or weapon:weapon_in_grenade_mode()) then
        return false
    end

    -- if magazine weapon
    local currentAmmo = weapon:get_ammo_in_magazine()
    if (currentAmmo == nil) then
        return true
    end

    local message = ""
    local clr = nil

    message, clr, is_simple_weapon = get_ammo_data(weapon, currentAmmo)

    disable_info("sleep_active")

    if is_simple_weapon then
        actor_menu.set_msg(1, message, 2, clr)
    else
        actor_menu.set_msg(1, "Loaded " .. message, 2, clr)
    end
    return true
end

function check_Ammo()
    local weapon = db.actor:active_item()

    print_dbg("[AC base] triggering base_ammo_check for weapon %s, %s, %s, %s", weapon, weapon == nil, IsWeapon(weapon), IsItem("fake_ammo_wpn", nil, weapon))  

    -- ends if no weapon is in hand etc.
    if (weapon == nil or (not IsWeapon(weapon)) or IsItem("fake_ammo_wpn", nil, weapon)) then
        return
    end

    local weaponId = weapon:id()

    local currentState = weapon:get_state()
    if not (currentState == 0 or weapon:weapon_in_grenade_mode()) then
        return
    end

    -- if magazine weapon
    local currentAmmo = weapon:get_ammo_in_magazine()
    if (currentAmmo == nil) then
        return
    end

    local message = ""
    local clr = nil
    message, clr = get_ammo_data(weapon, currentAmmo)

    -- local anm_ammo_check = common.has_animation(weapon:section(), "anm_ammo_check")
    local anm_ammo_check = (not weapon:weapon_in_grenade_mode()) and (is_supported_weapon(weapon:section()) or common.has_animation(weapon:section(), "anm_ammo_check"))
    -- anm_ammo_check = true
    if not anm_ammo_check then
        print_dbg("Ammo check for %s, playing holster/unholster animation (fallback 2).", weapon:section())
        local slot = db.actor:active_slot()
        db.actor:activate_slot(0)
        weapon_hidden = true
        CreateTimeEvent("ammo_check", "restore Weapon", 0, function(slot)
            if db.actor:active_item() then
                return false
            end
            if not weapon_hidden then
                db.actor:activate_slot(slot)
                weapon_hidden = false
                return true
            end
            db.actor:activate_slot(slot)

            return true
        end, slot)

        disable_info("sleep_active")
        CreateTimeEvent("ammo_check", "message delay", 1, function()
            actor_menu.set_msg(1, message, 2, clr)
            return true
        end)

        return
    end

    if do_ammo_check_anim then return end -- If an ammo check is already happening, don't do anything

    disable_info("sleep_active")
    CreateTimeEvent("ammo_check", "message delay", 1.5, function()
        actor_menu.set_msg(1, message, 2, clr)
        return true
    end)

    if common.has_animation(weapon:section(), "anm_bore") then
        print_dbg("Ammo check for %s, playing inspection animation.", weapon:section())
        weapon:switch_state(4)
        do_ammo_check_anim = 1
    else
    -- For weapons without inspection animations, we need to trigger a fake reload and reset the ammo
        print_dbg("Ammo check for %s, playing reload animation (fallback 1).", weapon:section())
        play_reload_animation(weapon)
        do_ammo_check_anim = nil
    end

    det_active = db.actor:active_detector() or nil
    if det_active then
        det_active:switch_state(2)
    end
end

function play_reload_animation(weapon)
    local wpn_id = weapon:id()
    local current_magazine = get_mag_loaded(wpn_id)
    -- If one in the chamber, reloading will substract a box of ammo and creat a new dummy mag
    if not current_magazine and magazines.retain_round(weapon) and (weapon:get_ammo_in_magazine() > 0) then
        local oitc_ammo_sec = magazines.get_sec_chambered(weapon)
        local oitc_ammo_type = weapon:get_ammo_type()
        print_dbg("One (%s: %s) in the chamber and no magazine loaded.", oitc_ammo_type, oitc_ammo_sec)   
        local pre_table = magazines.count_ammo(weapon)
        weapon:set_ammo_elapsed(0)
        weapon:switch_state(7)
        CreateTimeEvent("OneKey_AmmoCheck", "dummy_reload"..wpn_id, 0.1, restore_ammo_after_reload, weapon, 1, pre_table)
    else
        local pre_table = magazines.count_ammo(weapon)
        weapon:switch_state(7)
        CreateTimeEvent("OneKey_AmmoCheck", "dummy_reload"..wpn_id, 0.1, restore_ammo_after_reload, weapon, weapon:get_ammo_in_magazine(), pre_table)
    end
end

function restore_ammo_after_reload(weapon, count, pre_table)
    if weapon:get_state() == 0 then
        weapon:set_ammo_elapsed(count)
        magazines.refund_ammo(weapon, pre_table)
        print_dbg("Ammo refunded, count = %s.", count)
        return true
    end
    return false
end

-- Catch the inspect weapon state switch and do the ammo check animation instead
function mag_check_animation(anm_table, obj)
    if wpn_checking and obj and wpn_checking:id() == obj:id() and do_ammo_check_anim == 2 then
        wpn_checking = nil
        do_ammo_check_anim = nil

        if det_active then
            db.actor:show_detector(true)
        end
        det_active = nil
    end

    if do_ammo_check_anim == 1 and obj and obj:get_state() == 4 then
        wpn_checking = obj

        local is_supported, has_mag_data
        if magazine_binder then
            is_supported = is_supported_weapon(obj:section())
            has_mag_data = get_mag_loaded(obj:id())
        end
        if is_supported and not has_mag_data and common.has_animation(obj:section(), "anm_ammo_check_no_mag") then
            anm_table.anm_name = "anm_ammo_check_no_mag"
        elseif obj:get_ammo_in_magazine() == 0 and common.has_animation(obj:section(), "anm_ammo_check_empty") then
            anm_table.anm_name = "anm_ammo_check_empty"
        elseif common.has_animation(obj:section(), "anm_ammo_check") then
            anm_table.anm_name = "anm_ammo_check"
        else
            anm_table.anm_name = "anm_reload"
        end
        do_ammo_check_anim = 2
    end
end

local old_news = {}

function give_news(message)
    if old_news[message] then return end

    old_news[message] = true

    print_dbg(message)
    if db.actor then
        db.actor:give_game_news("AmmoCheck", message, "ui_inGame2_D_Ohotnik_na_mutantov", 0, 5000, 0)
    end
end