Editing Module:UnitData

local getArgs = require('Module:Arguments').getArgs
local unit = {}
local unitData = mw.loadData('Module:UnitData/data')

local function toPercent(amount)
	-- amount is a string when ToSI is used before calling this function
	if type(amount) == "string" then
		amount = tonumber(amount)
	end
	
	return string.format("%.0f", amount * 100) .. "%"
end

local function appendHeatInfo(frame, ud, currentStr)
	local args = {subbox = "yes"}
	args.header1 = "Heat Up"
	args.label2 = "Heat Per Shot (%)"
	args.data2 = ud.heatPerShot
	args.label3 = "Heat Decay (%/s)"
	args.data3 = ud.heatDecay
	args.label4 = "Heat Max Slow (%)"
	args.data4 = ud.heatMaxSlow
	local tbl = {title = "infobox", args = args}
	currentStr = currentStr .. frame:expandTemplate({title = "Infobox", args = args})
	return currentStr
end

function unit.getData(frame)
	if not frame then return '' end
	local unitDefName = frame.args.defname or frame.args[1] or "<no name specified>"
	local ud = unitData[unitDefName]
	if not (ud) then return 'unitdef ' .. unitDefName .. ' not found' end
	local property = (frame.args.defname == nil) and frame.args[2] or frame.args[1]
	local result = ud[property]
	return result or ''
end

function unit.getAbilityInfoboxTemplate(abType, args)
	local tbl = {title = "Infobox zkability " .. abType, args = args}
	return tbl
end

function unit.getInfoboxTemplate(frame)
	local args = frame.args
	local tbl = {title = "Infobox zkunit", args = {}}
	local u = tbl.args
	local defName = args.defname or args[1] or "<no name specified>"
	local ud = unitData[defName]
	if not (ud) then return 'unitdef ' .. defName  .. ' not found' end
	
	for key, value in pairs(ud) do
		u[key] = args[key] or value
	end
	u.abilities = unit.printAbilityInfoboxes(frame)
	u.weapons = unit.printWeaponInfoboxes(frame)
	u.weaponIDs = nil
	u.shieldIDs = nil
	u.drones = nil
	
	return tbl
end

