local DRONES_CAP = 2

local classIndices_Internal = {
	[1] = "Scout",
	[3] = "Soldier",
	[7] = "Pyro",
	[4] = "Demoman",
	[6] = "Heavy",
	[9] = "Engineer",
	[5] = "Medic",
	[2] = "Sniper",
	[8] = "Spy",
}

local function findInTable(table, value)
	for i, v in pairs(table) do
		if v == value then
			return i
		end
	end
end

local function dictionaryLength(dictionary)
	local length = 0

	for i, _ in pairs(dictionary) do
		length = length + 1
	end

	return length
end

local callbacks = {}
local weaponsData = {}
local weaponTimers = {}

local CUSTOM_WEAPONS_INDICES = { "Parry", "Drone", "PHD", "SmolShield", "WingerDash", "Scavenger", "Diver", "bt" }
for _, weaponIndex in pairs(CUSTOM_WEAPONS_INDICES) do
	callbacks[weaponIndex] = {}
	weaponsData[weaponIndex] = {}
	weaponTimers[weaponIndex] = {}
end

function ClearCallbacks(index, activator, handle)
	handle = handle or activator:GetHandleIndex()

	local weaponCallbacks = callbacks[index][handle]

	if not weaponCallbacks then
		return
	end

	if activator and IsValid(activator) then
		for _, callbackId in pairs(weaponCallbacks) do
			activator:RemoveCallback(callbackId)
		end
	end

	callbacks[index][handle] = nil
end

function ClearData(index, activator, handle)
	handle = handle or activator:GetHandleIndex()

	local weaponData = weaponsData[index][handle]

	if not weaponData then
		return
	end

	weaponsData[index][handle] = nil
end

function ClearTimers(index, activator, handle)
	handle = handle or activator:GetHandleIndex()

	local weaponTimer = weaponTimers[index][handle]

	if not weaponTimer then
		return
	end

	for _, timerId in pairs(weaponTimer) do
		timer.Stop(timerId)
	end

	weaponTimers[index][handle] = nil
end

function NoMercyFailSafe(exEntName, activator, caller)
	local primary = activator:GetPlayerItemBySlot(0)

	if primary == nil then
		return
	end

	local exEnt = ents.FindByName(exEntName)

    if not exEnt then
        return
    end

	exEnt:Remove()
	caller:Remove()
	activator.m_flGravity = 0
end

function InstaLvLRefunded(_, activator)
	activator.InstantLevel = false
end
function InstaLvLPurchase(_, activator)
	activator.InstantLevel = true
end

function BuildingBuilt(_, building)
	if building.m_iHighestUpgradeLevel >= 2 then
		return
	end

	if building.m_bMiniBuilding ~= 0 then
		return
	end

	if building.m_bDisposableBuilding ~= 0 then
		return
	end

	local owner = building.m_hBuilder

	if not owner.InstantLevel then
		return
	end

	building.m_iHighestUpgradeLevel = 2
end

local cooldown_players = {}

local function _getPlayerCameraAngleUp(player)
	local pitch = player["m_angEyeAngles[0]"]

	return Vector(0, 0, -pitch)
end

-- drone
function DroneFired(sentryName, projectile)
	local owner = projectile.m_hOwnerEntity

	local sentryEnt = ents.FindByName(sentryName)

	if owner.m_iTeamNum == 3 then
		timer.Simple(0, function()
			sentryEnt.m_nSkin = 1
		end)

		sentryEnt.m_iTeamNum = 3
	end

	timer.Simple(0.6, function()
		sentryEnt.m_hBuilder = owner
	end)

	local firerateUpgrade = owner:GetPlayerItemBySlot(LOADOUT_POSITION_PDA):GetAttributeValue("engy sentry fire rate increased") or 1

	local primary = owner:GetPlayerItemBySlot(0)
	local mothershipUpgrade = primary:GetAttributeValue("throwable damage")

	if mothershipUpgrade then
		sentryEnt.m_flModelScale = 2.5
		timer.Simple(0.1, function()
			sentryEnt:SetHealth(250)
		end)
		sentryEnt["$attributeoverride"] = 1
		sentryEnt["$damagemult"] = 0.8
		sentryEnt["$fireratemult"] = firerateUpgrade + 1
		-- sentryEnt["$rangemult"] = 0.5
		sentryEnt["$bulletweapon"] = "Upgradeable TF_WEAPON_ROCKETLAUNCHER"
	end

	local ownerHandle = owner:GetHandleIndex()

	local dronesData = weaponsData.Drone[ownerHandle]

	local dronesList = dronesData.DronesList

	if dictionaryLength(dronesList) >= dronesData.DronesCap then
		-- remove first drone
		-- automatically cleared from array
		dronesList[1]:Remove()
	end

	table.insert(dronesList, projectile)

	projectile:AddCallback(ON_REMOVE, function()
		if sentryEnt and IsValid(sentryEnt) then
			util.ParticleEffect("ExplosionCore_buildings", sentryEnt:GetAbsOrigin(), Vector(0, 0, 0))
			sentryEnt:Remove()
		end

		table.remove(dronesList, findInTable(dronesList, projectile))

		local stationaryId = dronesData.DronesStationaryIds[projectile]

		if stationaryId then
			dronesData.DronesStationaryIds[projectile] = nil
			timer.Stop(stationaryId)
		end
	end)
