-- Color Coded Healthbars (Working Final)
-- Description: Show color-coded healthbars based on enemy breed type
-- Author: Custom Mod
-- Requires: Original Healthbars mod by raindish
local mod = get_mod("ColorCodedHealthbars")
local Breeds = require("scripts/settings/breed/breeds")
local HealthExtension = require("scripts/extension_systems/health/health_extension")

-- Color definitions for different breed types
mod.breed_colors = {
	horde = { 150, 150, 150 },      -- Gray for horde enemies
	roamer = { 180, 180, 180 },     -- Light gray for roamers
	elite = { 255, 165, 0 },        -- Orange for elites
	special = { 255, 0, 255 },      -- Magenta for specials
	monster = { 255, 0, 0 },        -- Red for monsters
	captain = { 128, 0, 128 },      -- Purple for captains
	default = { 255, 255, 255 }     -- White for unknown types
}

-- Function to update colors from settings
local function update_colors_from_settings()
	mod.breed_colors.horde = { 
		mod:get("horde_color_r") or 150, 
		mod:get("horde_color_g") or 150, 
		mod:get("horde_color_b") or 150 
	}
	mod.breed_colors.elite = { 
		mod:get("elite_color_r") or 255, 
		mod:get("elite_color_g") or 165, 
		mod:get("elite_color_b") or 0 
	}
	mod.breed_colors.special = { 
		mod:get("special_color_r") or 255, 
		mod:get("special_color_g") or 0, 
		mod:get("special_color_b") or 255 
	}
	mod.breed_colors.monster = { 
		mod:get("monster_color_r") or 255, 
		mod:get("monster_color_g") or 0, 
		mod:get("monster_color_b") or 0 
	}
	mod.breed_colors.roamer = { 
		math.min(255, (mod:get("horde_color_r") or 150) + 30), 
		math.min(255, (mod:get("horde_color_g") or 150) + 30), 
		math.min(255, (mod:get("horde_color_b") or 150) + 30) 
	}
	mod.breed_colors.captain = { 
		math.floor((mod:get("monster_color_r") or 255) / 2), 
		0, 
		math.floor((mod:get("monster_color_r") or 255) / 2) 
	}
	-- Add subcategory colors
	mod.breed_colors.elite_ranged = {
		mod:get("elite_ranged_color_r") or 255,
		mod:get("elite_ranged_color_g") or 100,
		mod:get("elite_ranged_color_b") or 0
	}
	mod.breed_colors.elite_melee = {
		mod:get("elite_melee_color_r") or 255,
		mod:get("elite_melee_color_g") or 165,
		mod:get("elite_melee_color_b") or 0
	}
	mod.breed_colors.special_sniper = {
		mod:get("special_sniper_color_r") or 255,
		mod:get("special_sniper_color_g") or 0,
		mod:get("special_sniper_color_b") or 200
	}
	mod.breed_colors.special_disabler = {
		mod:get("special_disabler_color_r") or 200,
		mod:get("special_disabler_color_g") or 0,
		mod:get("special_disabler_color_b") or 255
	}
end

local show = {}

local function get_toggles()
	-- Get general breed type toggles only - individual breeds handled by original mod
	show.show_horde = mod:get("show_horde")
	show.show_roamer = mod:get("show_roamer")
	show.show_elite = mod:get("show_elite")
	show.show_special = mod:get("show_special")
	show.show_monster = mod:get("show_monster")
	show.show_captain = mod:get("show_captain")
	
	-- Update colors from settings
	update_colors_from_settings()
end

get_toggles()

mod.on_setting_changed = function()
	get_toggles()
end

