local UIUtil = import('/lua/ui/uiutil.lua')
local LayoutHelpers = import('/lua/maui/layouthelpers.lua')
local Bitmap = import('/lua/maui/bitmap.lua').Bitmap
local Group = import('/lua/maui/group.lua').Group
local Checkbox = import('/lua/maui/checkbox.lua').Checkbox
local Button = import('/lua/maui/button.lua').Button
local borders = import('/lua/ui/game/borders.lua').GetMapGroup()
local ToolTip = import('/lua/ui/game/tooltip.lua')
local EffectHelpers = import('/lua/maui/effecthelpers.lua')
local GPGSelection = import('/lua/ui/game/selection.lua')
local selectionSets = GPGSelection.selectionSets

-- Get options
local Prefs = import('/lua/user/prefs.lua')
local options = Prefs.GetFromCurrentProfile('options')

MaxSquad = 10
DefaultMaxSquad = options.gui_max_autogroups
SquadCycle = 0

CurrBattalion = 0
CurrSquad = false

autoselectinverse = false

local panel = false
local mouseover = false

-------------------------------------------------------------------------------
-- Functions to set global variables ------------------------------------------
-------------------------------------------------------------------------------

function SetCurrBattalion(num)
	CurrBattalion = num
    CurrSquad = false
	RefreshDisplay()
end

function SetCurrSquad(num)
	CurrSquad = num
	RefreshDisplay()
end

-------------------------------------------------------------------------------
-- Functions to get unit sets & stuff -----------------------------------------
-------------------------------------------------------------------------------

-- Return the label of desired squad
function GetSquadLabel(num, batt)
	batt = batt or CurrBattalion
	if (num == false) then num = 'X' end
	return batt .. ":" .. num
end

-- Return the squad num of battalion batt
function GetSquad(num, batt)
	batt = batt or CurrBattalion
	return selectionSets[GetSquadLabel(num, batt)]
end


-- Return all squads of battalion batt or all squads of current battalion
function GetAllSquads(batt)
	batt = batt or CurrBattalion
	local totalsel = {}
	for i=1,MaxSquad do
		local selectionset = GetSquad(i,batt)
		if selectionset then
			for i,unit in selectionset do
				table.insert(totalsel, unit)
			end
		end
	end
	return totalsel
end

-- Get the next unit, counting from the last unit in the selectionset.
-- Returning either the next unit in the unitset or the first unit not
-- included in the selectionset if the last unit in the unitset was
-- included in the selectionset.
function GetNextUnit(unitset, selunits)
	selunits = selunits or false
	local unit = false

	-- Select the last unit in the selected unit set
	if(selunits) then
		unit = selunits[table.getn(selunits)]
	end
	
	-- If unitset is empty, break
	local selectionpool = Filter(unitset)
	if not selectionpool then return end

	local returntable = {}
	-- if unit was null, then get the first unit
	if not unit then
		table.insert(returntable,selectionpool[1])
	else
		-- First, remove from the selectionpool table, all the units that are already selected
		-- except the one we're counting from.
		-- Make a hash table and check against a hash. This is O(n), not O(n^2) as most naive solutions
		local hashpool = {}
		for i,currunit in selunits do
			if(currunit ~= unit) then
				hashpool[currunit] = true
			end
		end
		-- Remove all that match the hash. Work from the end of the table so we don't have to worry about the indices changing
		for i=table.getn(selectionpool),1,-1 do
			if (hashpool[selectionpool[i]] == true) then
				table.remove(selectionpool,i)
			end
		end
		-- Then find the unit in the selectionpool and select the next unselected unit
		for i,currunit in selectionpool do
			if currunit == unit then
				-- If the unit is the last unit in the selectionset, add the first unit
				if not selectionpool[i+1] then
					table.insert(returntable, selectionpool[1])
				else
					table.insert(returntable, selectionpool[i+1])
				end
				break
			end
		end
	end
	return returntable
end

-- Starting from squad num, select the first nonempty squad
function GetFirstSquadWithUnits(num)
	num = num or 1
	local i = num
	while true do
		local selectionset = Filter(GetSquad(i))

		if table.getn(selectionset) > 0 then --(table.getn(selectionset) > 0) then
			SquadCycle = i
			return selectionset
		end
		i = i + 1
		if i > MaxSquad then
			i = 1
		end
		if (i == num) then
			return nil
		end
	end
end

