-- Temporary server-side profiler. Load after entering the world; never ship with the mod.
-- Captures future Queen's Torch impacts only. No damage, targeting, or FX suppression.
local G = rawget(getfenv(1), "GLOBAL") or _G
assert(G.TheWorld and G.TheWorld.ismastersim, "Run this script on the server console")
assert(G.os and G.os.clock, "This profiler requires os.clock")
local previous_profile = rawget(G, "NS_TORCH_PROFILE")
if previous_profile then
	previous_profile.Stop()
end

local Combat = G.require("components/combat")
local Utils = G.require("ns_utils")
local manager = G.TheWorld.components.horrorchainmanager
assert(manager and G.EntityScript, "Night Stories and EntityScript must be loaded")
local clock = G.os.clock
local unpack = G.unpack
local prefix = "[DEBUG-ns-torch]"
local active
local remaining = 0
local sequence = 0
local label = "capture"
local undo = {}
local reports = {}
local installed = true
local api = {}
G.NS_TORCH_PROFILE = api

local function pack(...)
	return { n = select("#", ...), ... }
end

local function count(name)
	if active then
		active.counts[name] = (active.counts[name] or 0) + 1
	end
end

local function measure(name, fn, ...)
	if not active then
		return fn(...)
	end
	local sample = active
	local parent = sample.frame
	local frame = { children = 0 }
	sample.frame = frame
	local start = clock()
	local result = pack(pcall(fn, ...))
	local elapsed = clock() - start
	sample.frame = parent
	if parent then
		parent.children = parent.children + elapsed
	end
	local stat = sample.stats[name]
	if not stat then
		stat = { calls = 0, total = 0, own = 0 }
		sample.stats[name] = stat
	end
	stat.calls = stat.calls + 1
	stat.total = stat.total + elapsed
	stat.own = stat.own + elapsed - frame.children
	if not result[1] then
		error(result[2], 0)
	end
	return unpack(result, 2, result.n)
end