end

function UpdateDroneCap(_, activator)
	local handle = activator:GetHandleIndex()

	local dronesData = weaponsData.Drone[handle]

	-- remove any existing drones
	for _, id in pairs(dronesData.DronesStationaryIds) do
		timer.Stop(id)
	end

	for _, projectile in pairs(dronesData.DronesList) do
		projectile:Remove()
	end

	local droneCount = DRONES_CAP

	local primary = activator:GetPlayerItemBySlot(0)
	local mothershipUpgrade = primary:GetAttributeValue("throwable damage")
	local droneCountUpgrade = primary:GetAttributeValue("throwable fire speed")

	if mothershipUpgrade then
		droneCount = 1
	elseif droneCountUpgrade then
		droneCount = droneCount + 1
	end

	local melee = activator:GetPlayerItemBySlot(LOADOUT_POSITION_MELEE)
	if droneCountUpgrade then
		melee:SetAttributeValue("mod wrench builds minisentry", 1)
	else
		melee:SetAttributeValue("mod wrench builds minisentry", nil)
	end

	dronesData.DronesCap = droneCount
end

local notificationsShown = {}

function DroneWalkerEquip(_, activator)
	-- fix weird quirk with template being spawned after you switch to a different class
	if classIndices_Internal[activator.m_iClass] ~= "Engineer" then
		return
	end

	print("drone walker equipped")
	local handle = activator:GetHandleIndex()

	if callbacks.Drone[handle] then
		DroneWalkerUnequip(activator, handle)
	end

	if not notificationsShown[handle] then
		notificationsShown[handle] = true
		activator["$DisplayTextCenter"](activator, "Press alt-fire to make drones stationary")
	end

	-- local meleeWeapon = activator:GetPlayerItemBySlot(2)
	-- local gunslingerEquipped = meleeWeapon.m_iClassname == "tf_weapon_robot_arm"

	-- if gunslingerEquipped then
	-- 	primary:SetAttributeValue("always crit", 1)
	-- 	-- primary:SetAttributeValue("engy sentry damage bonus", 1.25)
	-- else
	-- 	primary:SetAttributeValue("always crit", nil)
	-- 	-- primary:SetAttributeValue("engy sentry damage bonus", nil)
	-- end

	callbacks.Drone[handle] = {}
	weaponsData.Drone[handle] = {
		DronesList = {},

		DronesStationaryIds = {},

		-- DronesCap = not gunslingerEquipped and DRONES_CAP or DRONES_CAP + 1,
		DronesCap = DRONES_CAP,

		-- Buffed = gunslingerEquipped,
	}

	UpdateDroneCap(_, activator)

	local droneCallbacks = callbacks.Drone[handle]
	local dronesData = weaponsData.Drone[handle]

	-- on key press
	droneCallbacks.keyPress = activator:AddCallback(ON_KEY_PRESSED, function(_, key)
		if key ~= IN_ATTACK2 then
			return
		end

		if activator.m_hActiveWeapon.m_iClassname ~= "tf_weapon_shotgun_building_rescue" then
			return
		end

		if dictionaryLength(dronesData.DronesStationaryIds) > 0 then
			for projectile, id in pairs(dronesData.DronesStationaryIds) do
				timer.Stop(id)
				dronesData.DronesStationaryIds[projectile] = nil
			end

			return
		end

		for _, projectile in pairs(dronesData.DronesList) do
			local origin = projectile:GetAbsOrigin()

			projectile:SetLocalVelocity(Vector(0, 0, 0))

			dronesData.DronesStationaryIds[projectile] = timer.Create(0, function()
				projectile:SetAbsOrigin(origin)
			end, 0)
		end
	end)

	droneCallbacks.onRemoved = activator:AddCallback(ON_REMOVE, function()
		DroneWalkerUnequip(activator, handle)
	end)

	droneCallbacks.onDeath = activator:AddCallback(ON_DEATH, function()
		DroneWalkerUnequip(activator, handle)
	end)

	droneCallbacks.onSpawn = activator:AddCallback(ON_SPAWN, function()
		DroneWalkerUnequip(activator, handle)
	end)
end

function DroneWalkerUnequip(activator, handle)
	if not IsValid(activator) then
		activator = nil
	end

	local dronesData = weaponsData.Drone[handle]

	-- remove any existing drones
	for _, id in pairs(dronesData.DronesStationaryIds) do
		timer.Stop(id)
	end

	for _, projectile in pairs(dronesData.DronesList) do
		projectile:Remove()
	end

	ClearCallbacks("Drone", activator, handle)
	ClearData("Drone", activator, handle)
end