--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

function widget:GetInfo()
	return {
		name    = "Auto Engage",
		desc    = "Combat units act on the Engagement Predictor: they attack-move into fights it rates as won, " ..
		          "back away from fights it rates as lost, and walk back afterwards. Switch units on with the " ..
		          "Auto Engage order button. Needs Combat Advisor.",
		author  = "Thomas",
		date    = "2026-11-04",
		license = "GNU GPL, v2 or later",
		layer   = 0,
		enabled = true,
	}
end

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
-- How it works
--
-- Twice a second the managed units are grouped (any two within `LINK` elmos are
-- in the same group) and each group that has enemies nearby is scored with
-- WG.EngagementPredictor.Predict, the same Gather and Score the Combat
-- Advisor's panel and rings use. So the number this widget acts on is the
-- number you would see with that group selected.
--
--   * win chance >= "Attack above": attack-move (FIGHT) at the nearest enemy
--     group, if it is within "Engage range" and inside the leash.
--   * win chance <= "Avoid below": if any enemy could reach the group within a
--     couple of seconds, every unit steps back far enough to be outside that
--     reach, away from the enemy and bent towards your nearest factory. If no
--     enemy is that close yet, units walking in (attack-move) are stopped.
--   * in between: no new orders. That band is the hysteresis.
--
-- Fleeing is immediate; re-engaging after a retreat waits "Re-engage delay"
-- seconds, so a chance that wobbles round the threshold does not dance the
-- army back and forth. When the area is quiet again, units the widget moved
-- walk back to where they started, or pick up the attack-move you had given
-- them before it pulled them out.
--
-- Your own orders win. A unit following a move, attack, guard, build or any
-- other order you gave is left alone until it is idle again (attack-move and
-- patrol excepted: those are what the widget is for, so it may pull such a
-- unit out of a lost fight). Units the gadget is retreating, units in a
-- transport and units on hold position are never sent to attack.
--
-- Orders the widget gave are recognised by their coordinates, as in Auto
-- Retreat Plus: a queue whose first order is not one of ours is the player's.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

local spGetMyTeamID          = Spring.GetMyTeamID
local spGetSpectatingState   = Spring.GetSpectatingState
local spGetTeamUnits         = Spring.GetTeamUnits
local spGetUnitDefID         = Spring.GetUnitDefID
local spGetUnitPosition      = Spring.GetUnitPosition
local spGetUnitCommands      = Spring.GetUnitCommands
local spGetUnitHealth        = Spring.GetUnitHealth
local spGetUnitIsDead        = Spring.GetUnitIsDead
local spGetUnitTransporter   = Spring.GetUnitTransporter
local spGetUnitRulesParam    = Spring.GetUnitRulesParam
local spGetUnitStates        = Spring.GetUnitStates
local spGetUnitTeam          = Spring.GetUnitTeam
local spGetUnitsInCylinder   = Spring.GetUnitsInCylinder
local spGetSelectedUnits     = Spring.GetSelectedUnits
local spGiveOrderToUnit      = Spring.GiveOrderToUnit
local spGetGroundHeight      = Spring.GetGroundHeight
local spGetGaiaTeamID        = Spring.GetGaiaTeamID
local spWorldToScreenCoords  = Spring.WorldToScreenCoords
local spGetCmdDescIndex      = Spring.GetCmdDescIndex
local spGetActiveCmdDesc     = Spring.GetActiveCmdDesc
local spEcho                 = Spring.Echo

local floor, max, min, sqrt, huge = math.floor, math.max, math.min, math.sqrt, math.huge

local CUSTOM          = (Spring.Utilities and Spring.Utilities.CMD) or {}
local CMD_RAW_MOVE    = CUSTOM.RAW_MOVE or 31109
local CMD_MOVE        = CMD.MOVE
local CMD_FIGHT       = CMD.FIGHT
local CMD_PATROL      = CMD.PATROL
local CMD_STOP        = CMD.STOP
local CMD_AUTO_ENGAGE = 39417 -- unused by ZK's customcmds.lua (checked against v1.14.10.3)

local ENEMY_UNITS = Spring.ENEMY_UNITS or -4

local MAP_X, MAP_Z = Game.mapSizeX, Game.mapSizeZ
local EDGE = 64

local TICK       = 15   -- frames between decisions (0.5 s)
local LINK       = 450  -- elmos: two managed units this close are one group
local MAX_SCORED = 8    -- groups scored per tick; the rest wait their turn
local MATCH      = 8    -- elmos: an order this close to the one we gave is ours
local REISSUE    = 200  -- elmos a new target must move before the order is given again
local MIN_STEP   = 250  -- shortest backing-off move worth giving
local BLIP_REACH = 500  -- reach assumed for a radar contact of unknown type
local SPEED_TIME = 2    -- s of enemy movement added to its weapon range
local SPEED_CAP  = 300  -- elmos, cap on that

