---@diagnostic disable: inject-field, undefined-field
local WAVEBARGEN_DEBUG_MODE = true

local DEFAULT_CLASS_ICONS = {
	Scout = "scout",
	Soldier = "soldier",
	Pyro = "pyro",
	Demoman = "demo",
	Heavyweapons = "heavy",
	Engineer = "engineer",
	Medic = "medic",
	Sniper = "sniper",
	Spy = "spy",
}

local MISSION_FLAG_INTERVAL = 0.2

local DEFAULT_TANK_ICON = "tank"

local objectiveResource
local currentWaveBarEntries = {}

local activeFlashingTimer
local flashingIcons = {}

local activeExtraIcons = {} -- used for sentry buster & teleporter icon

local function makeKey(icon, flags)
	return icon .. "|" .. tostring(flags)
end

local function getClassFlag(template, subwaveData)
	local flag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NORMAL

	if template.IsTank then
		if template.Miniboss == false then
			flag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NORMAL
		else
			flag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NORMAL | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_MINIBOSS
		end

		if subwaveData.SupportInfinite then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT
		end

		if subwaveData.SupportLimited then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT_LIMITED
		end

		return flag
	end

	for _, attribute in pairs(template.Attributes or {}) do
		if attribute == BOT_ATTRIBUTES.MINIBOSS and not subwaveData.IsMission then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_MINIBOSS
		end

		if attribute == BOT_ATTRIBUTES.ALWAYS_CRIT then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_ALWAYSCRIT
		end
	end

	if subwaveData.IsMission then
		-- (royal): MVM_CLASS_FLAG_MISSION is not used on purpose
		flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT
	else
		if subwaveData.SupportInfinite then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT
		end

		if subwaveData.SupportLimited then
			flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT
			-- flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT_LIMITED
		end
	end

	return flag
end

function FindIconIndexFromIconName(classIconName, matchingFlagsAtLeastOne)
	if not objectiveResource then
		objectiveResource = ents.FindByClass("tf_objective_resource")
	end

	if not objectiveResource then
		util.PrintToChatAll("[WAVE-BAR-GEN]: NO tf_objective_resource FOUND")
		return nil
	end

	for i = 1, #objectiveResource.m_iszMannVsMachineWaveClassNames do
		local classIconAtIndex = objectiveResource.m_iszMannVsMachineWaveClassNames[i]
		local flagAtIndex = objectiveResource.m_nMannVsMachineWaveClassFlags[i]

		if classIconAtIndex == classIconName then
			local isMatching = true

			if matchingFlagsAtLeastOne then
				isMatching = false

				for _, flagToMatch in pairs(matchingFlagsAtLeastOne) do
					if flagToMatch == MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NONE then
						if flagAtIndex == flagToMatch then
							isMatching = true
							break
						end
					else
						if (flagAtIndex & flagToMatch) == flagToMatch then
							isMatching = true
							break
						end
					end
				end
			end

			if isMatching then
				return i
			end
		end
	end

	return nil
end

function AddToFlashingIcon(classIcon, flag, indexOverride, templateData)
	if flashingIcons[classIcon] then
		flashingIcons[classIcon].ActiveCount = flashingIcons[classIcon].ActiveCount + 1
		return
	end

	StartFlashingIcon(classIcon, flag, indexOverride, templateData)
end

function RemoveFromFlashingIcon(classIcon)
	if not flashingIcons[classIcon] then
		return
	end

	flashingIcons[classIcon].ActiveCount = flashingIcons[classIcon].ActiveCount - 1

	if flashingIcons[classIcon].ActiveCount > 0 then
		return
	end

	StopFlashingIcon(classIcon)
end