local function replace(object, key, wrapper)
	local previous = object[key]
	assert(type(previous) == "function", "Missing function: " .. key)
	local replacement = wrapper(previous)
	object[key] = replacement
	undo[#undo + 1] = function()
		if object[key] == replacement then
			object[key] = previous
		end
	end
end

local function report(sample)
	if not installed then
		return
	end
	print(
		string.format(
			"%s shot=%d label=%s CPU_ms=%.3f result=%s",
			prefix,
			sample.id,
			sample.label,
			sample.stats.impact.total * 1000,
			sample.failed and "error" or "ok"
		)
	)
	print(
		prefix
			.. " scope=server labelled phase; instrumented CPU time; excludes rendering and network delivery; chain-flush is reported separately"
	)
	print(prefix .. " inclusive columns overlap; exclusive columns exclude measured child calls")
	local names = {}
	for name in pairs(sample.stats) do
		names[#names + 1] = name
	end
	table.sort(names, function(a, b)
		return sample.stats[a].own > sample.stats[b].own
	end)
	for _, name in ipairs(names) do
		local stat = sample.stats[name]
		print(
			string.format(
				"%s %-30s calls=%d inclusive_ms=%.3f exclusive_ms=%.3f",
				prefix,
				name,
				stat.calls,
				stat.total * 1000,
				stat.own * 1000
			)
		)
	end
	names = {}
	for name in pairs(sample.counts) do
		names[#names + 1] = name
	end
	table.sort(names)
	for _, name in ipairs(names) do
		print(string.format("%s count %-30s %d", prefix, name, sample.counts[name]))
	end
end

-- Delayed chain damage must be measured separately, not mistaken for work removed.
if manager.QueueAttack and manager.FlushDamage then
	local pending_profile
	replace(manager, "QueueAttack", function(previous)
		return function(...)
			if active then
				pending_profile = { id = active.id, label = label .. "/chain-flush" }
			end
			return measure("chain.QueueAttack", previous, ...)
		end
	end)
	replace(manager, "FlushDamage", function(previous)
		return function(...)
			local capture = pending_profile
			pending_profile = nil
			if not capture or active or not installed then
				return previous(...)
			end
			local sample = { id = capture.id, label = capture.label, stats = {}, counts = {} }
			active = sample
			local result = pack(pcall(measure, "impact", previous, ...))
			active = nil
			sample.failed = not result[1]
			reports[#reports + 1] = G.TheWorld:DoTaskInTime(0, function()
				report(sample)
			end)
			if not result[1] then
				error(result[2], 0)
			end
			return unpack(result, 2, result.n)
		end
	end)
end

local function wrap_projectile(inst)
	if not inst or not inst.components.projectile then
		return
	end
	replace(inst.components.projectile, "Hit", function(previous)
		return function(self, target, ...)
			if not installed or remaining <= 0 or active then
				return previous(self, target, ...)
			end
			remaining = remaining - 1
			sequence = sequence + 1
			local sample = { id = sequence, label = label, stats = {}, counts = {} }
			active = sample
			local result = pack(pcall(measure, "impact", previous, self, target, ...))
			active = nil
			sample.failed = not result[1]
			reports[#reports + 1] = G.TheWorld:DoTaskInTime(0, function()
				report(sample)
			end)
			if not result[1] then
				error(result[2], 0)
			end
			return unpack(result, 2, result.n)
		end
	end)
end

replace(G, "SpawnPrefab", function(previous)
	return function(name, ...)
		local entity = measure("spawn." .. name, previous, name, ...)
		if remaining > 0 and name == "ns_queen_torch_projectile" then
			wrap_projectile(entity)
		end
		return entity
	end
end)

-- This seam also sees native attacks called directly by Horror Chain propagation.
replace(Utils, "WithAttackSource", function(previous)
	return function(source, fn, self, target, weapon, projectile, ...)
		local name = source and ("attack." .. (weapon and weapon.prefab or "unarmed")) or "attack.external"
		return measure(name, previous, source, fn, self, target, weapon, projectile, ...)
	end
end)

for _, method in ipairs({ "GetNearbyMembers", "AddMember" }) do
	local name = "chain." .. method
	replace(manager, method, function(previous)
		return function(...)
			return measure(name, previous, ...)
		end
	end)
end

replace(Combat, "GetAttacked_Internal", function(previous)
	return function(...)
		return measure("damage.GetAttacked", previous, ...)
	end
end)

local Weapon = G.require("components/weapon")
replace(Weapon, "OnAttack", function(previous)
	return function(self, ...)
		return measure("weapon.OnAttack." .. self.inst.prefab, previous, self, ...)
	end
end)

local function is_bag(inst)
	return inst.prefab == "punchingbag" or inst.prefab == "punchingbag_lunar" or inst.prefab == "punchingbag_shadow"
end

replace(G.EntityScript, "PushEvent", function(previous)
	return function(self, event, ...)
		if active and (event == "onhitother" or (is_bag(self) and (event == "healthdelta" or event == "attacked"))) then
			return measure("event." .. event, previous, self, event, ...)
		end
		return previous(self, event, ...)
	end
end)

replace(G.EntityScript, "DoTaskInTime", function(previous)
	return function(self, delay, fn, ...)
		if active and is_bag(self) then
			count("bag.tasks_scheduled")
			return measure("bag.schedule", previous, self, delay, fn, ...)
		end
		return previous(self, delay, fn, ...)
	end
end)

function api.Arm(shots, name)
	remaining = math.max(1, math.min(10, math.floor(tonumber(shots) or 3)))
	label = tostring(name or "capture"):gsub("[\r\n]", " ")
	print(string.format("%s armed: %d future projectile impacts; label=%s", prefix, remaining, label))
end

function api.Stop()
	if not installed then
		return
	end
	installed = false
	remaining = 0
	for _, task in ipairs(reports) do
		task:Cancel()
	end
	for i = #undo, 1, -1 do
		undo[i]()
	end
	G.NS_TORCH_PROFILE = nil
	print(prefix .. " stopped; hooks restored")
end

api.Arm(3, "enhanced-40")