local ICON = 'LuaUI/Images/commands/Bold/fight.png'

--------------------------------------------------------------------------------
-- Options

local Forget -- forward declaration, used by the scope option

options_path = 'Settings/Unit Behaviour/Auto Engage'
options_order = {'manage', 'newOn',
                 'attackAbove', 'avoidBelow', 'engageRange', 'leash', 'safety', 'reengage',
                 'returnHome', 'resumeFight', 'overrideOrders', 'countBuildings',
                 'commanders', 'gunships', 'showLabels', 'showMarks'}
options = {
	manage = {
		name = 'Units it controls',
		desc = 'Which of your combat units Auto Engage gives orders to. The order button switches the '
		    .. 'selected units on or off in either mode.',
		type = 'radioButton',
		value = 'toggled',
		items = {
			{key = 'toggled', name = 'Units I switch on', desc = 'Off until you press the Auto Engage button.'},
			{key = 'all',     name = 'All my combat units', desc = 'On unless you switch a unit off.'},
		},
		OnChange = function()
			if Forget then
				Forget()
			end
		end,
	},
	newOn = {
		name = 'New factory units start on',
		desc = 'Units leaving a factory are switched on, after they reach the rally point.',
		type = 'bool',
		value = false,
	},
	attackAbove = {
		name = 'Attack above (win %)',
		desc = 'A group the Engagement Predictor gives at least this chance attack-moves at the nearest enemies.',
		type = 'number',
		min = 50, max = 95, step = 5, value = 65,
	},
	avoidBelow = {
		name = 'Avoid below (win %)',
		desc = 'A group given this chance or less backs out of the enemy\'s reach. Between the two '
		    .. 'thresholds it gets no new orders.',
		type = 'number',
		min = 5, max = 50, step = 5, value = 35,
	},
	engageRange = {
		name = 'Engage range',
		desc = 'Elmos from the edge of a group to the nearest enemy. Farther than this and a winning '
		    .. 'group does not go looking for that fight.',
		type = 'number',
		min = 200, max = 3000, step = 50, value = 900,
	},
	leash = {
		name = 'Leash',
		desc = 'Never attack a target farther than this from where the unit was when the widget '
		    .. 'first moved it. 0 = no leash: a winning group keeps rolling.',
		type = 'number',
		min = 0, max = 6000, step = 100, value = 1800,
	},
	safety = {
		name = 'Safety margin',
		desc = 'Elmos kept between a retreating group and the reach of the enemy (weapon range plus '
		    .. 'two seconds of its movement).',
		type = 'number',
		min = 0, max = 600, step = 25, value = 150,
	},
	reengage = {
		name = 'Re-engage delay (s)',
		desc = 'After backing off, wait this long before attacking or walking back. Fleeing is always immediate.',
		type = 'number',
		min = 0, max = 15, step = 0.5, value = 3,
	},
	returnHome = {
		name = 'Walk back afterwards',
		desc = 'When nothing is left in range, units the widget moved go back to where they were.',
		type = 'bool',
		value = true,
	},
	resumeFight = {
		name = 'Resume my attack-move',
		desc = 'A unit pulled out of an attack-move or patrol you gave picks it up again once the area is quiet.',
		type = 'bool',
		value = true,
	},
	overrideOrders = {
		name = 'Avoid even during my orders',
		desc = 'Also pull out units following your own move/attack/guard orders when the fight is lost, '
		    .. 'and give the order back afterwards. Off leaves them entirely to you.',
		type = 'bool',
		value = false,
	},
	countBuildings = {
		name = 'Count turrets and buildings',
		desc = 'Score fights with structures on both sides, whatever the Engagement Predictor panel is '
		    .. 'set to. Off would walk an army into a turret line it never counted.',
		type = 'bool',
		value = true,
	},
	commanders = {
		name = 'Include commanders',
		desc = 'Let it move commanders too.',
		type = 'bool',
		value = false,
		OnChange = function()
			if Forget then
				Forget()
			end
		end,
	},
	gunships = {
		name = 'Include gunships',
		desc = 'Let it move gunships. Planes (fighters, bombers) are never included.',
		type = 'bool',
		value = true,
		OnChange = function()
			if Forget then
				Forget()
			end
		end,
	},
	showLabels = {
		name = 'Show decisions',
		desc = 'Write ENGAGE / AVOID / HOLD and the win chance over each group with enemies near.',
		type = 'bool',
		value = true,
	},
	showMarks = {
		name = 'Mark switched-on units',
		desc = 'A ring under the selected units that are switched on.',
		type = 'bool',
		value = true,
	},
}