-------------------------------------------------------------------------------
-- Autoselection functionality ------------------------------------------------
-------------------------------------------------------------------------------

function AutomaticSelection(labelToSelect)
	if labelToSelect == nil then return end
--	print("Doing automatic selection")
	if (options.gui_autoselect == 1) then
		CurrBattalion = tonumber(string.sub(labelToSelect,1,1))
		SelectSquad(tonumber(string.sub(labelToSelect,3,3)))
	end
	if (options.gui_autoselect == 2) then
		SelectBattalion(tonumber(string.sub(labelToSelect,1,1)))
	end
end


-------------------------------------------------------------------------------
-- Functions to select sets ---------------------------------------------------
-------------------------------------------------------------------------------

function SelectAllSquads(num)
	SelectUnits(Filter(GetAllSquads(num)))
end

function SelectSquad(num)
    local selunits = Filter(GetSelectedUnits())
    if (((num == CurrSquad) and (table.getn(selunits) > 0)) or ((CurrSquad == "B") and (SquadCycle == num))) then
		-- We're trying to reselect the current squad
		if table.getn(selunits) > 1 then
			-- Got a group selected. Go to first unit
			SelectUnits(GetNextUnit(GetSquad(num)))
		else
			SelectUnits(GetNextUnit(GetSquad(num), selunits))
		end
	else
		-- Selecting squad for the first time
		SelectUnits(Filter(GetSquad(num)))
    end
	-- This is always done
	SetCurrSquad(num)
end

function SelectBattalion(num)
	if ((num == CurrBattalion) and (CurrSquad == "B")) then
		-- Delegated to subroutine
		SelectNextSquadFromBattalion()
	else
		-- Selecting a battalion for the first time
		SelectUnits(Filter(GetAllSquads(num)))
		SquadCycle = 0
	end
	-- This is always done
	SetCurrBattalion(num)
	SetCurrSquad("B")
end

-- Reduce selection to only those units in current battalion
function FilterToBattalion()
	local totalsel = Filter(GetSelectedUnits())
	local finalsel = {}
	local label
	for i,unit in totalsel do
		for j=1,MaxSquad do
			label = GetSquadLabel(j)
			if unit:HasSelectionSet(label) then
				table.insert(finalsel, unit)
			end
		end
	end
	SelectUnits(finalsel)
	SetCurrSquad(false)
end

function SelectNextUnitFromBattalion(append)
	append = append or false
	local selunits = Filter(GetSelectedUnits())
	local addselection = {}
	-- If in any kind of squad selection mode, keep on doing that
	-- First simplify the following comparison a bit
	if (CurrSquad == "B") and (SquadCycle ~= 0) then
		SetCurrSquad(SquadCycle)
	end
	if (CurrSquad > 0) and (CurrSquad < MaxSquad +1) then
		if table.getn(selunits) > 0 then
			-- Look at the last unit and select the next
			addselection = GetNextUnit(Filter(GetSquad(CurrSquad)), selunits)
		else
			-- If selection is empty, select first
			addselection = GetNextUnit(Filter(GetSquad(CurrSquad)))
		end
	else
		-- Select from whole battalion in the same way
		if table.getn(selunits) > 0 then
			addselection = GetNextUnit(Filter(GetAllSquads(CurrBattalion)), selunits)
		else
			addselection = GetNextUnit(Filter(GetAllSquads(CurrBattalion)))
		end
	end
	if append then
		table.insert(selunits,addselection[1])
		SelectUnits(selunits)
	else
		SelectUnits(addselection)
	end
end

-- Ghaleon wanted this even though simply pressing the current group would do.
-- After my fixes he PROBABLY doesn't need this either, but here it is.
function SelectNextSquadFromBattalion(append)
	append = append or false
	-- See if a squad is selected and select the next if one is
	local addselection = {}
	if (SquadCycle > 0) and (SquadCycle < MaxSquad) then
		addselection = GetFirstSquadWithUnits(SquadCycle+1)
	else
		-- Select first squad
		addselection = GetFirstSquadWithUnits(1)
	end
	if append then
		AddSelectUnits(addselection)
	else
		SelectUnits(addselection)
	end
	SetCurrSquad("B")
end

function FactorySelect(num)
	num = num or CurrBattalion
	local unitset = GetAllSquads(num)
	unitset = EntityCategoryFilterDown(categories.FACTORY, ValidateUnitsList(unitset))
	SelectUnits(unitset)
	SetCurrBattalion(num)
