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


function gadget:GetInfo()
  return {
    name      = "Antigravity Upgrades",
    desc      = "Not sure what the point is yet :P",
    author    = "quantum",
    date      = "Jul 24, 2007",
    license   = "GNU GPL, v2 or later",
    layer     = 0,
    enabled   = true  --  loaded by default?
  }
end


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



if (gadgetHandler:IsSyncedCode()) then
--------------------------------------------------------------------------------
--  SYNCED
--------------------------------------------------------------------------------



--Speed-ups
local GetUnitPosition     = Spring.GetUnitPosition
local GetUnitsInCylinder  = Spring.GetUnitsInCylinder
local GetUnitVelocity     = Spring.GetUnitVelocity
local GetUnitHealth       = Spring.GetUnitHealth
local GetAllUnits         = Spring.GetAllUnits
local GetUnitTeam         = Spring.GetUnitTeam
local SetUnitHealth       = Spring.SetUnitHealth
local DestroyUnit         = Spring.DestroyUnit
local EditUnitCmdDesc     = Spring.EditUnitCmdDesc
local RemoveUnitCmdDesc   = Spring.RemoveUnitCmdDesc
local InsertUnitCmdDesc   = Spring.InsertUnitCmdDesc
local FindUnitCmdDesc     = Spring.FindUnitCmdDesc
local UseTeamResource     = Spring.UseTeamResource
local SetVelocity         = Spring.MoveCtrl.SetVelocity
local SetRotationVelocity = Spring.MoveCtrl.SetRotationVelocity
local SetGravity          = Spring.MoveCtrl.SetGravity
local Enable              = Spring.MoveCtrl.Enable
local SetLeaveTracks      = Spring.MoveCtrl.SetLeaveTracks
local Disable             = Spring.MoveCtrl.Disable
local SetCollideStop      = Spring.MoveCtrl.SetCollideStop
local SetTrackGround      = Spring.MoveCtrl.SetTrackGround



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


local commDefs = {
  [ UnitDefNames['armcom'].id ] = true,
  [ UnitDefNames['corcom'].id ] = true,
}


local isComm         = {}


--antigrav field
local isUpgraded       = {}
local isControlled     = {}
local isFalling       = {}
local confirm          = {}
local range            = 300
local damageMult       = 200
local CMD_U_ANTIGRAV   = 32201
local uAntiGravTooltip = 
[[Upgrade: Antigravity 
Upgrades your Commander with antigravity generators.
Energy Cost: 20k]]


local uAntigravCmdDesc = {
  id      = CMD_U_ANTIGRAV,
  type    = CMDTYPE.ICON,
  name    = "Upgrade",
  cursor  = 'normal',
  action  = 'u_antigrav',
  tooltip = uAntiGravTooltip,
}

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


local function Lift(unitID)
  local rx, ry, rz = 0.005, 0.025, 0.025
  --local vx, vy, vz = 0, 1, 0
  print("Lift", unitID)
  Enable(unitID)                 --take control
  SetGravity(unitID, -0.1)
  --SetVelocity(unitID, vx, vy, vz)
  SetRotationVelocity(unitID, rx, ry, rz)  
  SetLeaveTracks(unitID, false)
end


local function Drop(unitID)
  isFalling[unitID] = {} 
  print("Drop", unitID)
  SetGravity(unitID, 1)
  SetCollideStop(unitID, true)
  SetTrackGround(unitID, true)
end


local function GetUnitsNearUnit(unitID, proximity)
  local cx, cy, cz = GetUnitPosition(unitID)
  local isNearArray = GetUnitsInCylinder(cx, cz, proximity)
  local isNearTable = {}
  for _, nearUnitID in ipairs(isNearArray) do
    if (not isComm[nearUnitID]) then
      isNearTable[nearUnitID] = true
    end
  end
  return isNearTable
end 


local function Difference(a, b)     --   c = a - b
  local c = {}
  for k, _ in pairs(a) do
    if (not b[k]) then
      c[k] = true
    end
  end
  return c
end 


local function DamageUnit(unitID)
  local damage = isFalling[unitID][2]*damageMult
  print("damage", damage)
  local oldHealth, maxHealth = GetUnitHealth(unitID)
  print("oldHealth", oldHealth)
  print("maxHealth", maxHealth)
  local newHealth = oldHealth + damage
  print("newHealth", newHealth)
  if (newHealth < 0) then
    print("destroy")
    DestroyUnit(unitID, true)
  else
    SetUnitHealth(unitID, newHealth)
  end
  Disable(unitID)
  isControlled[unitID] = nil
  isFalling[unitID] = nil