--------------------------------------------------------------------------------
-- Unit definitions

-- customparams arrive stringified, so "false" and "0" are off.
local function Flag(cp, key)
	local v = cp and cp[key]
	return v ~= nil and v ~= false and v ~= "false" and v ~= "0" and v ~= 0
end

local armedDef  = {} -- defID -> true: mobile, armed, not a plane or drone
local commDef   = {}
local gunshipDef = {}
local flyDef    = {}
local factoryDef = {}
local skipEnemy = {} -- dontcount (terraform patches and the like): neither threat nor target
local costOf    = {}

for defID = 1, #UnitDefs do
	local ud = UnitDefs[defID]
	local cp = ud.customParams or {}
	costOf[defID] = ud.metalCost or 0
	skipEnemy[defID] = Flag(cp, "dontcount")
	factoryDef[defID] = ud.isFactory and true or nil
	local fly = ud.canFly and true or false
	flyDef[defID] = fly or nil
	local hovering = fly and (ud.hoverAttack or ud.isHoveringAirUnit)
	gunshipDef[defID] = hovering and true or nil
	commDef[defID] = (Flag(cp, "dynamic_comm") or Flag(cp, "commtype")) or nil

	local armed = false
	local weapons = ud.weapons or {}
	for i = 1, #weapons do
		local wd = weapons[i].weaponDef and WeaponDefs[weapons[i].weaponDef]
		if wd and not wd.isShield then
			local dmg = tonumber(wd.customParams and wd.customParams.stats_damage)
			            or (wd.damages and wd.damages[0]) or 0
			if dmg > 0 then
				armed = true
			end
		end
	end
	local mobile = (ud.speed or 0) > 0 and ud.canMove ~= false and not ud.isFactory
	local plane = fly and not hovering
	if armed and mobile and not plane and not Flag(cp, "is_drone") and not skipEnemy[defID] then
		armedDef[defID] = true
	end
end

