include('keysym.h.lua')

function widget:GetInfo()
	return {
		name      = "Graphic Unit Selector Redux",
		desc      = "Selects units when a user presses a certain key.",
		author    = "Shaman, mod by Helwor, NCG. Almost total rewrite ivand. Almost total rewrite esainane.",
		date      = "May, 2017",
		license   = "None",
		layer     = 15,
		enabled   = true,
	}
end

local Echo = Spring.Echo
local TableEcho = Spring.Utilities.TableEcho

local glBillboard = gl.Billboard
local glColor = gl.Color
local glDrawGroundCircle = gl.DrawGroundCircle
local glLineStipple = gl.LineStipple
local glLineWidth = gl.LineWidth
local glPushMatrix = gl.PushMatrix
local glPopMatrix = gl.PopMatrix
local glText = gl.Text
local glTranslate = gl.Translate

local GetModKeyState = Spring.GetModKeyState
local GetMouseState = Spring.GetMouseState
local GetMyTeamID = Spring.GetMyTeamID
local GetSpectatingState = Spring.GetSpectatingState
local GetUnitDefID = Spring.GetUnitDefID
local GetUnitPosition = Spring.GetUnitPosition
local GetUnitsInCylinder = Spring.GetUnitsInCylinder
local SelectUnitArray = Spring.SelectUnitArray
local IsReplay = Spring.IsReplay
local TraceScreenRay = Spring.TraceScreenRay
local ValidUnitID = Spring.ValidUnitID

------------------------------------------------------------------------------------------------------------
---  CONFIG
------------------------------------------------------------------------------------------------------------

local useIndividualHaloColors = true
local requireMouseCommit = true

local orderOfClass = {
	Constructor = 1,
	Raider = 2,
	Skirmisher = 3,
	Riot = 4,
	Assault = 5,
	Artillery = 6,
	["Weird Raider"] = 7, -- Also heavy raider
	["Anti Air"] = 8,
	Signature = 9,
	Special = 10,
	Utility = 11,
}

local keysByClass = {
	Constructor = KEYSYMS.Q,
	Raider = KEYSYMS.W,
	Skirmisher = KEYSYMS.E,
	Riot = KEYSYMS.R,
	Assault = KEYSYMS.T,
	Artillery = KEYSYMS.Y,
	["Weird Raider"] = KEYSYMS.S, -- Also heavy raider
	["Anti Air"] = KEYSYMS.D,
	Signature = KEYSYMS.F,
	Special = KEYSYMS.G,
	Utility = KEYSYMS.H
}

local classByKeys = {}
for k, v in pairs(keysByClass) do
	classByKeys[v] = k
end

local colorByClass = {
	Constructor = {0.8,0.68,0},
	Raider = {0,0.81,0.82},
	Skirmisher = {0.35,0.87,0.23},
	Riot = {0.86,0.28,0.34},
	Assault = {0.25,0.41,1},
	Artillery = {1,1,1},
	-- next row
	["Weird Raider"] = {0.34,0.89,0.94},
	["Anti Air"] = {0.34,0.59,.9},
	Signature = {1,.23,.8},
	Special = {1,0.58,0.25},
	Utility = {0.75,0.41,1},
	--growth = {0.19,0.5,0.08},
	--shrink = {1,0.08,0.58},
}
local colorByKey = {}
for k,v in pairs(colorByClass) do
	colorByKey[keysByClass[k]] = v
end

local keysAux={
	plus = KEYSYMS.KP6,
	minus = KEYSYMS.KP4,
}

local auxByKey = {}
for k, v in pairs(keysAux) do
	auxByKey[v] = k
end

local rad = 1000