end

-------------------------------------------------------------------------------
-- Functions to manipulate sets -----------------------------------------------
-------------------------------------------------------------------------------

-- Remove unwanted stuff such as factories and clean up the list
function Filter(unitset)
	return EntityCategoryFilterDown(categories.ALLUNITS - (categories.FACTORY - categories.MOBILE), ValidateUnitsList(unitset))
end

-- Create squad [num] in battalion batt
-- Format for label is currBatt:squadNum
function CreateSquad(num, batt)
	batt = batt or CurrBattalion
	ClearSquad(GetSelectedUnits())
	ConExecute("UI_MakeSelectionSet " .. GetSquadLabel(num,batt))
	SetCurrSquad(num)
end

function RemoveUnitFromSelectionSet(unit,set)
   if unit and set then
      for i, currunit in selectionSets[set] do
        if currunit == unit then
           table.remove(selectionSets[set], i)
        end
      end
   end
end

function ClearSquad(unitset)
	unitset = ValidateUnitsList(unitset)
	if unitset then
		for i,unit in unitset do
			local setlist = unit:GetSelectionSets()
			for i, set in setlist do
				unit:RemoveSelectionSet(set)
                RemoveUnitFromSelectionSet(unit, set)
			end
		end
	end
end

-- Insert selected units into squad, keep original selection
function AppendToSquad(num)
	local origsel = GetSelectedUnits()
	local totalsel = ValidateUnitsList(GetSquad(num))
	if origsel then
		for i,unit in origsel do
			table.insert(totalsel, unit)
		end
	end
	SelectUnits(totalsel)
	ClearSquad(totalsel)
	CreateSquad(num)
	SelectUnits(origsel)
	SetCurrSquad(false)
end


function FillCurrentBattalion(num, bytype)
	bytype = bytype or false
	num = num or DefaultMaxSquad
	FillBattalion(CurrBattalion,bytype,num)
end

function FillBattalion(num,bytype,numgroups)
	bytype = bytype or false
	numgroups = numgroups or DefaultMaxSquad
	SetCurrBattalion(num)

	local sel = GetSelectedUnits()
	local sortTables = {{},{},{},{},{},{},{},{},{},{}}

	if sel then
		-- Clear squads of selected units
		ClearSquad(sel)
		-- Clear squads of units in the battalion we are filling
		ClearSquad(GetAllSquads(num))
		-- If we have units, sort them into types
		table.sort(sel, function(unit1, unit2) bpid1 = unit1:GetBlueprint().BlueprintId bpid2 = unit2:GetBlueprint().BlueprintId return (bpid1 < bpid2) end)

		-- Split them into squads
		local index = 1

		if(bytype) then
			bpid1 = sel[1]:GetBlueprint().BlueprintId

			for i, unit in sel do
				bpid2 = unit:GetBlueprint().BlueprintId
				if bpid1 ~= bpid2 then
					bpid1 = bpid2
					if index < numgroups then
						index = index + 1
					else
						index = 1
					end
				end
				table.insert(sortTables[index], unit)
			end
		else
			for i, unit in sel do
				table.insert(sortTables[index], unit)
				if index < numgroups then
					index = index + 1
				else
					index = 1
				end
			end
		end

		-- Add them to squads
		for i=1, numgroups do
			SelectUnits(sortTables[i])
			if table.getn(sortTables[i]) > 0 then
				CreateSquad(i)
			end
		end
		SelectBattalion(num)
	else
		-- If no units selected, clear the battalion
		SelectBattalion(num)
		ClearSquad(GetSelectedUnits())
	end
end

function AppendSquadToSelection(num)
	AddSelectUnits(GetSquad(num))
end

function AppendBattalionToSelection(num)
	AddSelectUnits(GetAllSquads(num))
end

-------------------------------------------------------------------------------
-- Functions to modify settings -----------------------------------------------
-------------------------------------------------------------------------------

function ToggleAutoselectInversion()
	if autoselectinverse then
		autoselectinverse = false
		print('Autoselectinverse is Off')
	else
		autoselectinverse = true
		print('Autoselectinverse is On')
	end
end	

-------------------------------------------------------------------------------
-- Zeech's Stuff --------------------------------------------------------------
-------------------------------------------------------------------------------