-- Category sets, keyed for the reach cache. Same reading of the weapon tables as the Combat
-- Advisor and Auto Retreat Plus: stats_damage, and mount-level onlyTargets keyed by lowercase
-- category, where an empty table hits everything.
local catsCache = {}
local function Cats(defID)
	local c = catsCache[defID]
	if c then
		return c
	end
	local set, list = {}, {}
	for cat in pairs(UnitDefs[defID].modCategories or {}) do
		local lc = string.lower(cat)
		set[lc] = true
		list[#list + 1] = lc
	end
	table.sort(list)
	c = {set = set, key = table.concat(list, " "), fly = flyDef[defID] or false}
	catsCache[defID] = c
	return c
end

local reachCache = {} -- catKey -> enemy defID -> elmos it can shoot that kind of unit from (0 = cannot)
local function Reach(enemyDef, cats)
	local byKey = reachCache[cats.key]
	if not byKey then
		byKey = {}
		reachCache[cats.key] = byKey
	end
	local r = byKey[enemyDef]
	if r then
		return r
	end
	r = 0
	local ud = UnitDefs[enemyDef]
	local weapons = ud.weapons or {}
	for i = 1, #weapons do
		local wpn = weapons[i]
		local wd = wpn.weaponDef and WeaponDefs[wpn.weaponDef]
		if wd and not wd.isShield and not (cats.fly and wd.type == "TorpedoLauncher") then
			local dmg = tonumber(wd.customParams and wd.customParams.stats_damage)
			            or (wd.damages and wd.damages[0]) or 0
			if dmg > 0 and not (wd.customParams and wd.customParams.interceptor) then
				local ot = wpn.onlyTargets
				local hits = (not ot) or (next(ot) == nil)
				if not hits then
					for cat in pairs(ot) do
						if cats.set[string.lower(cat)] then
							hits = true
							break
						end
					end
				end
				if hits then
					r = max(r, wd.range or 0)
				end
			end
		end
	end
	byKey[enemyDef] = r
	return r
end

--------------------------------------------------------------------------------
-- State

local myTeam = spGetMyTeamID()
local gaiaTeam = spGetGaiaTeamID()
local spectating = false

local mine     = {} -- unitID -> defID, my armed mobiles that have finished building
local factories = {} -- unitID -> {x, z}
local S        = {} -- unitID -> per-unit record, created on first use
local labels   = {} -- what was decided last tick, for drawing: {x, z, p, word}
local selected = {} -- eligible own units in the selection
local rotate   = 0  -- which group scores first this tick
local warned   = false

local function Rec(uid)
	local s = S[uid]
	if not s then
		s = {flipAt = -huge, orderAt = -huge}
		S[uid] = s
	end
	return s
end

local function Eligible(defID)
	if not armedDef[defID] then
		return false
	end
	if commDef[defID] and not options.commanders.value then
		return false
	end
	if gunshipDef[defID] and not options.gunships.value then
		return false
	end
	return true
end

local function IsOn(uid)
	local s = S[uid]
	if s and s.on ~= nil then
		return s.on
	end
	return options.manage.value == 'all'
end

-- Drop what the widget was doing with a unit (one, or all of them), keeping the on/off switch.
-- The unit keeps whatever order it has.
Forget = function(uid)
	local function Clear(s)
		s.mode, s.oid, s.ox, s.oz, s.anchorX, s.anchorZ, s.resume = nil, nil, nil, nil, nil, nil, nil
	end
	if uid then
		if S[uid] then
			Clear(S[uid])
		end
	else
		for _, s in pairs(S) do
			Clear(s)
		end
	end
end

local function AddUnit(uid, defID, team)
	if team ~= myTeam then
		return
	end
	if armedDef[defID] then
		mine[uid] = defID
	end
	if factoryDef[defID] then
		local x, _, z = spGetUnitPosition(uid)
		if x then
			factories[uid] = {x, z}
		end
	end
end

local function RemoveUnit(uid)
	mine[uid] = nil
	factories[uid] = nil
	S[uid] = nil
end

local function Rescan()
	mine, factories, S = {}, {}, {}
	myTeam = spGetMyTeamID()
	spectating = spGetSpectatingState()
	if spectating then
		return
	end
	local units = spGetTeamUnits(myTeam) or {}
	for i = 1, #units do
		local uid = units[i]
		local defID = spGetUnitDefID(uid)
		local bp = select(5, spGetUnitHealth(uid))
		if defID and (bp == nil or bp >= 1) then
			AddUnit(uid, defID, myTeam)
		end
	end
end

--------------------------------------------------------------------------------
-- Orders

local MOVE_IDS = {[CMD_MOVE] = true, [CMD_RAW_MOVE] = true}

local function NearXZ(ax, az, bx, bz, d)
	local dx, dz = ax - bx, az - bz
	return dx * dx + dz * dz <= d * d
end

-- "idle", "ours", "fight" (the player's attack-move or patrol), or "busy" (any other order).
local function QueueState(uid, s)
	local q = spGetUnitCommands(uid, 1)
	local c = q and q[1]
	if not c then
		return "idle"
	end
	local id, p = c.id, c.params
	if s.oid and p and p[3] then
		local same = (id == s.oid) or (MOVE_IDS[id] and MOVE_IDS[s.oid])
		if same and NearXZ(p[1], p[3], s.ox, s.oz, MATCH) then
			return "ours"
		end
	end
	if id == CMD_FIGHT or id == CMD_PATROL then
		return "fight", c
	end
	return "busy", c
end

local function Clamp(x, z)
	return min(max(x, EDGE), MAP_X - EDGE), min(max(z, EDGE), MAP_Z - EDGE)
end

local function Give(uid, s, cmdID, x, z, now)
	x, z = Clamp(x, z)
	local y = spGetGroundHeight(x, z) or 0
	spGiveOrderToUnit(uid, cmdID, {x, y, z}, 0)
	s.oid, s.ox, s.oz, s.orderAt = cmdID, x, z, now
end

local function SetMode(s, mode, now)
	if s.mode ~= mode then
		s.mode = mode
		s.flipAt = now
	end
end

-- The player's order, kept so it can be given back after the widget pulled the unit out.
local function KeepResume(s, qs, c)
	if s.resume or not c then
		return
	end
	if qs == "fight" or (qs == "busy" and options.overrideOrders.value) then
		local p = {}
		for k = 1, #(c.params or {}) do
			p[k] = c.params[k]
		end
		s.resume = {id = c.id, params = p}
	end
end

--------------------------------------------------------------------------------
-- Grouping

local function Groups(list, n)
	local cell = LINK
	local grid = {}
	local parent = {}
	for i = 1, n do
		parent[i] = i
	end
	local function Find(i)
		while parent[i] ~= i do
			parent[i] = parent[parent[i]]
			i = parent[i]
		end
		return i
	end
	local STRIDE = 100003
	for i = 1, n do
		local e = list[i]
		local gx, gz = floor(e.x / cell), floor(e.z / cell)
		e.gx, e.gz = gx, gz
		local key = gx * STRIDE + gz
		local b = grid[key]
		if b then
			b[#b + 1] = i
		else
			grid[key] = {i}
		end
	end
	local l2 = LINK * LINK
	for i = 1, n do
		local e = list[i]
		for ox = -1, 1 do
			for oz = -1, 1 do
				local b = grid[(e.gx + ox) * STRIDE + (e.gz + oz)]
				if b then
					for k = 1, #b do
						local j = b[k]
						if j > i then
							local f = list[j]
							local dx, dz = e.x - f.x, e.z - f.z
							if dx * dx + dz * dz <= l2 then
								local a, c = Find(i), Find(j)
								if a ~= c then
									parent[a] = c
								end
							end
						end
					end
				end
			end
		end
	end
	local byRoot, groups = {}, {}
	for i = 1, n do
		local r = Find(i)
		local g = byRoot[r]
		if not g then
			g = {}
			byRoot[r] = g
			groups[#groups + 1] = g
		end
		g[#g + 1] = list[i]
	end
	-- Stable order (lowest unitID first) so the round-robin does not reshuffle every tick.
	for i = 1, #groups do
		local g = groups[i]
		local lo = huge
		for k = 1, #g do
			if g[k].uid < lo then
				lo = g[k].uid
			end
		end
		g.lo = lo
	end
	table.sort(groups, function(a, b) return a.lo < b.lo end)
	return groups
end

--------------------------------------------------------------------------------
-- One group

local function Home(cx, cz)
	local best, bx, bz = huge, nil, nil
	for _, f in pairs(factories) do
		local dx, dz = f[1] - cx, f[2] - cz
		local d = dx * dx + dz * dz
		if d < best then
			best, bx, bz = d, f[1], f[2]
		end
	end
	if bx then
		return bx, bz
	end
	local x, _, z = Spring.GetTeamStartPosition(myTeam)
	if x and x > 0 then
		return x, z
	end
	return nil
end

-- What is around a group: the nearest enemy, the enemies that could hit it soon and how far it
-- must step back to be clear of them, and the nearest enemy cluster to attack.
local function Survey(g)
	local cx, cz = 0, 0
	for k = 1, #g do
		cx, cz = cx + g[k].x, cz + g[k].z
	end
	cx, cz = cx / #g, cz / #g
	local r = 0
	local catsList, seenKey = {}, {}
	for k = 1, #g do
		local e = g[k]
		local dx, dz = e.x - cx, e.z - cz
		r = max(r, sqrt(dx * dx + dz * dz))
		local c = Cats(e.defID)
		if not seenKey[c.key] then
			seenKey[c.key] = true
			catsList[#catsList + 1] = c
		end
	end

	local engage = options.engageRange.value
	local scan = r + max(engage, 1200) + SPEED_CAP
	local found = spGetUnitsInCylinder(cx, cz, scan, ENEMY_UNITS) or {}
	local sv = {cx = cx, cz = cz, r = r, nearGap = huge, need = 0, threats = 0,
	            tx = 0, tz = 0, tw = 0, enemies = {}}
	local safety = options.safety.value
	for i = 1, #found do
		local eid = found[i]
		if spGetUnitTeam(eid) ~= gaiaTeam then
			local defID = spGetUnitDefID(eid)
			if not (defID and skipEnemy[defID]) then
				local x, _, z = spGetUnitPosition(eid)
				if x then
					local dx, dz = x - cx, z - cz
					local gap = max(0, sqrt(dx * dx + dz * dz) - r)
					local en = {x = x, z = z, gap = gap, w = defID and max(costOf[defID], 1) or 100}
					sv.enemies[#sv.enemies + 1] = en
					if gap < sv.nearGap then
						sv.nearGap = gap
					end
					local reach
					if defID then
						reach = 0
						for k = 1, #catsList do
							reach = max(reach, Reach(defID, catsList[k]))
						end
						if reach > 0 then
							reach = reach + min((UnitDefs[defID].speed or 0) * SPEED_TIME, SPEED_CAP)
						end
					else
						reach = BLIP_REACH
					end
					if reach > 0 and gap < reach + safety then
						sv.threats = sv.threats + 1
						sv.need = max(sv.need, reach + safety - gap)
						sv.tx, sv.tz, sv.tw = sv.tx + x * en.w, sv.tz + z * en.w, sv.tw + en.w
					end
				end
			end
		end
	end
	if sv.tw > 0 then
		sv.tx, sv.tz = sv.tx / sv.tw, sv.tz / sv.tw
	end

	-- Attack point: the value-weighted middle of the enemies within 400 of the nearest one's
	-- distance, i.e. the nearest enemy cluster rather than the whole field.
	if sv.nearGap < huge then
		local ax, az, aw = 0, 0, 0
		for i = 1, #sv.enemies do
			local en = sv.enemies[i]
			if en.gap <= sv.nearGap + 400 then
				ax, az, aw = ax + en.x * en.w, az + en.z * en.w, aw + en.w
			end
		end
		sv.ax, sv.az = ax / aw, az / aw
	end
	return sv
end

local function Retreat(g, sv, now)
	-- Away from the threats, bent towards home unless home is on the far side of them.
	local ux, uz = sv.cx - sv.tx, sv.cz - sv.tz
	local ul = sqrt(ux * ux + uz * uz)
	if ul < 1 then
		ux, uz, ul = 0, 1, 1
	end
	ux, uz = ux / ul, uz / ul
	local hx, hz = Home(sv.cx, sv.cz)
	if hx then
		local vx, vz = hx - sv.cx, hz - sv.cz
		local vl = sqrt(vx * vx + vz * vz)
		if vl > 300 then
			vx, vz = vx / vl, vz / vl
			if ux * vx + uz * vz > -0.3 then
				ux, uz = ux + vx * 0.8, uz + vz * 0.8
				ul = sqrt(ux * ux + uz * uz)
				ux, uz = ux / ul, uz / ul
			end
		end
	end
	local step = max(sv.need + 100, MIN_STEP)
	for k = 1, #g do
		local e = g[k]
		local s = e.s
		local tx, tz = e.x + ux * step, e.z + uz * step
		local again = e.qs == "ours" and s.mode == "avoid" and s.oid and MOVE_IDS[s.oid]
		              and NearXZ(tx, tz, s.ox, s.oz, REISSUE)
		if not again then
			KeepResume(s, e.qs, e.c)
			if not s.anchorX then
				s.anchorX, s.anchorZ = e.x, e.z
			end
			Give(e.uid, s, flyDef[e.defID] and CMD_MOVE or CMD_RAW_MOVE, tx, tz, now)
			SetMode(s, "avoid", now)
		end
	end
end

-- Lost, but nothing can reach the group yet: just stop walking in.
local function HoldBack(g, now)
	for k = 1, #g do
		local e = g[k]
		local s = e.s
		local walkingIn = e.qs == "fight" or (e.qs == "ours" and s.mode == "attack")
		if walkingIn then
			KeepResume(s, e.qs, e.c)
			spGiveOrderToUnit(e.uid, CMD_STOP, {}, 0)
			s.oid, s.orderAt = nil, now
			SetMode(s, "avoid", now)
		end
	end
end

local function HoldPosition(uid)
	local st = spGetUnitStates(uid)
	return type(st) == "table" and st.movestate == 0
end

local function Attack(g, sv, now)
	local leash = options.leash.value
	local wait = options.reengage.value
	for k = 1, #g do
		local e = g[k]
		local s = e.s
		local free = e.qs == "idle" or e.qs == "ours"
		local cooling = s.mode == "avoid" and now - s.flipAt < wait
		if free and not cooling and not HoldPosition(e.uid) then
			local ax, az = s.anchorX or e.x, s.anchorZ or e.z
			if leash <= 0 or NearXZ(sv.ax, sv.az, ax, az, leash) then
				local again = e.qs == "ours" and s.mode == "attack" and s.oid == CMD_FIGHT
				              and NearXZ(sv.ax, sv.az, s.ox, s.oz, REISSUE)
				if not again then
					if not s.anchorX then
						s.anchorX, s.anchorZ = e.x, e.z
					end
					Give(e.uid, s, CMD_FIGHT, sv.ax, sv.az, now)
					SetMode(s, "attack", now)
				end
			end
		end
	end
end

-- Nothing to fight or flee: give back the player's order, or walk back, or let go.
local function Settle(g, now)
	local wait = options.reengage.value
	for k = 1, #g do
		local e = g[k]
		local s = e.s
		if s.mode and e.qs == "idle" and now - s.flipAt >= wait then
			if s.mode == "return" then
				Forget(e.uid)
			elseif s.resume and options.resumeFight.value then
				local r = s.resume
				spGiveOrderToUnit(e.uid, r.id, r.params, 0)
				Forget(e.uid)
			elseif options.returnHome.value and s.anchorX
			       and not NearXZ(e.x, e.z, s.anchorX, s.anchorZ, 120) then
				local ax, az = s.anchorX, s.anchorZ
				Give(e.uid, s, flyDef[e.defID] and CMD_MOVE or CMD_RAW_MOVE, ax, az, now)
				s.mode = "return" -- not SetMode: arriving needs no re-engage wait
				s.resume = nil
			else
				Forget(e.uid)
			end
		end
	end
end

--------------------------------------------------------------------------------
-- The tick

local function Tick(frame)
	local ep = WG.EngagementPredictor
	if not (ep and ep.Predict) then
		if not warned then
			warned = true
			spEcho("Auto Engage: waiting for Combat Advisor (its Engagement Predictor scores the fights).")
		end
		labels = {}
		return
	end
	warned = false
	local now = frame / 30

	local list, n = {}, 0
	for uid, defID in pairs(mine) do
		if Eligible(defID) and IsOn(uid) and spGetUnitIsDead(uid) == false
		   and not spGetUnitTransporter(uid) then
			local s = Rec(uid)
			if spGetUnitRulesParam(uid, "retreat") == 1 then
				-- The retreat gadget has it: hands off, and forget what we were doing.
				if s.mode then
					Forget(uid)
				end
			else
				local qs, c = QueueState(uid, s)
				if qs == "busy" and not options.overrideOrders.value then
					if s.mode then
						Forget(uid) -- the player gave it something else to do
					end
				else
					local x, _, z = spGetUnitPosition(uid)
					if x then
						n = n + 1
						list[n] = {uid = uid, defID = defID, s = s, qs = qs, c = c, x = x, z = z}
					end
				end
			end
		end
	end

	local groups = Groups(list, n)
	local newLabels = {}
	local scored = 0
	local count = #groups
	local attackAt = options.attackAbove.value / 100
	local avoidAt = options.avoidBelow.value / 100
	local engage = options.engageRange.value
	local skipped = nil -- the first group that needed scoring but was over the cap
	for step = 0, count - 1 do
		local g = groups[(rotate + step) % count + 1]
		local sv = Survey(g)
		local near = sv.threats > 0 or sv.nearGap <= engage
		if not near then
			Settle(g, now)
		elseif scored >= MAX_SCORED then
			skipped = skipped or step
		else
			scored = scored + 1
			local ids = {}
			for k = 1, #g do
				ids[k] = g[k].uid
			end
			local res = ep.Predict(ids, {includeStatic = options.countBuildings.value})
			if res and res.p then
				local p = res.p
				local word
				if p >= attackAt then
					word = "ENGAGE"
					if sv.nearGap <= engage then
						Attack(g, sv, now)
					end
				elseif p <= avoidAt then
					word = "AVOID"
					if sv.threats > 0 then
						Retreat(g, sv, now)
					else
						HoldBack(g, now)
					end
				else
					word = "HOLD"
				end
				newLabels[#newLabels + 1] = {x = sv.cx, z = sv.cz, r = sv.r, p = p, word = word}
			else
				Settle(g, now)
			end
		end
	end
	-- Next tick starts at the first group this one had no room for, so none can starve.
	if skipped then
		rotate = (rotate + skipped) % count
	end
	labels = newLabels
end

--------------------------------------------------------------------------------
-- The order button

local function RecountSelection(sel)
	selected = {}
	if spectating then
		return
	end
	sel = sel or spGetSelectedUnits() or {}
	for i = 1, #sel do
		local uid = sel[i]
		local defID = mine[uid]
		if defID and Eligible(defID) then
			selected[#selected + 1] = uid
		end
	end
end

local function Toggle()
	if #selected == 0 then
		return
	end
	local anyOff = false
	for i = 1, #selected do
		if not IsOn(selected[i]) then
			anyOff = true
			break
		end
	end
	for i = 1, #selected do
		local uid = selected[i]
		local s = Rec(uid)
		s.on = anyOff
		if not anyOff then
			Forget(uid)
		end
	end
	spEcho(("Auto Engage %s for %d unit%s."):format(anyOff and "on" or "off", #selected,
	                                                 #selected == 1 and "" or "s"))
	if Spring.ForceLayoutUpdate then
		Spring.ForceLayoutUpdate()
	end
end

function widget:SelectionChanged(sel)
	RecountSelection(sel)
end

function widget:CommandsChanged()
	if spectating or #selected == 0 then
		return
	end
	local on = 0
	for i = 1, #selected do
		if IsOn(selected[i]) then
			on = on + 1
		end
	end
	local state
	if on == 0 then
		state = "off"
	elseif on == #selected then
		state = "ON"
	else
		state = ("on for %d of %d"):format(on, #selected)
	end
	-- A widget without `handler = true` gets ZK's restricted widgetHandler, which has no
	-- customCommands table (indexing it took the widget down on its first selection).
	-- AddLayoutCommand is the supported route, valid only inside CommandsChanged.
	local cmd = {
		id      = CMD_AUTO_ENGAGE,
		type    = CMDTYPE.ICON,
		name    = "Auto Engage",
		action  = "autoengage",
		texture = ICON,
		params  = {},
		tooltip = "Auto Engage: " .. state .. "\n"
		       .. "Switched-on units attack fights the Engagement Predictor rates as won and back away "
		       .. "from fights it rates as lost. Click to switch the selection "
		       .. ((on == #selected) and "off." or "on."),
	}
	if widgetHandler.AddLayoutCommand then
		widgetHandler:AddLayoutCommand(cmd)
	elseif widgetHandler.customCommands then
		local cc = widgetHandler.customCommands
		cc[#cc + 1] = cmd
	end
end

local function IsStateCommand(cmdID)
	local idx = spGetCmdDescIndex and spGetCmdDescIndex(cmdID)
	local desc = idx and spGetActiveCmdDesc(idx)
	return desc and desc.type == CMDTYPE.ICON_MODE
end

function widget:CommandNotify(cmdID, params, opts)
	if cmdID == CMD_AUTO_ENGAGE then
		Toggle()
		return true
	end
	-- A real order from the player to the selection: it replaces whatever we were doing,
	-- including the order we meant to give back.
	if not IsStateCommand(cmdID) then
		for i = 1, #selected do
			Forget(selected[i])
		end
	end
	return false
end

--------------------------------------------------------------------------------
-- Callins

function widget:GameFrame(n)
	if spectating or n % TICK ~= 0 then
		return
	end
	Tick(n)
end

function widget:UnitFinished(unitID, unitDefID, unitTeam)
	AddUnit(unitID, unitDefID, unitTeam)
end

function widget:UnitFromFactory(unitID, unitDefID, unitTeam)
	if unitTeam == myTeam and options.newOn.value and armedDef[unitDefID] then
		Rec(unitID).on = true
	end
end

function widget:UnitGiven(unitID, unitDefID, newTeam)
	RemoveUnit(unitID)
	local bp = select(5, spGetUnitHealth(unitID))
	if bp == nil or bp >= 1 then
		AddUnit(unitID, unitDefID, newTeam)
	end
end

function widget:UnitTaken(unitID, unitDefID, oldTeam, newTeam)
	RemoveUnit(unitID)
end

function widget:UnitDestroyed(unitID)
	RemoveUnit(unitID)
end

function widget:PlayerChanged()
	if spGetMyTeamID() ~= myTeam or spGetSpectatingState() ~= spectating then
		Rescan()
		RecountSelection()
	end
end

function widget:Initialize()
	Rescan()
	RecountSelection()
end

function widget:Shutdown()
	-- Units keep whatever order they have; nothing is cancelled.
	labels = {}
end

--------------------------------------------------------------------------------
-- Drawing

local COLOUR = {
	ENGAGE = "\255\120\255\120",
	AVOID  = "\255\255\100\90",
	HOLD   = "\255\220\220\160",
}

function widget:DrawScreen()
	if not options.showLabels.value or #labels == 0 then
		return
	end
	for i = 1, #labels do
		local l = labels[i]
		local y = spGetGroundHeight(l.x, l.z) or 0
		local sx, sy, sz = spWorldToScreenCoords(l.x, y + 40, l.z)
		if sx and sz and sz < 1 then
			gl.Text(("%s%s %d%%"):format(COLOUR[l.word] or "", l.word, floor(l.p * 100 + 0.5)),
			        sx, sy, 14, "co")
		end
	end
end

function widget:DrawWorldPreUnit()
	if not options.showMarks.value or #selected == 0 then
		return
	end
	gl.LineWidth(1.5)
	for i = 1, #selected do
		local uid = selected[i]
		if IsOn(uid) then
			local x, y, z = spGetUnitPosition(uid)
			if x then
				local s = S[uid]
				local mode = s and s.mode
				if mode == "avoid" then
					gl.Color(1, 0.45, 0.35, 0.7)
				elseif mode == "attack" then
					gl.Color(0.5, 1, 0.5, 0.8)
				else
					gl.Color(0.55, 0.85, 1, 0.55)
				end
				gl.DrawGroundCircle(x, y, z, 26, 16)
			end
		end
	end
	gl.Color(1, 1, 1, 1)
	gl.LineWidth(1)
end