local mcm_key_reload = DIK_keys.DIK_R
local mcm_modifier_unjam = 1
local mcm_doubletap_behavior = "situational"
local mcm_check_ammo_after_reload = false

function on_game_start()
    RegisterScriptCallback("actor_on_first_update", actor_on_first_update)
    RegisterScriptCallback("on_option_change", on_option_change)
    RegisterScriptCallback("on_before_key_press", on_before_key_press)
    RegisterScriptCallback("on_key_hold", on_key_hold)
end

function on_mcm_load()
    op = {
        id = "reload",
        sh = true,
        gr = {
            {id = "slide_onekey_reload", type = "slide", link = "ui_options_slider_player", text = "ui_mcm_title_onekey_reload", size = {512, 50}, spacing = 20},            
            {id = "desc_onekey_reload", type = "desc", text = "ui_mcm_onekey_key_reload_desc" },            
            {id = "key_reload", type = "key_bind", val = 2, def = DIK_keys.DIK_R, hint = "onekey_key_reload" },
            {id = "modifier_unjam", type = "radio_h", val = 2, def = 1, hint = "onekey_unjam_modifier", content = {{1,"mcm_kb_mod_shift"} , {2,"mcm_kb_mod_ctrl"}, {3,"mcm_kb_mod_alt"}}},
            {id = "doubletap_behavior", type = "radio_h", val = 0, def = "situational", hint = "onekey_doubletap_behavior", content = {{"reload", "dt_reload"}, {"next_ammo_type", "dt_next_ammo_type"}, {"ammo_wheel", "dt_ammo_wheel"}, {"situational", "dt_situational"}}},
            {id = "check_ammo_after_reload", type = "check", val = 1, def = false, hint = "onekey_check_ammo_after_reload"},
        }
    }
    return op, "onekey"
end

function actor_on_first_update()
    on_option_change()
end

