Guard-move 🛡️

function widget:GetInfo()
  return {
    name = "Smart Guard",
    desc = "An area-guard command that spreads guards evenly to the targets. Activate by double-pressing G.",
    author = "Thaddeus",
    layer = 0,
    enabled = true,
  }
end

local Debug = VFS.Include("LuaUI/Widgets/log.lua")

local G_KEY = string.byte("g")
local ESCAPE_KEY = 27
local LEFT_MOUSE_BUTTON = 1
local RIGHT_MOUSE_BUTTON = 3

local CIRCLE_DIVISIONS = 64
local CURSOR_LABEL = "Smart Guard"
local CURSOR_LABEL_OFFSET_X = 18
local CURSOR_LABEL_OFFSET_Y = -24
local CURSOR_LABEL_SIZE = 16
local DOUBLE_G_WINDOW = 0.4

local commandActive = false
local dragging = false
local lastGPressTime = -math.huge

local dragStartX
local dragStartZ
local dragCurrentX
local dragCurrentZ

local function GetGroundPosition(screenX, screenY)
  local hitType, hitData = Spring.TraceScreenRay(screenX, screenY, true)

  if hitType ~= "ground" or not hitData then
    return nil
  end

  return hitData[1], hitData[2], hitData[3]
end

local function GetDragRadius()
  if not dragging then
    return 0
  end

  local deltaX = dragCurrentX - dragStartX
  local deltaZ = dragCurrentZ - dragStartZ

  return math.sqrt(deltaX * deltaX + deltaZ * deltaZ)
end

local function CreateUnitSet(unitIDs)
  local unitSet = {}

  for i = 1, #unitIDs do
    unitSet[unitIDs[i]] = true
  end

  return unitSet
end

local function IsMobileNonBuilder(unitID)
  local unitDefID = Spring.GetUnitDefID(unitID)
  local unitDef = unitDefID and UnitDefs[unitDefID]

  return unitDef
    and unitDef.canMove
    and not unitDef.isBuilder
end

local function HasSelectedFactory()
  local selectedUnitIDs = Spring.GetSelectedUnits() or {}

  for i = 1, #selectedUnitIDs do
    local unitDefID = Spring.GetUnitDefID(selectedUnitIDs[i])
    local unitDef = unitDefID and UnitDefs[unitDefID]

    if unitDef and unitDef.isFactory then
      return true
    end
  end

  return false
end

local function IsFactoryUnitsTabKnownActive()
  if not HasSelectedFactory() then
    return false
  end

  local capsLockCommands = WG.CapsLockCommands

  return capsLockCommands
    and capsLockCommands.GetFactoryTab
    and capsLockCommands.GetFactoryTab() == "units"
end