-- Function to get breed color with subcategory support
function mod.get_breed_color(unit)
	if not unit then
		return mod.breed_colors.default
	end
	
	local unit_data_extension = ScriptUnit.has_extension(unit, "unit_data_system")
	if not unit_data_extension then
		return mod.breed_colors.default
	end
	
	local breed = unit_data_extension:breed()
	if not breed or not breed.tags then
		return mod.breed_colors.default
	end
	
	local tags = breed.tags
	local breed_name = breed.name
	
	-- Check for specific breed subcategories first
	if tags.elite then
		-- Elite subcategories - Ranged vs Melee
		if breed_name == "renegade_gunner" or breed_name == "renegade_shocktrooper" or 
		   breed_name == "cultist_gunner" or breed_name == "chaos_ogryn_gunner" or
		   breed_name == "cultist_shocktrooper" then
			return mod.breed_colors.elite_ranged
		else
			return mod.breed_colors.elite_melee
		end
	elseif tags.special then
		-- Special subcategories - Snipers vs Disablers
		if breed_name == "renegade_sniper" then
			return mod.breed_colors.special_sniper
		elseif breed_name == "renegade_netgunner" or breed_name == "chaos_hound" then
			return mod.breed_colors.special_disabler
		else
			return mod.breed_colors.special
		end
	elseif tags.monster then
		return mod.breed_colors.monster
	elseif tags.captain then
		return mod.breed_colors.captain
	elseif tags.horde then
		return mod.breed_colors.horde
	elseif tags.roamer then
		return mod.breed_colors.roamer
	else
		return mod.breed_colors.default
	end
end

-- Function to determine if unit should have healthbar (simplified - rely on original mod for individual breeds)
local function should_enable_healthbar(unit)
	if not unit then
		return false
	end

	local unit_data_extension = ScriptUnit.has_extension(unit, "unit_data_system")
	if not unit_data_extension then
		return false
	end
	
	local breed = unit_data_extension:breed()
	if not breed or not breed.tags then
		return false
	end
	
	-- Only check breed type category toggles - let original Healthbars mod handle individual breeds
	local tags = breed.tags
	if tags.monster and show.show_monster then
		return true
	elseif tags.captain and show.show_captain then
		return true
	elseif tags.elite and show.show_elite then
		return true
	elseif tags.special and show.show_special then
		return true
	elseif tags.horde and show.show_horde then
		return true
	elseif tags.roamer and show.show_roamer then
		return true
	end

	return false
end

-- Store units that should have colored healthbars
mod.tracked_units = {}

-- Hook HealthExtension to add our own healthbars AND track for coloring
mod:hook_safe("HealthExtension", "init", function(_self, _extension_init_context, unit, _extension_init_data, _game_object_data)
	if should_enable_healthbar(unit) then
		mod.tracked_units[unit] = mod.get_breed_color(unit)
		-- Add healthbar marker (this will work if original Healthbars mod is installed)
		Managers.event:trigger("add_world_marker_unit", "custom_healthbar", unit)
	end
end)

mod:hook_safe("HuskHealthExtension", "init", function(self, _extension_init_context, unit, _extension_init_data, _game_session, _game_object_id, _owner_id)
	-- Make sure husks have the methods needed
	self.set_last_damaging_unit = HealthExtension.set_last_damaging_unit
	self.last_damaging_unit = HealthExtension.last_damaging_unit
	self.last_hit_zone_name = HealthExtension.last_hit_zone_name
	self.last_hit_was_critical = HealthExtension.last_hit_was_critical
	self.was_hit_by_critical_hit_this_render_frame = HealthExtension.was_hit_by_critical_hit_this_render_frame

	if should_enable_healthbar(unit) then
		mod.tracked_units[unit] = mod.get_breed_color(unit)
		Managers.event:trigger("add_world_marker_unit", "custom_healthbar", unit)
	end
end)

-- Hook into the original Healthbars mod's marker template to apply colors
mod:hook_safe("HudElementWorldMarkers", "draw", function(self, dt, t, ui_renderer, render_settings, input_service)
	-- Check if we can access the custom_healthbar template and hook it
	if self._marker_templates and self._marker_templates["custom_healthbar"] then
		local template = self._marker_templates["custom_healthbar"]
		
		-- Hook the template's update function if we haven't already
		if template.update_function and not template._color_mod_hooked then
			template._color_mod_hooked = true
			local original_update = template.update_function
			
			template.update_function = function(parent, ui_renderer, widget, marker, template, dt, t)
				-- Call the original update function
				local result = original_update(parent, ui_renderer, widget, marker, template, dt, t)
				
				-- Apply our breed-specific color
				if marker and marker.unit and mod.tracked_units[marker.unit] then
					local color = mod.tracked_units[marker.unit]
					if widget and widget.style and widget.style.bar then
						widget.style.bar.color[2] = color[1]  -- R
						widget.style.bar.color[3] = color[2]  -- G
						widget.style.bar.color[4] = color[3]  -- B
					end
				end
				
				return result
			end
			
			if mod:get("show_startup_messages") then
				mod:echo("Successfully hooked into Healthbars mod template for color application!")
			end
		end
	end
end)