function on_option_change()
    if ui_mcm then
        mcm_key_reload = ui_mcm.get("onekey/reload/key_reload")
        mcm_modifier_unjam = ui_mcm.get("onekey/reload/modifier_unjam")
        mcm_doubletap_behavior = ui_mcm.get("onekey/reload/doubletap_behavior")
        mcm_check_ammo_after_reload = ui_mcm.get("onekey/reload/check_ammo_after_reload")
    end
    local show = false
    local pos = ActorMenu.get_maingame().m_ui_hud_states.m_ui_weapon_cur_ammo:GetWndPos()
    pos.x = ((not show and pos.x > 0) or (show 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)
end

function on_before_key_press(key, bind, dis, flags)
    if bind == key_bindings.kWPN_RELOAD or key == mcm_key_reload then
        -- passthrough if no weapon active
        local weapon = db.actor:active_item()
        if (weapon == nil or (not IsWeapon(weapon)) or IsItem("fake_ammo_wpn", nil, weapon)) then
            flags.ret_value = true
            return
        end

        if is_ammo_wheel_open() or weapon:get_state() ~= 0 then
            flags.ret_value = true
            return
        end

        -- Unjam or reload depending on modifier
        if ui_mcm.get_mod_key(mcm_modifier_unjam) then
            ui_mcm.simple_press("magazines_onekey_press_unjam", key, try_unjam_weapon)
        else
            ui_mcm.simple_press("magazines_onekey_press_reload", key, try_reload_weapon)
        end

        -- Double-tap to ammo-select
        if ui_mcm.double_tap("magazines_onekey_doubletap_changeammo", key) then
            -- printf("[OKR] doubletap: behavior = %s", mcm_doubletap_behavior)
            if mcm_doubletap_behavior == "situational" then
                _, _, num_avail = get_ammo_data(weapon, true)
                -- printf("[OKR] doubletap: behavior = %s, num_avail == %s", mcm_doubletap_behavior, num_avail)
                if num_avail > 2 then
                    switch_ammo_wheel()
                elseif num_avail == 2 then
                    switch_ammo_next()
                else
                    try_reload_weapon()
                end
            elseif mcm_doubletap_behavior == "ammo_wheel" then
                switch_ammo_wheel()
            elseif mcm_doubletap_behavior == "next_ammo_type" then
                switch_ammo_next()
            else -- reload
                try_reload_weapon()
            end
        end

        -- Suppress the keypress since it has been handled.
        flags.ret_value = false
    end
end

function on_key_hold(key)
    local bind = dik_to_bind(key)
    -- Hold reload key to ammo check
    if ((bind == key_bindings.kWPN_RELOAD or key == mcm_key_reload) and ui_mcm.key_hold("magazines_onekey_hold", key)) then
        ammo_check_mcm.check_Ammo()
    end
end

function try_unjam_weapon()
    local weapon = db.actor:item_in_slot(db.actor:active_slot())
    if weapon then
        arti_jamming.unjam(weapon)
    end
end

function try_reload_weapon()
    local weapon = db.actor:item_in_slot(db.actor:active_slot())
    -- Don't reload unless weapon is in idle state
    local state = weapon and weapon:get_state()
    if not weapon or weapon:weapon_in_grenade_mode() or state ~= 0 then
        return
    end

    if magazine_binder.is_supported_weapon(weapon) then
        magazines.actor_on_weapon_reload(db.actor, weapon)
    else
        db.actor:reload_weapon()
    end
    if mcm_check_ammo_after_reload then
        if (not magazine_binder.is_supported_weapon(weapon)) and (weapon:get_ammo_in_magazine() >= SYS_GetParam(2, weapon:section(), "ammo_mag_size")) then
            -- Do not ammo check for non-magazined fully loaded weapons (shotguns, revolvers, bolt actions, etc) since those won't be reloaded
            return
        end
        CreateTimeEvent("onekey_reload", "ammo_check_on_reload", 0.1, ammo_check_mcm.check_ammo_after_reload)
    end
end

-- Ammo type select
function switch_ammo_wheel()
    item_weapon.start_ammo_wheel()
end

function is_ammo_wheel_open()
    return (item_weapon.GUI and item_weapon.GUI:IsShown())
end

function switch_ammo_next()
    local success = try_switch_next_ammo()
    if not success then
        try_reload_weapon()
    end
end

function try_switch_next_ammo()
    local weapon = db.actor:active_item()
    if weapon and IsWeapon(weapon) and (not IsItem("fake_ammo_wpn", nil, weapon)) then
        local next_ammo_type
        local ammo_types
        local has_ammo_types
        ammo_types, has_ammo_types = get_ammo_data(weapon, false)
        local current_ammo_type = get_current_ammo_type(weapon, ammo_types) -- a number
        local ammo_map = invert_table(magazines.get_ammo_map(weapon:id()))
        local success = false

        -- printf("[OKR] current_ammo_type[%s] = %s", current_ammo_type, ammo_map[current_ammo_type])

        -- Find the next available ammo type
        for i = current_ammo_type + 1, #ammo_types do -- +1 because we need next type
            if has_ammo_types[i] then
                next_ammo_type = i
                break
            end
        end

        -- Loop around and find the next available ammo type
        if not next_ammo_type then
            for i = 1, current_ammo_type - 1 do
                if has_ammo_types[i] then
                    next_ammo_type = i
                    break
                end
            end
        end

        if next_ammo_type then
            -- printf("[OKR] found next_ammo_type[%s] = %s", next_ammo_type, ammo_map[next_ammo_type])
            success = try_switch_to_ammo_type(weapon, next_ammo_type)
        end
        return success
    end
end

function try_switch_to_ammo_type(weapon, ammo_type)
    if magazines and magazine_binder.is_supported_weapon(weapon) then
        local ammo_map = utils_item.get_ammo(nil, weapon:id())
        local magazine = magazines.find_magazine(weapon, ammo_map[ammo_type])
        if magazines.get_mag_data(weapon:id()) then
            magazines.eject_magazine(weapon)
        end
        local pre_table = magazines.count_ammo(weapon)

        local first_round = nil
        if magazines.retain_round(weapon) and weapon:get_ammo_in_magazine() > 0 then
            first_round = magazines.get_sec_chambered(weapon)
        end

        weapon:switch_state(7)
        disable_info("sleep_active")
        magazines.action_start_reload()
        CreateTimeEvent("mag_redux", "delay_weapon"..weapon:id(), 0.1, magazines.delay_load_weapon, weapon, magazine, pre_table, first_round)
    else
        weapon:unload_magazine(true)
        weapon:set_ammo_type(ammo_type - 1) -- ammotype is zero-indexed
        db.actor:reload_weapon()
    end
    if mcm_check_ammo_after_reload then
        CreateTimeEvent("magazines_onekey_reload", "ammo_check_on_reload", 0.2, ammo_check_mcm.check_ammo_after_reload)
    end
    return true
end

-- Cache the available ammo types for each weapon
local cached_ammo_types = {}

function get_ammo_data(weapon, include_loaded)
    local id = weapon:id()
    local section = weapon:section()

    -- Find all magazines (or loose ammo for unsupported weapons) in player inventory
    local inventory_ammo = {}
    if magazines and magazine_binder.is_supported_weapon(weapon) then
        inventory_ammo = magazines.count_magazines(weapon)
    else
        local function itr(temp, itm)
            local section = itm:section()
            if IsItem("ammo",section) or IsItem("grenade_ammo",section) then
                inventory_ammo[section] = (inventory_ammo[section] or 0) + itm:ammo_get_count()
            end
        end
        db.actor:iterate_inventory(itr, nil)
    end

    if include_loaded then
        local loaded_section = get_loaded_ammo_section(weapon)
        if loaded_section then
            inventory_ammo[loaded_section] = (inventory_ammo[loaded_section] or 0) + 1
            -- printf("[OKR] Loaded ammo is section %s, count is now %s", loaded_section, inventory_ammo[loaded_section])
        end
    end

    -- Fetch usable ammo types from config
    local ammo_types
    if (not cached_ammo_types[id]) then
        ammo_types = utils_item.get_ammo(section, id)
        cached_ammo_types[id] = ammo_types
    else
        ammo_types = cached_ammo_types[id]
    end

    -- Check if ammo is in inventory (ignoring old ammo)
    num_avail = 0
    local has_ammo_types = {}
    for i = 1, #ammo_types do
        local section = ammo_types[i]
        local is_bad_ammo_type = section and string.find(section, "verybad") and true or false
        has_ammo_types[i] = (not is_bad_ammo_type) and (inventory_ammo[section] and (inventory_ammo[section] > 0)) and true or false
        if has_ammo_types[i] then
            num_avail = num_avail + 1
        end
    end

    return ammo_types, has_ammo_types, num_avail
end

-- Return top round type (one-based index) in the magazine (ignoring the chambered round)
function get_current_ammo_type(weapon, ammo_types)
    local current_ammo_type
    if magazines and magazine_binder.is_supported_weapon(weapon) then
        local mag_data = magazine_binder.get_mag_loaded(weapon:id())
        -- No mag or empty mag: use the round in chamber (or the last one fired if none chambered)
        if not mag_data or weapon:get_ammo_in_magazine() <= 1 then
            current_ammo_type = weapon:get_ammo_type() + 1 -- convert to one-based index
        else
            local top_round
            -- one-chambered, peek below to check top round in mag
            if magazines.retain_round(weapon) then
                top_round = mag_data.loaded[#mag_data.loaded - 1]
            -- open bolt: check top round in mag
            else
                top_round = stack.peek(mag_data.loaded)
            end
            current_ammo_type = weapon:get_ammo_type() + 1 -- fallback if the loop below did not find anything
            for i, sec in ipairs(ammo_types) do
                if top_round == sec then
                    current_ammo_type = i -- keep as one-based
                    break
                end
            end
        end
    else
        current_ammo_type = weapon:get_ammo_type() + 1
    end
    return current_ammo_type
end

-- Return top round section (string)
function get_loaded_ammo_section(weapon)
    local section
    if weapon:get_ammo_in_magazine() == 0 then
        section = nil
        return section
    end

    if magazines and magazine_binder.is_supported_weapon(weapon) then
        local mag_data = magazine_binder.get_mag_loaded(weapon:id())
        -- No mag or empty mag: use the round in chamber (or the last one fired if none chambered)
        if not mag_data or (weapon:get_ammo_in_magazine() <= 1 and magazines.retain_round(weapon)) then
            section = nil
        else
            local top_round
            -- one-chambered, peek below to check top round in mag
            if magazines.retain_round(weapon) then
                top_round = mag_data.loaded[#mag_data.loaded - 1]
            -- open bolt: check top round in mag
            else
                top_round = stack.peek(mag_data.loaded)
            end
            section = top_round
        end
    else
        local ammo_map = invert_table(magazines.get_ammo_map(weapon:id()))
        section = ammo_map[weapon:get_ammo_type() + 1]
    end
    return section
end