local function GetGuardTargets(centerX, centerZ, radius, selectedUnitIDs)
  local selectedUnitSet = CreateUnitSet(selectedUnitIDs)
  local unitsInCircle = Spring.GetUnitsInCylinder(centerX, centerZ, radius) or {}
  local myAllyTeamID = Spring.GetMyAllyTeamID()
  local guardTargets = {}

  for i = 1, #unitsInCircle do
    local unitID = unitsInCircle[i]

    if not selectedUnitSet[unitID]
      and Spring.ValidUnitID(unitID)
      and not Spring.GetUnitIsDead(unitID)
      and Spring.GetUnitAllyTeam(unitID) == myAllyTeamID
      and IsMobileNonBuilder(unitID)
    then
      guardTargets[#guardTargets + 1] = unitID
    end
  end

  table.sort(guardTargets, function(unitA, unitB)
    local unitAX, _, unitAZ = Spring.GetUnitPosition(unitA)
    local unitBX, _, unitBZ = Spring.GetUnitPosition(unitB)

    if not unitAX then
      return false
    end

    if not unitBX then
      return true
    end

    local angleA = math.atan2(unitAZ - centerZ, unitAX - centerX)
    local angleB = math.atan2(unitBZ - centerZ, unitBX - centerX)

    if angleA == angleB then
      return unitA < unitB
    end

    return angleA < angleB
  end)

  return guardTargets
end

local function BuildRotatedTargetChain(guardTargets, startingIndex)
  local targetChain = {}

  for offset = 0, #guardTargets - 1 do
    local targetIndex = ((startingIndex + offset - 1) % #guardTargets) + 1
    targetChain[#targetChain + 1] = guardTargets[targetIndex]
  end

  return targetChain
end

local function IssueGuardChains(selectedUnitIDs, guardTargets)
  if #selectedUnitIDs == 0 then
    Debug.log("No selected guards")
    return
  end

  if #guardTargets == 0 then
    Debug.log("No valid guard targets")
    return
  end

  local assignments = {}

  for guardIndex = 1, #selectedUnitIDs do
    local guardUnitID = selectedUnitIDs[guardIndex]
    local startingTargetIndex = ((guardIndex - 1) % #guardTargets) + 1
    local targetChain = BuildRotatedTargetChain(
      guardTargets,
      startingTargetIndex
    )

    for chainIndex = 1, #targetChain do
      local commandOptions = {}

      if chainIndex > 1 then
        commandOptions = { "shift" }
      end

      Spring.GiveOrderToUnit(
        guardUnitID,
        CMD.GUARD,
        { targetChain[chainIndex] },
        commandOptions
      )
    end

    assignments[#assignments + 1] = {
      guardUnitID = guardUnitID,
      targetChain = targetChain,
    }
  end

  Debug.log("Guard chains", assignments)
end

local function ResetDrag()
  dragging = false
  dragStartX = nil
  dragStartZ = nil
  dragCurrentX = nil
  dragCurrentZ = nil
end

local function ActivateCommand()
  if commandActive then
    return false
  end

  local selectedUnitIDs = Spring.GetSelectedUnits() or {}

  if #selectedUnitIDs == 0 then
    Debug.log("Smart Guard not activated: no units selected")
    return false
  end

  if IsFactoryUnitsTabKnownActive() then
    Debug.log("Smart Guard not activated: factory Units tab active")
    return false
  end

  commandActive = true
  Debug.log("Smart Guard activated")
  return true
end

local function CancelCommand()
  if not commandActive then
    return
  end

  commandActive = false
  ResetDrag()
  Debug.log("Smart Guard cancelled")
end

function widget:Initialize()
  Debug.log("Smart Guard initialized")
end

function widget:Shutdown()
  Debug.log("Smart Guard shut down")
end

function widget:KeyPress(key, mods, isRepeat)
  if isRepeat then
    return false
  end

  if key == G_KEY and not mods.alt and not mods.ctrl and not mods.meta then
    local currentTime = os.clock()
    local isDoubleG = currentTime - lastGPressTime <= DOUBLE_G_WINDOW

    lastGPressTime = currentTime

    if isDoubleG then
      if ActivateCommand() then
        lastGPressTime = -math.huge
        Spring.SetActiveCommand(0)
        return true
      end

      return false
    end

    return false
  end

  if key == ESCAPE_KEY and commandActive then
    CancelCommand()
    return true
  end

  return false
end

function widget:MousePress(screenX, screenY, button)
  if not commandActive then
    return false
  end

  if button == RIGHT_MOUSE_BUTTON then
    CancelCommand()
    return true
  end

  if button ~= LEFT_MOUSE_BUTTON then
    return false
  end

  local worldX, _, worldZ = GetGroundPosition(screenX, screenY)

  if not worldX then
    return false
  end

  dragging = true
  dragStartX = worldX
  dragStartZ = worldZ
  dragCurrentX = worldX
  dragCurrentZ = worldZ

  Debug.log("Smart Guard drag started", {
    x = dragStartX,
    z = dragStartZ,
  })

  return true
end

function widget:MouseMove(screenX, screenY, dx, dy, button)
  if not dragging then
    return
  end

  local worldX, _, worldZ = GetGroundPosition(screenX, screenY)

  if not worldX then
    return
  end

  dragCurrentX = worldX
  dragCurrentZ = worldZ
end

function widget:MouseRelease(screenX, screenY, button)
  if not dragging or button ~= LEFT_MOUSE_BUTTON then
    return false
  end

  local worldX, _, worldZ = GetGroundPosition(screenX, screenY)

  if worldX then
    dragCurrentX = worldX
    dragCurrentZ = worldZ
  end

  local centerX = dragStartX
  local centerZ = dragStartZ
  local radius = GetDragRadius()
  local selectedUnitIDs = Spring.GetSelectedUnits() or {}
  local guardTargets = GetGuardTargets(
    centerX,
    centerZ,
    radius,
    selectedUnitIDs
  )

  Debug.log("Smart Guard drag finished", {
    centerX = centerX,
    centerZ = centerZ,
    radius = radius,
  })

  Debug.log("Selected guards", selectedUnitIDs)
  Debug.log("Valid targets", guardTargets)

  IssueGuardChains(selectedUnitIDs, guardTargets)

  commandActive = false
  ResetDrag()

  return true
end

function widget:DrawWorld()
  if not dragging then
    return
  end

  local radius = GetDragRadius()
  local centerY = Spring.GetGroundHeight(dragStartX, dragStartZ) + 3

  gl.LineWidth(4)
  gl.Color(0.3, 0.7, 1, 1)
  gl.DrawGroundCircle(
    dragStartX,
    centerY,
    dragStartZ,
    radius,
    CIRCLE_DIVISIONS
  )

  gl.LineWidth(1)
  gl.Color(1, 1, 1, 1)
end

function widget:DrawScreen()
  if not commandActive then
    return
  end

  Spring.SetMouseCursor("Guard")

  local mouseX, mouseY = Spring.GetMouseState()

  gl.Color(0.3, 0.7, 1, 1)
  gl.Text(
    CURSOR_LABEL,
    mouseX + CURSOR_LABEL_OFFSET_X,
    mouseY + CURSOR_LABEL_OFFSET_Y,
    CURSOR_LABEL_SIZE,
    "o"
  )
  gl.Color(1, 1, 1, 1)
end