end


local function AddUAntigravCmdDesc(unitID)
  if (not FindUnitCmdDesc(unitID, CMD_U_ANTIGRAV)) then
    local insertID = 
      FindUnitCmdDesc(unitID, CMD.CLOAK)      or
      FindUnitCmdDesc(unitID, CMD.ONOFF)      or
      FindUnitCmdDesc(unitID, CMD.TRAJECTORY) or
      FindUnitCmdDesc(unitID, CMD.REPEAT)     or
      FindUnitCmdDesc(unitID, CMD.MOVE_STATE) or
      FindUnitCmdDesc(unitID, CMD.FIRE_STATE) or
      123456 -- back of the pack
    InsertUnitCmdDesc(unitID, insertID + 1, uAntigravCmdDesc)
  end
end


local function Upgrade(unitID, teamID)
  local cmdDescID = FindUnitCmdDesc(unitID, CMD_U_ANTIGRAV)
  if (not confirm[unitID]) then
    EditUnitCmdDesc(unitID, cmdDescID, {
      tooltip = "Confirm antigravity upgrade",
      name    = "Confirm"
    })
    confirm[unitID] = true
  else
    if (not UseTeamResource(teamID, {e = 20000})) then
      EditUnitCmdDesc(unitID, cmdDescID, {
        tooltip = "uAntiGravTooltip",
        name    = "Upgrade"
      })
      return
    end
    isUpgraded[unitID] = true
    RemoveUnitCmdDesc(unitID, cmdDescID)
    confirm[unitID] = nil
  end  
end
  
   
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------


function gadget:Initialize()
  gadgetHandler:RegisterCMDID(CMD_U_ANTIGRAV)  
  for _, unitID in ipairs(GetAllUnits()) do
    local teamID = GetUnitTeam(unitID)
    local unitDefID = GetUnitDefID(unitID)
    gadget:UnitCreated(unitID, unitDefID, teamID)
  end
end


function gadget:Shutdown()
  for _, unitID in ipairs(GetAllUnits()) do
    local cmdDescID = FindUnitCmdDesc(unitID, CMD_RETREAT)
    if (cmdDescID) then
      RemoveUnitCmdDesc(unitID, cmdDescID)
    end
  end
end


function gadget:UnitCreated(unitID, unitDefID)
  print("unit created", unitID)
  for comDefID, _ in pairs(commDefs) do
    if (comDefID == unitDefID) then
      isComm[unitID] = true
      isControlled[unitID] = {}
      AddUAntigravCmdDesc(unitID)
    end
  end
end


function gadget:UnitDestroyed(unitID)
  isComm[unitID]       = nil
  isControlled[unitID] = nil
  isFalling[unitID]   = nil
  isUpgraded[unitID]   = nil
end


function gadget:AllowCommand(unitID, unitDefID, teamID, cmdID, cmdParams)
  if (cmdID ~= CMD_U_ANTIGRAV) then
    return true
  end
  Upgrade(unitID, teamID)  
  return false
end

 
function gadget:GameFrame(n)
  if ((n % 34) < 0.1) then
    for commID, _ in pairs(isUpgraded) do
      local isInRange = GetUnitsNearUnit(commID, range)
      local entered   = Difference(isInRange, isControlled[commID])
      local left      = Difference(isControlled[commID], isInRange)
      for unitID, _ in pairs(entered) do
        isControlled[commID][unitID] = true
        Lift(unitID)
      end
      for unitID, _ in pairs(left) do
        if (not isFalling[unitID]) then
          Drop(unitID)
        end
      end
    end
  end
  for unitID, velocity in pairs(isFalling) do
    local vx, vy, vz = GetUnitVelocity(unitID)
    if (((vx == 0) and (vy == 0) and (vz == 0)) and   --stationary
        ((velocity[2] or 1) < 0)) then                 --was moving downwards
      DamageUnit(unitID)
    else
      isFalling[unitID] = {vx, vy, vz}
    end
  end
end 


--------------------------------------------------------------------------------
--  UNSYNCED
--------------------------------------------------------------------------------
end 


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