-- This widget allows units to disengage from combat. Specifically this is
-- useful for disengaging damaged units and sending them back to a safe location
-- where they can be repaired and returned to combat.
--
-- When a unit is disengaged, it will be deselected, removed from its unit
-- group, and (with UnicodeSnowman's patched smart_select) will not be
-- selectable with smart_select. This is to avoid accidentally re-engaging the
-- unit prematurely when executing sweaty micro. Once a unit is re-engaged
-- (reaches the fallback points or is given a new command or stopped) it will be
-- added back to the unit group and become selectable again.

function widget:GetInfo()
  return {
    name = "Disengage",
    desc = "Allows units to disengage from combat.",
    author = "UnicodeSnowman",
    date = "May 10, 2023",
    license = "GNU GPL, v2 or later",
    handler = true,
    layer = 5,
    enabled = true  --  loaded by default?
  }
end

widgetHandler = assert(widgetHandler)
widget = assert(widget)
Spring = assert(Spring)
WG = assert(WG)
CMD = assert(CMD)
CMDTYPE = assert(CMDTYPE)
gl = assert(gl)
GL = assert(GL)
UnitDefs = assert(UnitDefs)

local spacePerUnit = 70

local spSetUnitGroup = Spring.SetUnitGroup
local spGetUnitGroup = Spring.GetUnitGroup
local spGetCameraPosition = Spring.GetCameraPosition
local spGetViewGeometry = Spring.GetViewGeometry
local spTraceScreenRay = Spring.TraceScreenRay
local GL_LINE_STRIP = GL.LINE_STRIP
local glVertex = gl.Vertex
local glLineStipple = gl.LineStipple
local glLineWidth = gl.LineWidth
local glColor = gl.Color
local glBeginEnd = gl.BeginEnd

local ceil = math.ceil
local sqrt = math.sqrt
local max = math.max


local CMD_MOVE = CMD.MOVE
local CMD_DISENGAGE = 1133
local CMD_SET_FALLBACK_POSITION = 1134

local spGetUnitPosition = Spring.GetUnitPosition
local spGiveOrderToUnit = Spring.GiveOrderToUnit
local spGetSelectedUnits = Spring.GetSelectedUnits
local spGetGroundHeight = Spring.GetGroundHeight
local spGetTeamStartPosition = Spring.GetTeamStartPosition

local unitFallbackPosition = {}

local function tVerts(verts)
  for i = 1, #verts do
    local v = verts[i]
    glVertex(v[1], v[2], v[3])
  end
end

local function disengageUnit(uid)
  local pos = unitFallbackPosition[uid]
  if not pos then
    -- By default, run back to base.
    pos = spGetTeamStartPosition(0)
    if not pos then
      return false
    end
  end

  spGiveOrderToUnit(uid, CMD_MOVE, { pos.x, pos.y, pos.z }, 0)

  local group = spGetUnitGroup(uid)
  spSetUnitGroup(uid, -1)

  WG.disengaged[uid] = { originalGroup = group }
  return true
end

local function reengageUnit(uid)
  local unitInfo = WG.disengaged[uid]
  if unitInfo then
    WG.disengaged[uid] = nil
    spSetUnitGroup(uid, unitInfo.originalGroup or -1)
  end
end

local function disengageUnits()
  for _, u in pairs(spGetSelectedUnits()) do
    disengageUnit(u)
  end
end

function widget:UnitDestroyed(uid)
  unitFallbackPosition[uid] = nil
end

-- local function dump(o)
--   if type(o) == 'table' then
--     local s = '{ '
--     for k,v in pairs(o) do
--       if type(k) ~= 'number' then k = '"'..k..'"' end
--       s = s .. '['..k..'] = ' .. dump(v) .. ','
--     end
--     return s .. '} '
--   else
--     return tostring(o)
--   end
-- end

local function getLocationsForFront(n, center, right)
  -- The command give the center point and the rightmost point. We'll normalize
  -- this so the leftmost point is at the origin and this normalized vector is
  -- the rightmost position. It makes the math MUCH easier.
  local normalized = {
    x = (right.x - center.x) * 2,
    y = (right.y - center.y) * 2,
    z = (right.z - center.z) * 2 }

  -- Takes a vector relative to the normalized vector and denormalizes it back
  -- to the world coordinates.
  local function denorm(pos)
    return { x = (pos.x - normalized.x / 2) + center.x,
      y = (pos.y - normalized.y / 2) + center.y,
      z = (pos.z - normalized.z / 2) + center.z }
  end

  -- Get the length battle-front line.
  local xsq = normalized.x
  xsq = xsq * xsq
  local zsq = normalized.z
  zsq = zsq * zsq
  local llen2d = sqrt(xsq + zsq)

  local max_units_in_rank = llen2d / spacePerUnit
  local n_ranks = ceil(n / max_units_in_rank)

  local xstep = normalized.x / n
  local pos

  local function getY(x)
    return { x = x, y = normalized.y, z = (normalized.z / normalized.x) * x }
  end

  local x = 0
  if n_ranks <= 1 then
    -- If there is only one rank, spread the units along that front.
    pos = {}
    x = xstep / 2
    for i = 1, n do
      local p = denorm(getY(x))
      p.y = spGetGroundHeight(p.x, p.z)
      p.u = i
      table.insert(pos, p)
      x = x + xstep
    end
  else
    pos = {}

    -- Vector perpendicular to the drawn front vector. Used to create the ranks.
    local perp_vec = { x = - normalized.z, 0, z = normalized.x }

    -- Normalize the perpendicuar vector to be of length "spacePerUnit"
    -- This vector * rank will be added to each unit.
    local perp_vec_len = sqrt(
      perp_vec.x * perp_vec.x +
      perp_vec.z * perp_vec.z)
    perp_vec.x = perp_vec.x / perp_vec_len
    perp_vec.z = perp_vec.z / perp_vec_len
    perp_vec.x = perp_vec.x * spacePerUnit
    perp_vec.z = perp_vec.z * spacePerUnit

    xstep = normalized.x / max_units_in_rank
    for r = 1, n_ranks do
      x = 0
      for j = 1, max_units_in_rank do
        if #pos <= n then
          local point = getY(x)
          point.x = point.x + perp_vec.x * r
          point.z = point.z + perp_vec.z * r
          local p = denorm(point)
          p.u = #pos
          p.y = spGetGroundHeight(p.x, p.z)
          table.insert(pos, p)
          x = x + xstep
        end
      end
    end
  end
  return pos
end

-- When the Unit is done executing a command, the unit is then considered
-- re-engaged.
function widget:UnitCmdDone(uid)
  reengageUnit(uid)
end

function widget:UnitCommand(uid, _, _, cid)
  Spring.Echo("Unit Command: " .. uid .. ", " .. cid)

  if cid == CMD_DISENGAGE then
    disengageUnit(uid)
    return true
  end
end

function widget:CommandNotify(id, params, options)
  Spring.Echo("Command Notify: " .. id)

  if id == CMD_SET_FALLBACK_POSITION then
    -- Sets the unit's fallback point.
    --
    -- The command doesn't actually change what the unit is doing, it just sets
    -- the fallback point for disengaged units.
    local selectedUnits = Spring.GetSelectedUnits()
    if #params < 6 then
      for _, uid in selectedUnits do
        unitFallbackPosition[uid] = { x = params[1], y = params[2], z = params[3] }
      end
    else
      local locs = getLocationsForFront(
        #selectedUnits,
        { x = params[1], y = params[2], z = params[3] },
        { x = params[4], y = params[5], z = params[6] })

      for i = 1, #selectedUnits do
        unitFallbackPosition[selectedUnits[i]] = locs[i]
      end
    end

    return true
  elseif id == CMD_DISENGAGE then
    -- Tells the units to disengage from combat. Disengaged units will start
    -- moving to the fallback point and will no longer be selectable using smart
    -- select or the unit's group.
    --
    -- Once the unit is re-engaged, it will se selectable again.
    disengageUnits()
    return true
  else
    -- If the command is neither FALLBACK or SET_FALLBACK, re-engage the unit.
    for _, u in pairs(spGetSelectedUnits()) do
      reengageUnit(u)
    end
    return false
  end
end

local function drawGroundQuad(x,y,z,size)
  gl.TexCoord(0,0)
  gl.Vertex(x-size,y,z-size)
  gl.TexCoord(0,1)
  gl.Vertex(x-size,y,z+size)
  gl.TexCoord(1,1)
  gl.Vertex(x+size,y,z+size)
  gl.TexCoord(1,0)
  gl.Vertex(x+size,y,z-size)
end

local dotImage			= "LuaUI/Images/formationDot.dds"

local function drawCircle(pos, size, disengaged)
  if disengaged == true then
    glColor(1, 0.6, 0.6, 0.85)
  else
    glColor(0.75, 0.75, 0.75, 0.75)
  end
  gl.Texture(dotImage)
  gl.BeginEnd(GL.QUADS,drawGroundQuad, pos[1], pos[2], pos[3], size)
  gl.Texture(false)
end

local Xs, Ys = spGetViewGeometry()
Xs, Ys = Xs*0.5, Ys*0.5
function widget:ViewResize(viewSizeX, viewSizeY)
  Xs, Ys = spGetViewGeometry()
  Xs, Ys = Xs*0.5, Ys*0.5
end

function widget:DrawWorld()
  local camX, camY, camZ = spGetCameraPosition()
  local at, p = spTraceScreenRay(Xs,Ys,true,false,false)
  local zoomY

  if at == "ground" then
    local dx, dy, dz = camX-p[1], camY-p[2], camZ-p[3]
    zoomY = sqrt(dx*dx + dy*dy + dz*dz)
  else
    zoomY = camY - max(spGetGroundHeight(camX, camZ), 0)
  end

  if zoomY < 6 then zoomY = 6 end
  glLineStipple(2, 8031)
  glLineWidth(1.0)
  for _, uid in pairs(spGetSelectedUnits()) do
    local pos = unitFallbackPosition[uid]
    if pos then
      local x, y, z = spGetUnitPosition(uid)

      if WG.disengaged[uid] then
        glColor(1, 0.6, 0.6, 0.8)
      else
        glColor(0.75, 0.75, 0.75, 0.5)
      end

      glBeginEnd(GL_LINE_STRIP, tVerts, { {x, y, z}, { pos.x, pos.y, pos.z } })
      local dotSize = sqrt(zoomY*0.12)
      drawCircle(
        { pos.x, pos.y, pos.z },
          dotSize * (WG.disengaged[uid] and 2 or 1),
          WG.disengaged[uid] and true or false)
    end
  end
end

function widget:Initialize()
  WG['disengaged'] = {}

  WG.CMD_DISENGAGE = CMD_DISENGAGE
  WG.CMD_SET_FALLBACK_POSITION = CMD_SET_FALLBACK_POSITION
end

function widget:CommandsChanged()
  local selectedUnits = Spring.GetSelectedUnits()
  if #selectedUnits > 0 then
    local customCommands = widgetHandler.customCommands
    for i = 1, #selectedUnits do
      local udid = Spring.GetUnitDefID(selectedUnits[i])
      if UnitDefs[udid].canMove then
        customCommands[#customCommands + 1] = {
          id = CMD_SET_FALLBACK_POSITION,
          type = CMDTYPE.ICON_FRONT,
          tooltip = 'Set a fallback position for this unit.',
          name = 'SetFallback',
          cursor = 'areamex',
          action = 'setfallback',
        }

        customCommands[#customCommands + 1] = {
          id = CMD_DISENGAGE,
          type = CMDTYPE.ICON,
          tooltip = 'Disengage and fallback to a previously defined fallback position',
          name = 'Disengage',
          cursor = 'areamex',
          action = 'disengage',
        }
        return
      end
    end
  end
end