-- almost entirely faithful reproductions of where unit icons going on the build panel.
-- Exceptions:
-- - Planes are all over the place;
--   - Raptor from scout to Anti-Air
--   - Phoenix from Riot to Scout/Heavy Raider
--   - Thunderbird from Special G to Riot
--   - Sparrow from Artillery to Special G
-- - Striders are all over the place, so go to where they'd feel most sensible
local unitTypes = {
	Constructor = {
		cloakcon = true,
		shieldcon = true,
		vehcon = true,
		hovercon = true,
		gunshipcon = true,
		planecon = true,
		spidercon = true,
		jumpcon = true,
		tankcon = true,
		amphcon = true,
		shipcon = true,
		athena = true, -- sure why not
		striderfunnelweb = true -- ditto
	},
	Raider = {
		cloakraid = true,
		shieldraid = true,
		vehraid = true,
		hoverraid = true,
		gunshipraid = true,
		planefighter = true,
		spiderscout = true,
		jumpraid = true,
		tankheavyraid = true,
		amphraid = true,
		shiptorpraider = true,
		striderantiheavy = true, -- sure, why not? ...there might be good reasons why not, but we'll see
	},
	Skirmisher = {
		cloakskirm = true,
		shieldskirm = true,
		vehsupport = true,
		hoverskirm = true,
		gunshipskirm = true,
		-- planes have no skirm
		spiderskirm = true,
		jumpskirm = true,
		-- tanks have no skirm
		amphfloater = true,
		shipskirm = true,
		striderbantha = true, -- lol
	},
	Riot = {
		cloakriot = true,
		shieldriot = true,
		vehriot = true,
		hoverriot = true,
		gunshipbomb = true,
		bomberdisarm = true, -- Thunderbird is better suited to the riot slot
		spiderriot = true,
		jumpblackhole = true,
		tankriot = true,
		amphriot = true,
		shipriot = true,
		striderdante = true, -- sure, why not
	},
	Assault = {
		cloakassault = true,
		shieldassault = true,
		vehassault = true,
		hoverassault = true,
		gunshipassault = true,
		bomberprec = true,
		spiderassault = true,
		jumpassault = true,
		tankassault = true,
		-- amph has no assault, it's in the signature slot; amphassault = true,
		shipassault = true,
		striderdetriment = true, -- sure, why not
	},
	Artillery = {
		cloakarty = true,
		shieldarty = true,
		veharty = true,
		hoverarty = true,
		gunshipheavyskirm = true,
		-- plane has no arty
		-- spider has no arty
		jumparty = true,
		tankarty = true,
		amphlaunch = true,
		shiparty = true,
		striderarty = true, -- sure, why not
	},
	["Weird Raider"] = {
		cloakheavyraid = true,
		shieldscout = true,
		vehscout = true,
		hoverheavyraid = true,
		gunshipemp = true,
		bomberriot = true, -- more of a heavy raider than a riot
		spideremp = true,
		jumpscout = true,
		tankraid = true,
		amphimpulse = true,
		shipscout = true,
		striderscorpion = true, -- sure, why not
	},
	["Anti Air"] = {
		cloakaa = true,
		shieldaa = true,
		vehaa = true,
		hoveraa = true,
		gunshipaa = true,
		planeheavyfighter = true,
		spideraa = true,
		jumpaa = true,
		tankaa = true,
		amphaa = true,
		shipaa = true,
		-- no strider anti air
	},
	Signature = {
		cloaksnipe = true,
		shieldfelon = true,
		vehheavyarty = true,
		-- no hover signature unit
		gunshipkrow = true,
		bomberheavy = true,
		spidercrabe = true,
		jumpsumo = true,
		tankheavyassault = true,
		amphassault = true,
		-- no ship signature unit
		shipheavyarty = true, -- sure, why not
	},
	Special = {
		cloakbomb = true,
		shieldbomb = true,
		vehcapture = true,
		hoverdepthcharge = true,
		gunshiptrans = true,
		planelightscout = true,
		spiderantiheavy = true,
		jumpbomb = true,
		-- tank has no special
		amphbomb = true,
		subraider = true,
		subtacmissile = true, -- sure, why not
	},
	Utility = {
		cloakjammer = true,
		shieldshield = true,
		-- rover has no utility
		-- hover has no utility
		gunshipheavytrans = true,
		planescout = true,
		spiderscout = true,
		-- jump has no utility
		tankheavyarty = true,
		amphtele = true,
		-- ship has no utility
		shipcarrier = true,
	},
}

local keyByUnitDefID = {}
for class,entries in pairs(unitTypes) do
	local key = keysByClass[class]
	for unitDefName,_ in pairs(entries) do
		local ud = UnitDefNames[unitDefName]
		if not ud then
			Echo('Could not find', unitDefName)
		end
		--Echo(ud, unitDefName, '->', key)
		keyByUnitDefID[ud.id] = key
	end
end

--[[
for unitDefID,unitDef in ipairs(UnitDefs) do
	if keyByUnitDefID[unitDefID] == nil and not string.match(UnitDefs[unitDefID].name, "dyn") and not string.match(UnitDefs[unitDefID].name, "comm_") and not string.match(UnitDefs[unitDefID].name, "chicken") and not string.match(UnitDefs[unitDefID].name, "pw_") then
    Echo('Warning, no type set for', unitDef.name)
  end
end
--]]

------------------------------------------------------------------------------------------------------------
---  END OF CONFIG
------------------------------------------------------------------------------------------------------------

