-- @author DoomTero 2026

ZLMain = {};

local ZomboRut = require("ZomboRut/ZomboRut");
require("ZomboRut/ZomboRutAnimationUtils")
local ZomboRutScene = require "ZomboRut/ZomboRutScene"

local ISWorldObjectContextMenu = ISWorldObjectContextMenu
local ISContextMenu = ISContextMenu
local ISToolTip = ISToolTip

local getText = getText
local string = string

return function(ContextMenu, playerObj, context, worldobjects)
	if not SandboxVars.ZomboRut.enableMod then return; end

	local player = playerObj
	local isFemale = playerObj:isFemale()

	local animationList
	if playerObj:isSitOnGround() then
		animationList = ContextMenu.Client.AnimationUtils:getAnimations(1, isFemale and 0 or 1, isFemale and 1 or 0, {"MastSit"})
	else
		animationList = ContextMenu.Client.AnimationUtils:getAnimations(1, isFemale and 0 or 1, isFemale and 1 or 0, {"Mast"})
	end

	local stats = player:getStats()
	local unhappiness = stats:get(CharacterStat.UNHAPPINESS)
	local stress = stats:get(CharacterStat.STRESS)
	local panic = stats:get(CharacterStat.PANIC)
	local sickness = stats:get(CharacterStat.FOOD_SICKNESS)
	local pain = stats:get(CharacterStat.PAIN)

	local function isComfortable()
		if pain >= 10 then return false; end
		if unhappiness >= 46 then return false; end
		if panic >= 6 then return false; end
		if sickness >= 25 then return false; end
		if stress >= 25 then return false; end
		return true
	end

	local function isInside()
		if ZR_Trait.has(player, "Innocent") and player:isOutside() then return false; end
		return true
	end

	local function isPrivate()
		local canSee = player:getSpottedList()
		local count = {}
		local maxDistance = 6
		local maxWatching = 1
		for i = 0, canSee:size() - 1 do
			local b = canSee:get(i)
			if b:isAlive() and b:DistTo(player) <= maxDistance then
				if b:getVariableBoolean("Bandit") then
					table.insert(count, b)
				end
			end
		end
		if ZR_Trait.has(player, "promiscuous") or ZR_Trait.has(player, "swinger") or ZR_Trait.has(player, "lewd") then
			maxWatching = math.huge
		end
		return #count <= maxWatching
	end

	local function isZombieFree()
		local zombs = player:getSpottedList()
		local maxDistance = 8
		for i = 0, zombs:size() - 1 do
			local z = zombs:get(i)
			if z:isAlive() and z:isZombie() and not z:getVariableBoolean("Bandit") and z:DistTo(player) <= maxDistance then
				return false
			end
		end
		return true
	end

	local SexOption = context:addOption(getText("ContextMenu_Main"), worldobjects)
	local SexSubMenu = ISContextMenu:getNew(context)
	context:addSubMenu(SexOption, SexSubMenu)

	if not isComfortable() or not isInside() or not isPrivate() or not isZombieFree() then
		local masturbateOption = SexSubMenu:addOption(getText("ContextMenu_Masturbate"), worldobjects)
		local tooltip = ISToolTip:new()
		tooltip:initialise()
		tooltip:setVisible(false)
		masturbateOption.notAvailable = true
		if not isComfortable() then
			tooltip.description = " <RED>" .. getText("ContextMenu_Mast_disc")
		elseif not isInside() then
			tooltip.description = " <RED>" .. getText("ContextMenu_Mast_inno")
		elseif not isPrivate() then
			tooltip.description = " <RED>" .. getText("ContextMenu_Mast_bandi")
		elseif not isZombieFree() then
			tooltip.description = " <RED>" .. getText("ContextMenu_Mast_zomb")
		end
		masturbateOption.toolTip = tooltip
	end

	if isComfortable() and isInside() and isPrivate() and isZombieFree() then
		local masturbateOption = SexSubMenu:addOption(getText("ContextMenu_Masturbate"), worldobjects)
		local masturbationSubMenu = ISContextMenu:getNew(context)
		SexSubMenu:addSubMenu(masturbateOption, masturbationSubMenu)

		local function playMasturbation(worldobjects, actors, animationData)
			local stripped = {}
			local callbacks = {
				Start = function(action)
					if action and action.character then
						ZomboRutScene._stripLowerBody(action.character, stripped)
					end
				end,
				Stop = function(action)
					if action and action.character then
						ZomboRutScene._reequipItems(action.character, stripped)
					end
				end,
				Perform = function(action)
					if action and action.character then
						ZomboRutScene._reequipItems(action.character, stripped)
					end
				end,
			}
			return ContextMenu.Client.AnimationHandler.Play(worldobjects, actors, animationData, false, true, callbacks)
		end

		for i = 1, #animationList do
			local animation = animationList[i]
			local key = string.format("ContextMenu_%s%s", animation.prefix, animation.id)
			local text = getText(key)
			if text then
				local toolTip = ISToolTip:new()
				toolTip.description = ""
				toolTip:setName(text)
				toolTip:initialise()
				toolTip:setVisible(false)
				local animationOption = masturbationSubMenu:addOption(text, worldobjects, playMasturbation, {playerObj}, animation)
				animationOption.toolTip = toolTip
			end
		end
	end
end