function unit.printAbilityInfoboxes(frame)
	if not frame then return '' end
	local unitDefName = frame.args.defname or frame.args[1] or "<no name specified>"
	local ud = unitData[unitDefName]
	if not (ud) then return 'unitdef ' .. unitDefName .. ' not found' end
	
	local result = ''
	if ud.drones then
		for index, entry in ipairs(ud.drones) do
			local args = {
				drone = "[[" .. entry.drone .. "]]", 
				maxdrones = entry.maxDrones,
				range = entry.range,
				interval = entry.interval,
				spawnsize = entry.spawnSize,
				buildtime = entry.buildTime,
				maxbuilding = entry.maxBuilding
			}
			result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("drone", args))
		end
	end
	
	if ud.buildSpeed then 
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("construction", {buildpower = ud.buildSpeed}))
	end
	if ud.transportLightSpeed or ud.transportMediumSpeed or ud.transportHeavySpeed then
		local light = ud.transportLightSpeed
		local medium = ud.transportMediumSpeed
		local heavy = ud.transportHeavySpeed
		local args = {}
		if light then args.light = light .. "%" end
		if medium then args.medium = medium .. "%" end
		if heavy then args.heavy = heavy .. "%" end
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("transport", args))
	end
	if ud.cloakCost or ud.idleCloak then
		local args = {}
		if (ud.cloakCost or 0) > 0 then
			args.upkeepidle = ud.cloakCost
			args.upkeepmobile = ud.cloakCostMoving
		elseif ud.idleCloak then
			args.customdata1 = "Only when idle"
			args.customdata2 = "Free and automated"
		end
		args.decloakradius = ud.decloakDistance
		if ud.decloakOnFire == false then
			args.customdata3 = "No decloak while shooting"
		end
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("cloak", args))
	end
	if ud.areaCloakUpkeep then
		local args = {name = "Area Cloak", radius = ud.areaCloakRadius, upkeep = ud.areaCloakUpkeep}
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("cloak", args))
	end
	if ud.shieldIDs then
		local weaponData = mw.loadData('Module:WeaponData/data')
		for index, shieldID in pairs(ud.shieldIDs or {}) do
			local args = {}
			local shieldDef = weaponData[shieldID]
			args.name = shieldDef.name
			args.strength = shieldDef.strength
			args.regen = shieldDef.regen
			args.regencost = shieldDef.regencost
			args.radius = shieldDef.radius
			result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("shield", args))
		end
	end
	if ud.radarDistance or ud.jammerDistance then
		local args = {}
		if ud.jammerDistance then
			if not ud.radarDistance then
				args.name = "Radar Jammer"
			end
			args.jam = ud.jammerDistance
		end
		if ud.radarDistance then
			args.radar = ud.radarDistance
		end
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("intel", args))
	end
	if ud.jumpRange then
		local args = {range = ud.jumpRange, reload = ud.jumpReload, speed = ud.jumpSpeed, midairjump = ud.midairJump and "Yes" or "No"}
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("jump", args))
	end
	if ud.idleRegen or ud.combatRegen or ud.waterRegen or ud.armoredRegen then
		local args = {}
		if ud.idleRegenTime then
			args.idleregen = ud.idleRegen
			args.timetoenable = ud.idleRegenTime
		elseif ud.combatRegen then
			args.combatregen = ud.combatRegen
		end
		if ud.waterRegen then
			args.waterregen = ud.waterRegen
			args.atdepth = ud.waterRegenDepth
		end
		if ud.armoredRegen then
			args.customlabel1 = "Closed regen (HP/s)"
			args.customdata1 = ud.armoredRegen
		end
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("regen", args))
	end
	if ud.morphTo then
		local args = {to = "[[" .. ud.morphTo .. "]]", cost = ud.morphCost, time = ud.morphTime, disabled = ud.combatMorph and "No" or "Yes"}
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("morph", args))
	end
	if ud.armorDamageReduction then
		local args = {reduction = ud.armorDamageReduction .. '%'}
		abilities = abilities .. "\n\t" .. writeHeaderLine("Infobox zkability armored")
		if ud.armorForceClose then
			args.special1 = "Forced for " .. ud.armorForceClose .. "s on damage"
		end
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("armored", args))
	end
	
	local args = {}
	local num = 1
	do
		if ud.isMex then
			args["customdata".. num] = "Extracts metal"
			num = num + 1
		end
		if ud.stealth then
			args["customdata".. num] = "Invisible to radar"
			num = num + 1
		end
		if ud.fireproof then
			args["customdata".. num] = "Immunity to afterburn"
			num = num + 1
		end
		if ud.dontFireAtRadar then
			args["customdata".. num] = "Can ignore unidentified targets"
			num = num + 1
		end
		if ud.instaSelfDestruct then
			args["customdata".. num] = "Instant self-destruction"
			num = num + 1
		end
	end
	if num > 1 then
		result = result .. frame:expandTemplate(unit.getAbilityInfoboxTemplate("line", args))
	end
	
	return result
end

function unit.printWeaponInfoboxes(frame)
	if not frame then return '' end
	local unitDefName = frame.args.defname or frame.args[1] or "<no name specified>"
	local ud = unitData[unitDefName]
	if not (ud) then return 'unitdef ' .. unitDefName .. ' not found' end
	
	local result = ''
	local weaponModule = require('Module:WeaponData')
	
	for index,weaponID in ipairs(ud.weaponIDs or {}) do
		local tempFrame = mw.getCurrentFrame():newChild{title="Module:WeaponData",args={defname=weaponID}}
		result = result .. weaponModule.printInfobox(tempFrame)
	end
	
	if ud.heatPerShot then
		result = appendHeatInfo(frame, ud, result)
	end
	
	if result == '' then result = "None" end
	return result
end

function unit.printInfobox(frame)
	return frame:expandTemplate(unit.getInfoboxTemplate(frame))
end

return unit