function widget:Initialize()
	--Unload if in replay or if mod is not Zero-K
	if IsReplay() then
		widgetHandler:RemoveWidget(widget)
	end
end

function widget:GameOver(winningAllyTeams)
	--GameOver is irreversable with cheats, thus removing
	widgetHandler:RemoveWidget(widget)
end

local myTeamID = GetMyTeamID()
local unloaded = false
local function CheckIfSpectator()
	if GetSpectatingState() then
		myTeamID = nil
	else
		myTeamID = GetMyTeamID()
	end
	--spectator state is reversable with cheats
	-- Lets see what its like as a spectator
	--[[if GetSpectatingState() then
		widgetHandler:RemoveCallIn("KeyPress")
		widgetHandler:RemoveCallIn("KeyRelease")
		widgetHandler:RemoveCallIn("Update")
		widgetHandler:RemoveCallIn("DrawWorld")
		unloaded = true
	else
		if unloaded then
			widgetHandler:UpdateCallIn("KeyPress")
			widgetHandler:UpdateCallIn("KeyRelease")
			widgetHandler:UpdateCallIn("Update")
			widgetHandler:UpdateCallIn("DrawWorld")
			unloaded = false
		end
	end--]]
end

function widget:TeamChanged(teamID)
	CheckIfSpectator()
end

function widget:PlayerChanged(playerID)
	CheckIfSpectator()
end

function widget:PlayerAdded(playerID)
	CheckIfSpectator()
end

function widget:PlayerRemoved(playerID)
	CheckIfSpectator()
end

function widget:TeamDied(teamID)
	CheckIfSpectator()
end

function widget:TeamChanged(teamID)
	CheckIfSpectator()
end

local currentSelection = {}

function widget:SelectionChanged(selection)
	currentSelection = selection
end

local function CompareButtonOrder(l,r)
	return orderOfClass[l] < orderOfClass[r]
end

-- Key -> true | nil
local selectingKeys = {}
local activeSelectionKeys = 0