function StartFlashingIcon(classIcon, flag, indexOverride, templateData)
	if flashingIcons[classIcon] then
		return
	end

	local index = indexOverride or FindIconIndexFromIconName(classIcon)

	if not index then
		util.PrintToChatAll(
			"[WAVE-BAR-GEN][ERROR]: StartFlashingIcon failed, unable to find icon from " .. tostring(classIcon)
		)
		return
	end

	local isCrit = templateData.Attributes and TableContains(templateData.Attributes, BOT_ATTRIBUTES.ALWAYS_CRIT)

	local iconInfo = {
		OriginalFlag = flag,
		ActiveCount = 1,
		Index = index,
		IsCrit = isCrit,
	}
	flashingIcons[classIcon] = iconInfo

	if activeFlashingTimer then
		return
	end

	local isRed = false
	activeFlashingTimer = timer.Create(MISSION_FLAG_INTERVAL, function()
		isRed = not isRed

		for otherIcon in pairs(flashingIcons) do
			local newFlag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT

			if isRed then
				newFlag = newFlag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_MINIBOSS
			end

			local otherIconInfo = flashingIcons[otherIcon]

			if otherIconInfo.IsCrit then
				newFlag = newFlag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_ALWAYSCRIT
			end

			objectiveResource.m_nMannVsMachineWaveClassFlags[otherIconInfo.Index] = newFlag
		end
	end, 0)
end

function StopFlashingIcon(classIcon)
	if flashingIcons[classIcon] then
		objectiveResource.m_nMannVsMachineWaveClassFlags[flashingIcons[classIcon].Index] =
			flashingIcons[classIcon].OriginalFlag

		flashingIcons[classIcon] = nil
	end

	if not next(flashingIcons) and activeFlashingTimer then
		timer.Stop(activeFlashingTimer)
		activeFlashingTimer = nil
	end
end

function HideIcon(classIcon, flagRequired)
	if not objectiveResource then
		objectiveResource = ents.FindByClass("tf_objective_resource")
	end

	if not objectiveResource then
		util.PrintToChatAll("[WAVE-BAR-GEN]: NO tf_objective_resource FOUND")
		return
	end

	local foundIndex
	for i = 1, #objectiveResource.m_iszMannVsMachineWaveClassNames do
		local iconAtIndex = objectiveResource.m_iszMannVsMachineWaveClassNames[i]
		local flagAtIndex = objectiveResource.m_nMannVsMachineWaveClassFlags[i]

		if iconAtIndex == classIcon then
			if flagRequired then
				if (flagAtIndex & flagRequired) == flagRequired then
					foundIndex = i
					break
				end
			else
				foundIndex = i
				break
			end
		end
	end

	if not foundIndex then
		print("HideIcon failed due to missing foundIndex")
		util.PrintToChatAll("[WAVE-BAR-GEN]: HideIcon FAILED")
		return
	end

	objectiveResource.m_nMannVsMachineWaveClassFlags[foundIndex] = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NONE
end

function SentryBusterSpawned(templateData)
	AddToExtraIcon(templateData, true)
end

function SentryBusterDied(templateData)
	RemoveFromExtraIcon(templateData)
end

function AddToExtraIcon(templateData, shouldFlash)
	local classIcon = templateData.ClassIcon

	if activeExtraIcons[classIcon] then
		activeExtraIcons[classIcon].ActiveCount = activeExtraIcons[classIcon].ActiveCount + 1
		return
	end

	local flag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT
	if templateData.Attributes and TableContains(templateData.Attributes, BOT_ATTRIBUTES.ALWAYS_CRIT) then
		flag = flag | MVM_CLASS_FLAGS.MVM_CLASS_FLAG_ALWAYSCRIT
	end

	for i = 1, #objectiveResource.m_iszMannVsMachineWaveClassNames do
		local iconAtIndex = objectiveResource.m_iszMannVsMachineWaveClassNames[i]
		if not iconAtIndex or iconAtIndex == "" then
			objectiveResource.m_iszMannVsMachineWaveClassNames[i] = classIcon
			objectiveResource.m_nMannVsMachineWaveClassCounts[i] = 1
			objectiveResource.m_nMannVsMachineWaveClassFlags[i] = flag

			if shouldFlash then
				StartFlashingIcon(templateData.ClassIcon, flag, i, templateData)
			end

			activeExtraIcons[classIcon] = {
				Index = i,
				ActiveCount = 1,
				Flashing = shouldFlash,
			}

			break
		end
	end