-- Also apply colors during the world markers update
mod:hook_safe("HudElementWorldMarkers", "update", function(self, dt, t, ui_renderer, render_settings, input_service)
	if self._active_markers then
		for marker_id, marker in pairs(self._active_markers) do
			if marker and marker.template and marker.template.name == "custom_healthbar" then
				if marker.unit and mod.tracked_units[marker.unit] and marker.widget then
					local color = mod.tracked_units[marker.unit]
					local widget = marker.widget
					
					-- Apply color to the bar style
					if widget.style and widget.style.bar then
						widget.style.bar.color[2] = color[1]  -- R
						widget.style.bar.color[3] = color[2]  -- G
						widget.style.bar.color[4] = color[3]  -- B
					end
				end
			end
		end
	end
end)

-- Debug commands
mod:command("cc_test", "Test color coded healthbars", function()
	mod:echo("Color Coded Healthbars (Working Final) loaded!")
	
	-- Check if original Healthbars mod is available
	local original_mod = get_mod("Healthbars")
	if original_mod then
		mod:echo("✓ Original Healthbars mod detected")
	else
		mod:echo("⚠ Original Healthbars mod NOT detected - healthbars may not appear")
		mod:echo("  Install the 'Healthbars' mod by raindish for best results")
	end
	
	mod:echo("Tracking " .. table.size(mod.tracked_units) .. " units")
	
	-- Display current settings
	mod:echo("Current breed type toggles:")
	for setting, value in pairs(show) do
		if string.find(setting, "show_") then
			mod:echo("  " .. setting .. ": " .. tostring(value))
		end
	end
	
	mod:echo("Current colors:")
	for breed_type, color in pairs(mod.breed_colors) do
		mod:echo("  " .. breed_type .. ": RGB(" .. color[1] .. ", " .. color[2] .. ", " .. color[3] .. ")")
	end
end)

mod:command("cc_list", "List tracked units", function()
	mod:echo("Currently tracked units:")
	local count = 0
	for unit, color in pairs(mod.tracked_units) do
		if count < 10 and unit and HEALTH_ALIVE[unit] then
			local unit_data_extension = ScriptUnit.has_extension(unit, "unit_data_system")
			if unit_data_extension then
				local breed = unit_data_extension:breed()
				if breed then
					mod:echo("  " .. breed.name .. " - RGB(" .. color[1] .. ", " .. color[2] .. ", " .. color[3] .. ")")
					count = count + 1
				end
			end
		elseif count >= 10 then
			mod:echo("  ... and more")
			break
		end
	end
	if count == 0 then
		mod:echo("  No units currently tracked")
	end
end)

mod:command("cc_status", "Check integration status", function()
	local original_mod = get_mod("Healthbars")
	local hud_manager = Managers.ui and Managers.ui:active_hud()
	local world_markers = hud_manager and hud_manager._elements and hud_manager._elements.HudElementWorldMarkers
	
	mod:echo("=== Integration Status ===")
	mod:echo("Original Healthbars mod: " .. (original_mod and "✓ Found" or "✗ Missing"))
	mod:echo("HudElementWorldMarkers: " .. (world_markers and "✓ Found" or "✗ Missing"))
	
	if world_markers and world_markers._marker_templates then
		local has_template = world_markers._marker_templates["custom_healthbar"] ~= nil
		mod:echo("custom_healthbar template: " .. (has_template and "✓ Found" or "✗ Missing"))
		
		if has_template then
			local template = world_markers._marker_templates["custom_healthbar"]
			local is_hooked = template._color_mod_hooked == true
			mod:echo("Template color hook: " .. (is_hooked and "✓ Applied" or "✗ Not applied"))
		end
	end
	
	mod:echo("Tracked units: " .. table.size(mod.tracked_units))
end)

-- Startup message with toggle
if mod:get("show_startup_messages") then
	mod:echo("Color Coded Healthbars loaded!")
	mod:echo("⚠ Requires the original 'Healthbars' mod to display healthbars")
	mod:echo("Type /cc_status to check integration status")
end