local selectionStrings
local selectionColors
local function activeKeysChanged()
	selectionStrings = {}
	selectionColors = {}
	for k,_ in pairs(selectingKeys) do
		selectionStrings[#selectionStrings + 1] = classByKeys[k]
		selectionColors[#selectionColors+1] = colorByKey[k]
	end
	table.sort(selectionStrings, CompareButtonOrder)
end

local isPregame = true
function widget:GameFrame(n)
	if n > 2 then
		isPregame = false
		widgetHandler:RemoveCallIn("GameFrame")
	end
end

-- {unitID} -> key | nil
local pendingSelection = {}
-- [unitID]
local pendingSelectionArray = {}

local acc = 0
function widget:KeyPress(key, mods, isRepeat)
	if isPregame then return end -- pregame effectively has your starting unit "selected"
	if selectingKeys[key] then return end -- No dupes
	if mods.alt then return end -- we want to be able to quickly select factories with no selection
	if mods.ctrl then return end -- also many global commands
	if auxByKey[key] then
		if key == keysAux.plus then
			rad = math.min(3000, rad+100)
		end

		if key == keysAux.minus then
			rad = math.max(100, rad-100)
		end
	elseif not classByKeys[key]then
		-- Not a key we're interested in
		return
	elseif activeSelectionKeys == 0 then
		if #currentSelection ~= 0 then
			-- If we aren't already selecting types, and we already have an existing selection, do nothing - we'd more likely want these keys to issue orders to them.
			return
		end
		-- New selection! How exciting.
		acc = 0
	end
	selectingKeys[key] = true
	activeSelectionKeys = activeSelectionKeys + 1
	activeKeysChanged()
	return true
end

local function DoCommitSelection()
	local _, _, _, shiftKey = GetModKeyState()
	if #pendingSelectionArray then
		SelectUnitArray(pendingSelectionArray, shiftKey)
	end
	table.clear(pendingSelection)
	table.clear(pendingSelectionArray)
end

function widget:KeyRelease(key) -- Called whenever user stops pressing a key.
	if not selectingKeys[key] then return end -- Not a key we're using
	selectingKeys[key] = nil
	activeSelectionKeys = activeSelectionKeys - 1
	activeKeysChanged()
	if activeSelectionKeys > 0 then return true end -- Still selecting more types with other keys
	-- We've released all type selection buttons. Commit the pending selection.
	if not requireMouseCommit then
		DoCommitSelection()
	end
	return true
end

local mouseAdding = false

function widget:MousePress(x,y,btn)
	if btn == 1 and activeSelectionKeys > 0 then
		-- Start adding
		mouseAdding = true
	elseif btn == 3 and mouseAdding then
		-- Cancel adding
		mouseAdding = false
		table.clear(pendingSelection)
		table.clear(pendingSelectionArray)
	else
		-- Don't interfere with dragging or other potential hotkeys
		return false
	end
	return true
end

function widget:MouseRelease(x,y,btn)
	if btn ~= 1 or not mouseAdding then return end
	mouseAdding = false
	DoCommitSelection()
end

local function GetCameraHeight()
	local cs = Spring.GetCameraState()
	local gy = Spring.GetGroundHeight(cs.px, cs.pz)
	local testHeight = cs.py - gy
	if cs.name == "ov" then
		testHeight = WG.AllUnitIcon.iconheight * 2
	elseif cs.name == "ta" then
		testHeight = cs.height - gy
	end
	return testHeight
end

local radAdj = 1
local x, y, z
local unitsPos = {}
function widget:Update(dt)
	if not (activeSelectionKeys or (requireMouseCommit and mouseAdding)) then return end
	acc = acc + dt
	local mouseX, mouseY = GetMouseState()
	local desc, mpos = TraceScreenRay(mouseX, mouseY, true)
	radAdj = GetCameraHeight() / 6000
	-- Is the mouse over the map?
	if desc ~= nil then
		-- Save its position
		x, y, z = unpack(mpos)

		-- If we're currently holding down left mouse, or if we don't require mouse clicks to add units, add all units in the area
		if activeSelectionKeys and (mouseAdding or not requireMouseCommit) then
			for _, uID in ipairs(GetUnitsInCylinder(x, z, rad * radAdj, myTeamID)) do
				-- Make sure not to add units multiple times
				if not pendingSelection[uID] then
					local key = keyByUnitDefID[GetUnitDefID(uID)]
					-- If this is the sort of unit we're looking for, add it
					if key and selectingKeys[key] then
						pendingSelection[uID] = key
					  pendingSelectionArray[#pendingSelectionArray+1] = uID
					end
				end
			end
		end
	else
		x = nil
	end
	-- We can continue to update unit halo information even when the mouse is off the world
	table.clear(unitsPos)
	for _,uID in ipairs(pendingSelectionArray) do
		if ValidUnitID(uID) then
			local pos = {GetUnitPosition(uID)}
			unitsPos[uID] = pos
		end
	end
end

local lineHeight = 80
local textSize = 66
local textX = 85

function widget:DrawWorld() -- this is used for openGL stuff.
	if (activeSelectionKeys == 0 and not (requireMouseCommit and mouseAdding)) or not x then return end
	local color = activeSelectionKeys == 0 and {.46,.41,.41} or selectionColors[1 + ((math.floor(acc * 5)) % #selectionColors)]
	glPushMatrix() --This is the start of an openGL function.
	glLineStipple(true)
	glLineWidth(2.0)
	glTranslate(x, y, z)
	glColor(color[1], color[2], color[3], 1)
	glDrawGroundCircle(0, 0, 0, rad * radAdj, 40) -- draws a simple circle.
	glColor(1,1,1,1)
	glBillboard()
	local description = "Will select:"
	if mouseAdding or not requireMouseCommit then
		description = "Selecting:"
	end
	local textY = (#selectionStrings * lineHeight * .5 + 25)
	glText(description, textX, textY, textSize, "v") -- Displays text. First value is the string, second is a modifier for x (in this case it's x-25), third is a modifier for y, fourth is the size, then last is a modifier for the text itself. "v" means vertical align.
	for i=1,#selectionStrings do
		local text = selectionStrings[i]
		color = colorByClass[text]
		glColor(color[1], color[2], color[3], 1)
		textY = textY - lineHeight
		glText(text, textX, textY, textSize, "v")
	end
	glColor(1, 1, 1, 1) -- we have to reset what we did here.
	glLineWidth(1.0)
	glLineStipple(false)
	glPopMatrix() -- end of function. Have to use this with after a push!

	glPushMatrix()
	glLineWidth(3.0)
	for uID, key in pairs(pendingSelection) do
		local pos = unitsPos[uID]
		if pos ~= nil then
			if useIndividualHaloColors then
				color = colorByKey[key]
			end
			glColor(color[1], color[2], color[3], 1)
			local px,py,pz = unpack(pos)
			glDrawGroundCircle(px, py, pz, 80, 80)
		end
	end
	glColor(1, 1, 1, 1)
	glPopMatrix()
end