end

function RemoveFromExtraIcon(templateData)
	local classIcon = templateData.ClassIcon

	local activeIconData = activeExtraIcons[classIcon]
	if not activeIconData then
		return
	end

	activeIconData.ActiveCount = activeIconData.ActiveCount - 1

	if activeIconData.ActiveCount <= 0 then
		if activeIconData.Flashing then
			StopFlashingIcon(classIcon)
		end

		local index = activeIconData.Index
		local iconAtIndex = objectiveResource.m_iszMannVsMachineWaveClassNames[index]

		if iconAtIndex == classIcon then
			objectiveResource.m_iszMannVsMachineWaveClassNames[index] = ""
			objectiveResource.m_nMannVsMachineWaveClassCounts[index] = 0
			objectiveResource.m_nMannVsMachineWaveClassFlags[index] = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NONE

			activeExtraIcons[classIcon] = nil
		end
	end
end

function MissionBotSpawned(templateData, subwaveData)
	print("mission bot spawned with icon ", templateData.ClassIcon)

	local classFlag

	if subwaveData.MissionNoHideIcon then
		classFlag = getClassFlag(templateData, subwaveData)
	else
		classFlag = MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NONE
	end

	AddToFlashingIcon(
		templateData.ClassIcon,
		classFlag,
		FindIconIndexFromIconName(
			templateData.ClassIcon,
			{ MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT, MVM_CLASS_FLAGS.MVM_CLASS_FLAG_NONE }
		),
		templateData
	)
end

function MissionBotDied(templateData)
	print("mission bot died with icon ", templateData.ClassIcon)
	RemoveFromFlashingIcon(templateData.ClassIcon)
end

