--[[
	=====================================================================
	Addon      : Tosox Mini Mods Repo: Auto Equip Gear
	Link       : https://www.moddb.com/mods/stalker-anomaly/addons/tosox-mini-mods-repo/
	Author     : Tosox
	Date       : 11.01.2024
	Last Edit  : 02.12.2024
	=====================================================================
--]]

function IsBinoc(o, c)
	if not (c) then
		c = o and o:clsid()
	end
	return c and (c == clsid.wpn_binocular_s or c == clsid.wpn_binocular)
end

function IsTorch(o, c)
	if not (c) then
		c = o and o:clsid()
	end
	return c and (c == clsid.device_torch_s or c == clsid.device_torch)
end

function IsPDA(o, c)
	if not (c) then
		c = o and o:clsid()
	end
	return c and (c == clsid.obj_pda_s or c == clsid.device_pda)
end

--=====================================================================

-- Flag to pause animations
pause_animations = false

-- General function to patch methods
function patch_function(target_table, function_name)
    if target_table and target_table[function_name] then
        local original_function = target_table[function_name]
        target_table[function_name] = function(...)
            if pause_animations then
                return
            end
            return original_function(...)
        end
    end
end

-- Patch specific functions
patch_function(actor_effects, "play_item_fx")
patch_function(outfit_animations, "on_item_to_slot")
patch_function(outfit_animations, "on_item_to_ruck")
patch_function(headgear_animations, "on_item_to_slot")
patch_function(headgear_animations, "on_item_to_ruck")
patch_function(headgear_animations, "on_item_drop")

--=====================================================================

function equip_gear()
	local function gather_gear(actor, itm)
		local sec = itm:section()
		if IsMelee(itm) then
			db.actor:move_to_slot(itm, 1)
		elseif IsGrenade(itm) then
			db.actor:move_to_slot(itm, 4)
		elseif IsBinoc(itm) then
			db.actor:move_to_slot(itm, 5)
		elseif IsBolt(itm) then
			db.actor:move_to_slot(itm, 6)
		elseif IsOutfit(itm) then
			db.actor:move_to_slot(itm, 7)
		elseif IsPDA(itm) then
			db.actor:move_to_slot(itm, 8)
		elseif IsTorch(itm) then
			db.actor:move_to_slot(itm, 10)
		elseif IsItem("device", sec) then
			db.actor:move_to_slot(itm, 9)
		elseif IsHeadgear(itm) then
			db.actor:move_to_slot(itm, 12)
		elseif IsItem("backpack", sec) then
			db.actor:move_to_slot(itm, 13)
		end
	end
	db.actor:iterate_inventory(gather_gear, db.actor)
end

function empty_slots()
	for i = 1, 13 do
		local itm = db.actor:item_in_slot(i)
		if itm then
			db.actor:move_to_ruck(itm)
		end
	end
end

--=====================================================================

equipped_slots = {}
equipped_belt = {}
function save_equipped_slots()
	-- Save all items in the slots
	for i = 1, 13 do
		local itm = db.actor:item_in_slot(i)
		if itm then
			equipped_slots[i] = itm:id()
		end
	end

	-- Save all items in the belt
	db.actor:iterate_belt(function(actor, itm)
		table.insert(equipped_belt, itm:id())
	end)
end

function restore_equipped_slots()
	-- Restore all items from the slots
	for i = 1, 13 do
		if equipped_slots[i] then
			local itm = level.object_by_id(equipped_slots[i])
			if itm then
				db.actor:move_to_slot(itm, i)
			end
		end
	end
	equipped_slots = {}

	-- Wait until we equipped our outfit
	CreateTimeEvent("tsx_auto_equip_gear", "equip_belt", 0.25, function()
		-- Restore all items from the belt
		for _, v in pairs(equipped_belt) do
			local itm = level.object_by_id(v)
			if itm then
				db.actor:move_to_belt(itm)
			end
		end
		equipped_belt = {}
		return true
	end)
end

--=====================================================================

orig_new_game_equipment = itms_manager.new_game_equippment
itms_manager.new_game_equippment = function(...)
	orig_new_game_equipment(...)

	-- Disable item animations
	pause_animations = true
	
	equip_gear()
	
	-- Wait until we equipped our outfit
	CreateTimeEvent("tsx_auto_equip_gear", "equip_gear", 0.25, function()
		-- Override autosave again
		exec_console_cmd("save " .. user_name() .. " - autosave")

		-- Restore item animations
		pause_animations = false
		return true
	end)
	
	return true
end

orig_xr_effects_bar_arena_teleport = xr_effects.bar_arena_teleport
xr_effects.bar_arena_teleport = function(...)
	-- Save the slots before entering the arena
	save_equipped_slots()

	-- Disable item animations
	pause_animations = true

    orig_xr_effects_bar_arena_teleport(...)

	-- Wait until we received all arena items
	CreateTimeEvent("tsx_auto_equip_gear", "equip_arena_enter", 0.25, function()
		equip_gear()
		return true
	end)
end

orig_xr_effects_bar_arena_teleport_2 = xr_effects.bar_arena_teleport_2
xr_effects.bar_arena_teleport_2 = function(...)
	orig_xr_effects_bar_arena_teleport_2(...)
	
	-- MBo - Arena rebalance patch
	if mbo_arena_monkeypatch then
		-- give actor back his items that were taken
		local box = get_story_object("bar_arena_inventory_box")
		if (box) then
			local function transfer_object_item(box, item)
				box:transfer_item(item, db.actor)
			end
			box:iterate_inventory_box(transfer_object_item, box)
		end
	end

	-- Wait until we received our items back
	CreateTimeEvent("tsx_auto_equip_gear", "equip_arena_exit", 0.25, function()
		-- Move the items that were equipped by default back into the backpack
		empty_slots()

		-- Wait until all items are back in the backpack
		CreateTimeEvent("tsx_auto_equip_gear", "restore_equipped_slots", 0.25, function()
			restore_equipped_slots()

			-- Restore item animations
			pause_animations = false
			
			return true
		end)
		return true
	end)
end

function save_state(m_data)
	m_data.tsx_auto_equip_gear = {
		["pause_animations"] = pause_animations,
		["equipped_slots"] = equipped_slots,
		["equipped_belt"] = equipped_belt
	}
end

function load_state(m_data)
	if m_data.tsx_auto_equip_gear then
		pause_animations = m_data.tsx_auto_equip_gear.pause_animations

		equipped_slots = m_data.tsx_auto_equip_gear.equipped_slots
		equipped_belt = m_data.tsx_auto_equip_gear.equipped_belt

		m_data.tsx_auto_equip_gear = nil
	end
end

function on_game_start()
	RegisterScriptCallback("save_state", save_state)
	RegisterScriptCallback("load_state", load_state)
end