-- @author DoomTero 2026

local ZomboRutScene = require "ZomboRut/ZomboRutScene"

ZR_Surrender = ZR_Surrender or {}

local DEBUG = false
local function log(fmt, ...)
	if DEBUG then print(string.format("[ZR Surrender] " .. fmt, ...)) end
end

local TRIGGER_RANGE_FIRST = 1.5
local TRIGGER_RANGE_CHAIN = 5

local BANDIT_TRIGGER_RANGE = 8

local PACIFY_RANGE = 8

local SCENE_POOL = {
	"DefeatPunch_DefeatPunch_DefeatPunch",
	"DefeatBJ_DefeatBJ_DefeatBJ",
	"DefeatLunge_DefeatLunge_DefeatLunge",
	"DefeatLesbianFaceSit_DefeatLesbianFaceSit_DefeatLesbianFaceSit",
	"DefeatForcedCowGirl_DefeatForcedCowGirl_DefeatForcedCowGirl",
}

local function rollSceneId(player, zombie)
	local pg = (player and player.isFemale and player:isFemale()) and "F" or "M"
	local zg = (zombie and zombie.isFemale and zombie:isFemale()) and "F" or "M"
	local exactSuffix = "_" .. pg .. zg

	local exact = {}
	local fallback = {}
	for i = 1, #SCENE_POOL do
		local id = SCENE_POOL[i]
		if ZomboRutScene.get(id .. exactSuffix) then
			exact[#exact + 1] = id
		elseif ZomboRutScene._resolveGender(player, zombie, id) then
			fallback[#fallback + 1] = id
		end
	end
	local pool = (#exact > 0) and exact or fallback
	if #pool == 0 then return nil end
	return pool[ZombRand(#pool) + 1]
end

local DEFAULT_SURRENDER_DURATION_TICKS = 300

local DEFAULT_ESCAPE_DURATION_TICKS = 600

local active = {}
local tickCounter = 0

local function nowTick() return tickCounter end

local endSurrender

local fireCallbacks

local function applyPlayerLock(player, sceneActive)
	pcall(function()
		if player.setSprinting        then player:setSprinting(false)        end
		if player.setRunning          then player:setRunning(false)          end
		if player.setForceSprint      then player:setForceSprint(false)      end
		if player.setIgnoreMovement   then player:setIgnoreMovement(true)    end
		if player.setBlockMovement    then player:setBlockMovement(true)     end
		if player.setIgnoreContextKey then player:setIgnoreContextKey(true)  end
		if player.setIgnoreAutoVault  then player:setIgnoreAutoVault(true)   end
	end)
	if not sceneActive then
		pcall(function()
			if player.StopAllActionQueue then player:StopAllActionQueue() end
		end)
	end
end

local function applyEscapeMobility(player)
	pcall(function()
		if player.setIgnoreMovement   then player:setIgnoreMovement(false)   end
		if player.setBlockMovement    then player:setBlockMovement(false)    end
		if player.setIgnoreContextKey then player:setIgnoreContextKey(false) end
		if player.setIgnoreAutoVault  then player:setIgnoreAutoVault(false)  end
	end)
end

local function clearPlayerEffects(player)
	pcall(function()
		if player.setIgnoreMovement   then player:setIgnoreMovement(false)   end
		if player.setBlockMovement    then player:setBlockMovement(false)    end
		if player.setIgnoreContextKey then player:setIgnoreContextKey(false) end
		if player.setIgnoreAutoVault  then player:setIgnoreAutoVault(false)  end
	end)
end

local function isHostileBandit(z)
	if not z:getVariableBoolean("Bandit") then return false end
	if not (Bandit and BanditBrain and Bandit.IsHostile) then return false end
	return Bandit.IsHostile(z) and true or false
end

local function isAggressor(z)
	if z:getVariableBoolean("Bandit") then return isHostileBandit(z) end
	return true
end

local function isAggressorFor(z, paciSet)
	if not z:getVariableBoolean("Bandit") then return true end
	if Bandit and Bandit.IsHostile and Bandit.IsHostile(z) then return true end
	local rec = paciSet[z]
	if rec and rec.isBandit
		and (rec.wasHostile or rec.wasHostileP)
	then
		return true
	end
	return false
end

local function pacifyOne(z, paciSet)
	if not paciSet[z] then
		local snap = {
			wasNoTeeth = (z.isNoTeeth and z:isNoTeeth()) or false,
		}
		if z:getVariableBoolean("Bandit") and BanditBrain then
			local brain = BanditBrain.Get(z)
			if brain then
				snap.isBandit      = true
				snap.wasHostile    = brain.hostile
				snap.wasHostileP   = brain.hostileP
				snap.wasStationary = brain.stationary
			end
		end
		paciSet[z] = snap
	end
	pcall(function() if z.setNoTeeth then z:setNoTeeth(true) end end)
	if z:getVariableBoolean("Bandit") and Bandit then
		pcall(function()
			if Bandit.SetHostile      then Bandit.SetHostile(z, false)     end
			if Bandit.SetHostileP     then Bandit.SetHostileP(z, false)    end
			if Bandit.ForceStationary then Bandit.ForceStationary(z, true) end
			if Bandit.ClearTasks      then Bandit.ClearTasks(z)            end
			pcall(function() z:clearAggroList() end)
			pcall(function() z:setTarget(nil)   end)
		end)
	end
end

local function pacifyNearbyZombies(player, paciSet)
	local cell = getCell()
	local zList = cell and cell:getZombieList()
	if not zList then return end
	for i = 0, zList:size() - 1 do
		local z = zList:get(i)
		if z and z:isAlive() and isAggressorFor(z, paciSet) then
			local inRange    = z:DistTo(player) <= PACIFY_RANGE
			local isAttacker = z.isZombieAttacking and z:isZombieAttacking(player)
			local isTargeter = z.getTarget and z:getTarget() == player
			if inRange or isAttacker or isTargeter then
				pacifyOne(z, paciSet)
			end
		end
	end
end

local function unpacifyZombies(paciSet)
	for z, rec in pairs(paciSet) do
		pcall(function()
			if z and z:isAlive() and z.setNoTeeth then
				z:setNoTeeth(rec.wasNoTeeth)
			end
			if rec.isBandit and z and z:isAlive() and BanditBrain then
				local brain = BanditBrain.Get(z)
				if brain then
					brain.hostile    = rec.wasHostile
					brain.hostileP   = rec.wasHostileP
					brain.stationary = rec.wasStationary
				end
			end
		end)
	end
end

local function ensureEmote(player)
	pcall(function()
		if player.playEmote then player:playEmote("surrender") end
	end)
end

local function stopEmote(player)
	pcall(function()
		if player.setVariable   then player:setVariable("EmotePlaying", false) end
		if player.clearVariable then player:clearVariable("emote")             end
	end)
end

local function forceClearHitReaction(player)
	pcall(function()
		if IdleState and IdleState.instance and player.changeState then
			local state = player.getCurrentStateName and player:getCurrentStateName() or ""
			if state == "PlayerHitReactionState"
				or state == "PlayerHitReactionPVPState"
			then
				player:changeState(IdleState.instance())
			end
		end
	end)
end

local function knockDown(z)
	pcall(function()
		if z.setKnockedDown then z:setKnockedDown(true) end
		local fallState = z.tryGetAIState and z:tryGetAIState("falldown")
		if fallState and z.changeState then
			z:changeState(fallState)
		end
		if z.clearAggroList then z:clearAggroList() end
		if z.setTarget      then z:setTarget(nil)     end
		if z:getVariableBoolean("Bandit") and Bandit and Bandit.ClearTasks then
			Bandit.ClearTasks(z)
		end
	end)
end

local function pronePropeNearbyZombies(player, paciSet)
	local count = 0
	local seen  = {}

	if paciSet then
		for z, _ in pairs(paciSet) do
			if z and z:isAlive() and isAggressorFor(z, paciSet) then
				knockDown(z)
				seen[z] = true
				count = count + 1
			end
		end
	end

	local cell = getCell()
	local zList = cell and cell:getZombieList()
	if zList then
		for i = 0, zList:size() - 1 do
			local z = zList:get(i)
			if z and not seen[z] and z:isAlive()
				and isAggressor(z)
				and z:DistTo(player) <= PACIFY_RANGE
			then
				knockDown(z)
				count = count + 1
			end
		end
	end

	if count > 0 then
		log("pronePropeNearbyZombies: knocked down %d aggressor(s)", count)
	end
	return count
end

local function transitionToEscape(player, st)
	st.phase = "escape"
	st.startTick = nowTick()
	local escDur = DEFAULT_ESCAPE_DURATION_TICKS
	if SandboxVars and SandboxVars.ZomboRut
		and SandboxVars.ZomboRut.escapeDuration
	then
		escDur = SandboxVars.ZomboRut.escapeDuration * 60
	end
	st.endTick = nowTick() + escDur
	applyEscapeMobility(player)
	stopEmote(player)
	pronePropeNearbyZombies(player, st.pacified)

	st.participated = {}
	log("ESCAPE phase started (duration=%d ticks); participated[] reset", escDur)
	fireCallbacks(ZR_Surrender._escapeCallbacks, player, st)
end

function ZR_Surrender.isActive(player)
	return active[player] ~= nil
end

function ZR_Surrender.getPhase(player)
	local st = active[player]
	return st and st.phase or nil
end

local function countAggroedZombies(player)
	local cell = getCell()
	local zList = cell and cell:getZombieList()
	if not zList then return 0 end
	local count = 0
	for i = 0, zList:size() - 1 do
		local z = zList:get(i)
		if z and z:isAlive()
			and isAggressor(z)
			and z:DistTo(player) <= 8
		then
			count = count + 1
		end
	end
	return count
end

function ZR_Surrender.start(player)
	if not player or not player:isAlive() then return false end

	local isSprinting   = player.isSprinting    and player:isSprinting()
	local isRunning     = player.isRunning      and player:isRunning()
	local isMoving      = player.isPlayerMoving and player:isPlayerMoving()
	if isSprinting or isRunning or isMoving then
		log("start refused: player still moving (sprinting=%s running=%s isMoving=%s)",
			tostring(isSprinting), tostring(isRunning), tostring(isMoving))
		player:Say(getText("ContextMenu_ZR_Surrender_RefuseStillRunning")
			or "I need to stop first.")
		return false
	end

	if ZR_Trait and ZR_Trait.has and ZR_Trait.has(player, "dominate") then
		local bd = player.getBodyDamage and player:getBodyDamage()
		local hp = bd and bd.getHealth and bd:getHealth() or 100
		if hp > 33 then
			log("start refused: dominate trait + HP %d > 33", hp)
			player:Say(getText("ContextMenu_Dominate_RefuseSurrender") or "I'm not done fighting yet!")
			return false
		end
	end

	if active[player] and active[player].phase == "surrender" then
		log("start refused: already in surrender phase")
		return false
	end

	if active[player] and active[player].phase == "escape" then
		log("start: ending active escape phase to re-enter surrender")
		endSurrender(player, active[player])
	end

	if ZomboRutScene.isPlayerInScene(player) then
		log("start refused: already in a scene")
		return false
	end

	local duration = DEFAULT_SURRENDER_DURATION_TICKS
	if SandboxVars and SandboxVars.ZomboRut
		and SandboxVars.ZomboRut.surrenderDuration
	then
		duration = SandboxVars.ZomboRut.surrenderDuration * 60
	end

	local credits = math.max(1, countAggroedZombies(player))

	active[player] = {
		phase           = "surrender",
		startTick       = nowTick(),
		endTick         = nowTick() + duration,
		pacified        = {},
		sceneActive     = false,
		scenesRemaining = credits,
		participated    = {},
	}
	applyPlayerLock(player, false)
	forceClearHitReaction(player)
	pacifyNearbyZombies(player, active[player].pacified)
	ensureEmote(player)
	log("SURRENDER started for player (duration=%d ticks, %d aggroed scene credits)",
		duration, credits)
	fireCallbacks(ZR_Surrender._startCallbacks, player, active[player])
	return true
end

function ZR_Surrender.standUp(player)
	local st = active[player]
	if not st then return false end
	if st.phase ~= "surrender" then return false end
	transitionToEscape(player, st)
	log("SURRENDER manually cancelled by player → ESCAPE")
	return true
end

ZR_Surrender._startCallbacks  = ZR_Surrender._startCallbacks  or {}
ZR_Surrender._escapeCallbacks = ZR_Surrender._escapeCallbacks or {}
ZR_Surrender._endCallbacks    = ZR_Surrender._endCallbacks    or {}

function ZR_Surrender.registerStartCallback(fn)
	if type(fn) ~= "function" then return end
	table.insert(ZR_Surrender._startCallbacks, fn)
end

function ZR_Surrender.registerEscapeCallback(fn)
	if type(fn) ~= "function" then return end
	table.insert(ZR_Surrender._escapeCallbacks, fn)
end

function ZR_Surrender.registerEndCallback(fn)
	if type(fn) ~= "function" then return end
	table.insert(ZR_Surrender._endCallbacks, fn)
end

fireCallbacks = function(list, player, st)
	for i = 1, #list do
		local ok, err = pcall(list[i], player, st)
		if not ok then log("callback error: %s", tostring(err)) end
	end
end

endSurrender = function(player, st)
	fireCallbacks(ZR_Surrender._endCallbacks, player, st)
	clearPlayerEffects(player)
	stopEmote(player)
	unpacifyZombies(st.pacified or {})
	active[player] = nil
	log("surrender mechanic ended — player back to normal")
end

local function onPlayerUpdate(player)
	tickCounter = tickCounter + 1

	local st = active[player]
	if not st then return end

	if st.phase == "surrender" then
		applyPlayerLock(player, st.sceneActive)
		pacifyNearbyZombies(player, st.pacified)

		if st.sceneActive then return end

		ensureEmote(player)

		if nowTick() >= st.endTick then
			transitionToEscape(player, st)
		end

	elseif st.phase == "escape" then
		applyEscapeMobility(player)
		pacifyNearbyZombies(player, st.pacified)

		if nowTick() >= st.endTick then
			endSurrender(player, st)
		end
	end
end

Events.OnPlayerUpdate.Add(onPlayerUpdate)

local POST_SCENE_GRACE_TICKS = 120

ZomboRutScene.registerSceneEndCallback(function(player, _active, result)
	local st = active[player]
	if not st then return end
	st.sceneActive = false
	if st.phase ~= "surrender" then return end

	forceClearHitReaction(player)
	pacifyNearbyZombies(player, st.pacified)

	if result == "Released" then
		st.scenesRemaining = (st.scenesRemaining or 1) - 1
		log("scene-end: %d scene credits remaining (result=%s)",
			st.scenesRemaining, tostring(result))
	end

	if (st.scenesRemaining or 0) <= 0 then
		log("scene credits exhausted → forcing ESCAPE phase")
		transitionToEscape(player, st)
	else
		st.endTick = math.max(st.endTick, nowTick() + POST_SCENE_GRACE_TICKS)
	end
end)

local function canFireScene(player, zombie)
	local st = active[player]
	if not st or st.phase ~= "surrender" or st.sceneActive then return false end
	if not zombie:isAlive() then return false end

	if st.participated[zombie] then return false end

	if zombie:getVariableBoolean("Bandit") then
		local rec = st.pacified[zombie]
		local wasHostile = rec and rec.isBandit
			and (rec.wasHostile or rec.wasHostileP)
		if not wasHostile then return false end
		if zombie:DistTo(player) > BANDIT_TRIGGER_RANGE then return false end
	else
		local isFirst = true
		for _ in pairs(st.participated) do isFirst = false; break end
		local triggerRange = isFirst and TRIGGER_RANGE_FIRST or TRIGGER_RANGE_CHAIN
		if zombie:DistTo(player) > triggerRange then return false end

		local isAttacking = zombie.isZombieAttacking and zombie:isZombieAttacking(player)
		local isTargeting = zombie.getTarget and zombie:getTarget() == player
		if not isAttacking and not isTargeting then return false end
	end

	if SandboxVars and SandboxVars.ZomboRut
		and SandboxVars.ZomboRut.enableMod == false
	then
		return false
	end
	return true
end

local function forcePacify(zombie, st)
	pacifyOne(zombie, st.pacified)
end

local function onZombieUpdate(zombie)
	if not zombie or not zombie:isAlive() then return end

	for player, st in pairs(active) do
		if player and player:isAlive()
			and zombie:DistTo(player) <= PACIFY_RANGE
			and isAggressorFor(zombie, st.pacified)
		then
			forcePacify(zombie, st)

			if canFireScene(player, zombie) then
				st.sceneActive = true
				st.participated[zombie] = true
				local sceneId = rollSceneId(player, zombie)
				local launcher = zombie:getVariableBoolean("Bandit")
					and ZomboRutScene.playSceneFromLive
					or ZomboRutScene.playSceneFromAttacker
				log("aggressor id=%s (%s) attacking surrendering player → firing scene %s",
					tostring(zombie:getID()),
					zombie:getVariableBoolean("Bandit") and "bandit" or "zombie",
					sceneId)
				local started = launcher(player, zombie, sceneId)
				if not started then
					st.sceneActive = false
					forceClearHitReaction(player)
					pacifyNearbyZombies(player, st.pacified)
					st.scenesRemaining = (st.scenesRemaining or 1) - 1
					log("aggressor id=%s no scene variant available → counted as completed, %d credits left",
						tostring(zombie:getID()), st.scenesRemaining)
					if (st.scenesRemaining or 0) <= 0 then
						log("scene credits exhausted (all completed) → forcing ESCAPE phase")
						transitionToEscape(player, st)
					else
						st.endTick = math.max(st.endTick, nowTick() + POST_SCENE_GRACE_TICKS)
					end
				end
				return
			end
		end
	end
end

Events.OnZombieUpdate.Add(onZombieUpdate)

return ZR_Surrender