function WaveBarGen(waveData, inWave)
	if not objectiveResource then
		objectiveResource = ents.FindByClass("tf_objective_resource")
	end

	if not objectiveResource then
		util.PrintToChatAll("[WAVE-BAR-GEN]: NO tf_objective_resource FOUND")
		return
	end

	-- clear existing wave bar
	for i = 1, #objectiveResource.m_nMannVsMachineWaveClassCounts do
		objectiveResource.m_nMannVsMachineWaveClassCounts[i] = 0
	end

	if WAVEBARGEN_DEBUG_MODE then
		util.PrintToChatAll("[WAVE-BAR-GEN]: CLEARED EXISTING WAVE BAR")
	end

	if WAVEBARGEN_DEBUG_MODE then
		util.PrintToChatAll("[WAVE-BAR-GEN]: GENERATING WAVE BAR")
	end

	local templateTotals = {}
	local preAllocatedTemplates = {}

	local function templateIndexKey(subwaveIndex, templateIndex)
		return tostring(subwaveIndex) .. "_" .. tostring(templateIndex)
	end

	local function preAllocatedIndexKey(subwaveIndex, classIcon)
		return tostring(subwaveIndex) .. "_" .. classIcon
	end

	local function applyDefaultClassIcon(templateData)
		if not templateData.ClassIcon then
			if templateData.IsTank then
				templateData.ClassIcon = DEFAULT_TANK_ICON
			else
				templateData.ClassIcon = DEFAULT_CLASS_ICONS[templateData.Class]
			end
		end
	end

	for subwaveIndex, subwaveData in ipairs(waveData) do
		local templateCount = #subwaveData.Templates

		for templateIndex, templateData in ipairs(subwaveData.Templates) do
			applyDefaultClassIcon(templateData)

			local botCount = 1

			if not subwaveData.SupportInfinite and not subwaveData.SupportLimited then
				botCount = math.floor(subwaveData.TotalCount / templateCount)

				if templateIndex <= (subwaveData.TotalCount % templateCount) then
					botCount = botCount + 1
				end
			end

			templateTotals[templateIndexKey(subwaveIndex, templateIndex)] = botCount
		end

		if subwaveData.PreAllocatedTemplates then
			for _, templateData in ipairs(subwaveData.PreAllocatedTemplates) do
				applyDefaultClassIcon(templateData)

				local indexKey = preAllocatedIndexKey(subwaveIndex, templateData.ClassIcon)

				if not preAllocatedTemplates[indexKey] then
					preAllocatedTemplates[indexKey] = {
						Template = templateData,
						Count = 0,
					}
				end

				preAllocatedTemplates[indexKey].Count = preAllocatedTemplates[indexKey].Count + 1
			end
		end
	end

	local curIndex = 0
	local merged = {}
	local orderedKeys = {}

	local function handleIcon(subwaveData, icon, count, flags)
		local canMerge = not subwaveData.DoNotMergeIcon

		local keyData = {
			icon = icon,
			count = count,
			flags = flags,
			support = subwaveData.SupportInfinite or subwaveData.SupportLimited,
			mission = subwaveData.IsMission,
		}

		local shouldHideIcon = subwaveData.IsMission

		if subwaveData.MissionNoHideIcon then
			shouldHideIcon = false
		end

		keyData.shouldHideInWave = shouldHideIcon

		local key = makeKey(icon, flags)
		if canMerge then
			if not merged[key] then
				merged[key] = keyData
				table.insert(orderedKeys, key)
			else
				merged[key].count = merged[key].count + count
			end
		else
			merged[key] = keyData
			table.insert(orderedKeys, key)
		end
	end

	for subwaveIndex, subwaveData in ipairs(waveData) do
		for templateIndex, template in ipairs(subwaveData.Templates) do
			curIndex = curIndex + 1

			local icon = template.ClassIcon

			local flags = getClassFlag(template, subwaveData)

			local count
			if not subwaveData.SupportInfinite and not subwaveData.SupportLimited then
				count = templateTotals[templateIndexKey(subwaveIndex, templateIndex)] or 1
			else
				count = 1
			end

			handleIcon(subwaveData, icon, count, flags)
		end

		if subwaveData.PreAllocatedTemplates then
			for _, preAllocatedTemplateData in pairs(preAllocatedTemplates) do
				local template = preAllocatedTemplateData.Template

				local icon = template.ClassIcon

				local flags = getClassFlag(template, subwaveData)

				local count
				if not subwaveData.SupportInfinite and not subwaveData.SupportLimited then
					count = preAllocatedTemplateData.Count
				else
					count = 1
				end

				handleIcon(subwaveData, icon, count, flags)
			end
		end
	end

	local totalRobots = 0

	for i, key in ipairs(orderedKeys) do
		local entry = merged[key]

		objectiveResource.m_iszMannVsMachineWaveClassNames[i] = entry.icon
		objectiveResource.m_nMannVsMachineWaveClassCounts[i] = entry.count
		objectiveResource.m_nMannVsMachineWaveClassFlags[i] = entry.flags

		if not entry.support then
			totalRobots = totalRobots + entry.count
		end
	end

	currentWaveBarEntries = merged

	objectiveResource.m_nMannVsMachineWaveEnemyCount = totalRobots

	if WAVEBARGEN_DEBUG_MODE then
		util.PrintToChatAll("total robots: " .. tostring(totalRobots))
	end

	if inWave then
		for _, key in pairs(currentWaveBarEntries) do
			if key.shouldHideInWave then
				HideIcon(key.icon, MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT)
			end
		end
	end
end

local function waveInit(_wave)
	if flashingIcons then
		for icon in pairs(flashingIcons) do
			StopFlashingIcon(icon)
		end
	end

	activeFlashingTimer = nil
	flashingIcons = {}

	activeExtraIcons = {}
end

local function waveStart(_wave)
	-- for _, key in pairs(currentWaveBarEntries) do
	-- 	if key.mission then
	-- 		HideIcon(key.icon, MVM_CLASS_FLAGS.MVM_CLASS_FLAG_SUPPORT)
	-- 	end
	-- end
end

if OnWaveInitFuncs then
	table.insert(OnWaveInitFuncs, waveInit)
else
	function OnWaveInit(wave)
		waveInit(wave)
	end
end

if OnWaveStartFuncs then
	table.insert(OnWaveStartFuncs, waveStart)
else
	function OnWaveStart(wave)
		waveStart(wave)
	end
end