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

function gadget:GetInfo()
  return {
    name      = "Ammo Limiter",
    desc      = "Gives each unit a personal 'ammo' storage that it draws from to fire, when empty it fires much more slowly",
    author    = "Quantum, screwed up horribly by Nemo",
    date      = "Feb 01, 2007",
    license   = "CC attribution-noncommerical 3.0 or later",
    layer     = 0,
    enabled   = true  --  loaded by default?
  }
end
	
  
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------


--[[
a note on the customParams (custom FBI tags) used by this script:

maxAmmo=the total ammo capacity of this unit;
ammoSupplier=boolean; //is this a ammo supplying unit?
supplyRange=x; //how far this supply unit can supply (at this distance ammoRegen is default)
ammoWarningLevel=0-1; the level (% of total ammo) when this unit displays some kind of "low ammo" graphic
weaponswithammo; // Number of weapons that use ammo. Must be the first ones. Default is 2.

------
TODO:
implement trucks and halftracks and such somehow
visuals (low ammo graphic, supply range graphic, ect)
make more customizable/flexible (custom params for ammoPenalty, multi-weapon support, 
]]--

--and here we go
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

local SetWeaponState     = Spring.SetWeaponState
local GetWeaponState     = Spring.GetWeaponState
local GetUnitsInCylinder = Spring.GetUnitsInCylinder
local GetUnitPosition    = Spring.GetUnitPosition
local GetUnitSeparation  = Spring.GetUnitSeparation


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

local vehicles      = {}
local ammoSuppliers = {}

--what reload time is multiplied by when the unit runs out of ammo
local ammoPenalty = 1.5


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

local function AmmoRegen(defaultRegen, supplyRange, nearestDepotDistance)
  return ((0-1)*(defaultRegen/supplyRange)*nearestDepotDistance + defaultRegen)
end


local function DefaultRegen(unitID)
  local reload = UnitDefs[unitID].weapons[1].reload
  return 2/reload
end


local function CheckReload(unitID, reloadFrame, weaponNum)
  local oldReloadFrame = vehicles[unitID].reloadFrame[weaponNum]
  if (oldReloadFrame == reloadFrame or
      not reloadFrame ) then
    return false
  else
    vehicles[unitID].reloadFrame[weaponNum] = reloadFrame
    return true
  end
end


local function ProcessWeapon(unitID, weaponNum)
  
  local _, _, reloadFrame = GetUnitWeaponState(unitID, weaponNum)
  local ammoLevel = vehicles[unitID].ammoLevel
  local reload = WeaponDefs[weaponID].reload
  local weaponID = UnitDefs[unitID].weapons[weaponNum].weaponDef
  
  if (CheckReload(unitID, reloadFrame, weaponNum)) then
    if (ammoLevel > 0) then
      vehicles[unitID].ammoLevel = ammoLevel - 1
    end
  end
  
  if (ammoLevel < 1) then
    SetUnitWeaponStates(unitID, weaponNum, {reloadtime = ammoPenalty * reload})
  else
    SetUnitWeaponStates(unitID, weaponNum, {reloadtime = reload})
  end
  
end


local function FindSupplier(vehicleID)
  local closestSupplier
  local closestDistance = math.huge
  
  for supplierID in pairs(ammoSuppliers) do
    local separation  = GetUnitSeparation(vehicleID, supplierID, true)
    if (separation < closestDistance) then
      closestSupplier = supplierID
      closestDistance = separation
    end
  end
  
  return closestSupplier, closestDistance
end


local function Resupply(unitID)
  
  local supplierID, distanceFromSupplier = FindSupplier(unitID)
  
  if (not supplierID) then
    return
  end
  
  local supplyRange = UnitDefs[supplierID].customParams.supplyRange
  local maxAmmo = UniDefs[unitID].customParams.maxAmmo
  local defaultRegen = DefaultRegen(unitID)
  local ammoRegen = AmmoRegen(defaultRegen, supplyRange, nearestDepotDistance)
  local oldAmmo = vehicles[unitID].ammoLevel
  
  local newAmmo = oldAmmo + ammoRegen
  if (newAmmo > maxAmmo) then
    newAmmo = maxAmmo
  end
  
  vehicles[unitID].ammoLevel = newAmmo
  
end
    
    
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

function gadget:UnitCreated(unitID, unitDefID, unitTeam, builderID)
  local ud = UnitDefs[unitDefID]
  if (string.find(ud.moveData.name, "TANK")) then
    vehicles[unitID] = {
      ammoLevel       = UnitDefs[unitID].customParams.maxammo,
      weaponsWithAmmo = UnitDefs[unitID].customParams.weaponswithammo,
    }
  end
  if (ud.customParams.ammoSupplier == 1) then
    ammoSuppliers[unitID] = true
  end
end


function gadget:UnitDestroyed(unitID) -- you can omit unneeded arguments if they
  vehicles[unitID] = nil              -- are at the end
  ammoSuppliers[unitID] = nil
end


function gadget:GameFrame(n)
  if ((n+3 % 30*2) > 0.01) then
    return
  end
  for unitID in pairs(vehicles) do
    local weaponsWithAmmo = vehciles[unitID].weaponsWithAmmo or 2
    for weaponNumber=1, weaponsWithAmmo do
      ProcessWeapon(unitID, weaponNum)
    end
    Resupply(unitID)
  end
end


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