function CascadeFire(delay)
   local battNum = CurrBattalion
   local totalSel = {}
   for i=1,MaxSquad do
      label = GetSquadLabel(i)
      local selectionSet = Filter(selectionSets[label])
      if selectionSet then
         local squad = i
         local BattalionFireThread = ForkThread(function()
                                                   local unitSet = selectionSet
                                                   SetFireState(unitSet, 1) -- Hold Fire
                                                   WaitSeconds(((squad-1) * delay) + 1)
                                                   unitSet = ValidateUnitsList(unitSet)
                                                   SetFireState(unitSet, 0) -- Return Ground
                                                end)
      end
   end
end

function HoldFire()
   local unitSet = EntityCategoryFilterDown(categories.ALLUNITS - categories.FACTORY, GetSelectedUnits())
   SetFireState(unitSet, 1) -- Hold Fire
end

function ReturnFire()
   local unitSet = EntityCategoryFilterDown(categories.ALLUNITS - categories.FACTORY, GetSelectedUnits())
   SetFireState(unitSet, 0) -- Return Fire
end

-------------------------------------------------------------------------------
-- Display functions ----------------------------------------------------------
-------------------------------------------------------------------------------

function RefreshDisplay()
	local batt = "X"
	local squad = "X"
	if CurrBattalion then
		batt = CurrBattalion
	end
	if CurrSquad then
		squad = CurrSquad
	end
	CurrBattalionDisplay:SetText(batt .. " : " .. squad)
end

function GetCurrentLabel()
	return GetSquadLabel(CurrSquad,CurrBattalion)
end

function CreatePanel(parent)
	local panel = Bitmap(parent)
	panel:SetTexture('/mods/W_Groups/textures/group_indicator_right.dds')
	panel.Right:Set(function() return parent.Right() - 270 end)
	panel.Top:Set(function() return parent.Top() + 5 end)
	panel.Depth:Set(100)
	panel.HandleEvent = function(self, event)
		if event.Type == 'MouseExit' then
			CloseToolTip()
		end
	end
	return panel
end

function CreateSplitControls(parent)
	local splitgroup = Group(parent)
	splitgroup.Left:Set(function() return parent.Left() + 42 end)
	splitgroup.Top:Set(function() return parent.Top() + 5 end)
	splitgroup.Right:Set(parent.Right)
	splitgroup.Height:Set(28)
	local title = UIUtil.CreateText(splitgroup, "X : X", 28)
	CurrBattalionDisplay = title
	title.HandleEvent = function(self, event)
		if event.Type == 'MouseEnter' then
			CreateToolTip(self, {Title = "Battalion : Squad", Body = "This panel displays the current Battalion and Squad"})
		end
	end
	LayoutHelpers.AtLeftTopIn(title, splitgroup, 5, 1)
	return splitgroup
end

function CreateToolTip(control, data, togglebutton)
	if mouseover then mouseover:Destroy() mouseover = false end
	local text = data.Title
	local body = data.Body
	mouseover = ToolTip.CreateExtendedToolTip(control, text, body)
	mouseover.Top:Set(function() return panel.Top() + 2 end)
	mouseover.Right:Set(function() return panel.Left() - 1 end)
	mouseover:SetAlpha(0, true)
	mouseover:SetNeedsFrameUpdate(true)

	local delayTime = 0.5
	local totaltime = 0
	mouseover.OnFrame = function(self, deltaTime)
		if totaltime > delayTime then
			local alpha = self:GetAlpha() + (deltaTime * 4)
			if alpha < 1 then
				self:SetAlpha(alpha, true)
			else
				self:SetAlpha(1, true)
				self:SetNeedsFrameUpdate(false)
			end
		end
		totaltime = totaltime + deltaTime
	end
end

function CloseToolTip()
	if mouseover then mouseover:Destroy() mouseover = false end
end

function hideEffect(control, subcontrol1)
	control:DisableHitTest(true)
    control:SetNeedsFrameUpdate(true)

	local timefade = 0.1
	local time1 = 0.1
	local time2 = 0.2
	local minHeight = 16
    local timeAccum = 0

	EffectHelpers.FadeOut(subcontrol1, timefade)

    control.OnFrame = function(self, elapsedTime)
        timeAccum = timeAccum + elapsedTime
		if timeAccum >= (time1 + time2) then
            self:SetNeedsFrameUpdate(false)
            self:Hide()
			self.Left:Set(function() return self.Right() - 3 end)
		elseif timeAccum >= time1 then
			self.Left:Set(function() return self.Right() - self.BitmapWidth()*((time2+time1 - timeAccum)/(time2)) end)
			self.Bottom:Set(function() return self.Top() + minHeight end)
		else
			self.Bottom:Set(function() return self.Top() + minHeight + (self.BitmapHeight() - minHeight)*((time1 - timeAccum)/(time1)) end )
		end
    end
end

function showEffect(control, subcontrol1)
	control:Show()
    control:SetNeedsFrameUpdate(true)

	local timefade = 0.1
	local time1 = 0.2
	local time2 = 0.1
	local minHeight = 16
    local timeAccum = 0

    control.OnFrame = function(self, elapsedTime)
        timeAccum = timeAccum + elapsedTime
		if timeAccum >= (time1 + time2) then
			self:EnableHitTest(true)
            self:SetNeedsFrameUpdate(false)
			LayoutHelpers.ResetLeft(self)
			LayoutHelpers.ResetBottom(self)
			EffectHelpers.FadeIn(subcontrol1, timefade)
		elseif timeAccum >= time1 then
			self.Bottom:Set(function() return self.Top() + minHeight + (self.BitmapHeight() - minHeight)*((timeAccum - time1)/(time2)) end )
			self.Left:Set(function() return self.Right() - self.BitmapWidth() end)
        else
			self.Left:Set(function() return self.Right() - self.BitmapWidth()*((timeAccum)/(time1)) end)
        end
    end
end

function Init()
	panel = CreatePanel(borders)
	splitPanel = CreateSplitControls(panel)
	panelToggle = UIUtil.CreateCheckbox(borders,
        UIUtil.UIFile('/game/avatar-tab-btn/tab-close_btn_up.dds'),
        UIUtil.UIFile('/game/avatar-tab-btn/tab-open_btn_up.dds'),
        UIUtil.UIFile('/game/avatar-tab-btn/tab-close_btn_over.dds'),
        UIUtil.UIFile('/game/avatar-tab-btn/tab-open_btn_over.dds'),
        UIUtil.UIFile('/game/avatar-tab-btn/tab-close_btn_dis.dds'),
        UIUtil.UIFile('/game/avatar-tab-btn/tab-open_btn_dis.dds'))
	LayoutHelpers.AtRightTopIn(panelToggle, borders, 259, 15)
	panelToggle.Height:Set(20)
	panelToggle.Width:Set(12)
	panelToggle.Depth:Set(101)
    panelToggle.OnCheck = function(self, checked)
        if checked then
			hideEffect(panel, splitPanel)
		else
			showEffect(panel, splitPanel)
		end
    end
	panelToggle.oldHandleEvent = panelToggle.HandleEvent
	panelToggle.HandleEvent = function(self, event)
		if event.Type == 'ButtonPress' then
			CloseToolTip()
			if self:IsChecked() then
				CreateToolTip(self, {Title = "Minimize", Body = "Minimize the group indicator."}, true)
			else
				CreateToolTip(self, {Title = "Maximize", Body = "Maximize the group indicator."}, true)
			end
		elseif event.Type == 'MouseEnter' then
			if self:IsChecked() then
				CreateToolTip(self, {Title = "Maximize", Body = "Maximize the group indicator."}, true)
			else
				CreateToolTip(self, {Title = "Minimize", Body = "Minimize the group indicator."}, true)
			end
		elseif event.Type == 'MouseExit' then
			CloseToolTip()
		end
		panelToggle.oldHandleEvent(self, event)
	end

	--AddConsoleOutputReciever(function(output) LOG(output) end) -- For con_listcommands

	-- The mod will simply crash if I refer to variables which are not set, so checking them now
	if (options.gui_max_autogroups == nil) then 
		options.gui_max_autogroups = 10
		DefaultMaxSquad = 10
	end

	if (options.gui_autoselect == nil) then
		options.gui_autoselect = 1
	end

	if (options.gui_autoselectinverse ~= 0) then
		autoselectinverse = true
	else
		autoselectinverse = false
	end

	if (options.gui_display_grouping_indicator == nil) then
		options.gui_display_grouping_indicator = 1
	end

	-- Hide on startup if disabled.
	if (options.gui_display_grouping_indicator == 0) then
--		hideEffect(panel, splitPanel)
		panel:Hide()
		panelToggle:Hide()
	end

	RefreshDisplay()
end
