-- This script reproduces the C functions that SRB2 uses in order to build cooperative AI.
-- Base AI logic written in lua by CobaltBW, converted to C by Lach, then converted back into lua by CobaltBW (lol)

-- ***
-- Creates constants: AI_STANDBY, AI_FOLLOW, AI_CATCHUP, AI_THINKFLY, AI_FLYSTANDBY, AI_FLYCARRY, AI_SPINFOLLOW
-- Creates functions: B_UpdateBotLeader, B_BuildTailsTiccmd, B_HandleFlightIndicator
-- If B_BuildTailsTiccmd or B_HandleFlightIndicator are used, the table player.botmemory is also created.

-- These functions will be mainly be useful if you need to override the default AI but still be able to call the logic it uses,
-- or it can be useful if you simply want to use the default bot behavior as a base for your own AI logic.


if type(B_BuildTailsTiccmd) == "function" -- Either we've already done this, or this version of SRB2 already exposes B_BuildTailsTiccmd
	return
end

local makeBotMemory = function(player)
	-- player.botmem is the struct used in C to store various time-related AI thought processes.
	-- However, the struct and its contents are read-only, so we must assign this table under a different name instead.
	player.botmemory = $ or { 
		lastForward = false,
		lastBlocked = false,
		blocked = false,
		catchup_tics = 0,
		thinkstate = 0
	}
end

addHook("BotTiccmd", function(player, cmd) -- This hook jumpstarts the AI to run the custom lua instead of the native C code. You can modify or replace this hook to suit your own needs.
	if gamestate != GS_LEVEL return end
	B_UpdateBotleader(player)
	if player.botleader
		B_BuildTailsTiccmd(player.botleader.mo, player.mo, cmd)
	end
	B_HandleFlightIndicator(player)
	return true -- Must return true, or else the vanilla AI will write over the cmd.
end)

-- This simulates the function used to update the botleader. The logic behind the source code is pretty half-assed, so you'll probably want to modify it to suit your own needs.
rawset(_G, "B_UpdateBotleader", function(player)
	local i
	local dist
	local neardist = INT32_MAX
	local nearplayer
	//Find new botleader
	for p in players.iterate() do
		if p.bot or p.playerstate != PST_LIVE or p.spectator or not (p.mo and p.mo.valid)
			continue
		end
		if (not player.botleader)
			player.botleader = p // set default
			return
		end

		if (not player.mo)
			return
		end
		//Update best candidate based on nearest distance
		dist = R_PointToDist2(player.mo.x, player.mo.y, p.mo.x, p.mo.y)
		if (neardist > dist)
			neardist = dist
			nearplayer = p
		end
	end
	//Set botleader to best candidate (or null if none available)
	player.botleader = nearplayer
end)


-- Simulate the constants used by the source ai code.
rawset(_G, "AI_STANDBY", 0)
rawset(_G, "AI_FOLLOW", 1)
rawset(_G, "AI_CATCHUP", 2)
rawset(_G, "AI_THINKFLY", 3)
rawset(_G, "AI_FLYSTANDBY", 4)
rawset(_G, "AI_FLYCARRY", 5)
rawset(_G, "AI_SPINFOLLOW", 6)

-- Simulate the function used by the source ai code to build cooperative AI
rawset(_G, "B_BuildTailsTiccmd", function(sonic, tails, cmd)

	local	forward,	backward,	left,	right,	jump,	spin	=
			false,		false,		false,	false,	false,	false

	local player, bot = sonic.player, tails.player
	makeBotMemory(bot) -- bot.botmemory needs to be set, for this function to work.
	local pcmd = player.cmd
	local mem = bot.botmemory
	local water = (tails.eflags & MFE_UNDERWATER)
	local  flip = P_MobjFlip(tails)
	local _2d = (tails.flags2 & MF2_TWOD) or twodlevel
	local  scale = tails.scale
	local jump_last = (bot.lastbuttons & BT_JUMP)
	local spin_last = (bot.lastbuttons & BT_SPIN)

	local  dist = P_AproxDistance(sonic.x - tails.x, sonic.y - tails.y)
	local  zdist = flip * (sonic.z - tails.z)
	local  ang = sonic.angle
	local  pmom = P_AproxDistance(sonic.momx, sonic.momy)
	local  bmom = P_AproxDistance(tails.momx, tails.momy)
	local  followmax = 128 * 8 * scale -- Max follow distance before AI begins to enter catchup state
	local  followthres = 92 * scale -- Distance that AI will try to reach
	local  followmin = 32 * scale
	local  comfortheight = 96 * scale
	local  touchdist = 24 * scale
	local stalled = (bmom < scale >> 1) and dist > followthres -- Helps to see if the AI is having trouble catching up
	local samepos = (sonic.x == tails.x and sonic.y == tails.y)
	local blocked = bot.blocked
	
	if (not samepos)
		ang = R_PointToAngle2(tails.x, tails.y, sonic.x, sonic.y)
	end
	-- We can't follow Sonic if he's not aroundnot 
	if (not sonic or sonic.health <= 0)
	
		mem.thinkstate = AI_STANDBY
		return
	elseif (mem.thinkstate == AI_STANDBY)
		mem.thinkstate = AI_FOLLOW
	end
	if (tails.player.powers[pw_carry] == CR_MACESPIN or tails.player.powers[pw_carry] == CR_GENERIC)
	
		local isrelevant = (sonic.player.powers[pw_carry] == CR_MACESPIN or sonic.player.powers[pw_carry] == CR_GENERIC)
		if (sonic.player.cmd.buttons & BT_JUMP and (sonic.player.pflags & PF_JUMPED) and isrelevant)
			cmd.buttons = $ |  BT_JUMP
		end
		if (isrelevant)
		
			cmd.forwardmove = sonic.player.cmd.forwardmove
			cmd.angleturn = abs(tails.angle - sonic.angle)>>16
			if (sonic.angle < tails.angle)
				cmd.angleturn = -cmd.angleturn
			end
		elseif (dist > FixedMul(512*FRACUNIT, tails.scale))
			cmd.buttons = $ |  BT_JUMP
		end
		return
	end

	-- Adapted from CobaltBW's tails_AI.wad

	-- Check water
	if (water)
	
		followmin = 0
		followthres = 16*scale
		followmax = $ >>   1
		if (mem.thinkstate == AI_THINKFLY)
			mem.thinkstate = AI_FOLLOW
		end
	end
	-- Update catchup_tics
	if (mem.thinkstate == AI_SPINFOLLOW)
	
		mem.catchup_tics = 0
	elseif (dist > followmax or zdist > comfortheight or stalled)	
		mem.catchup_tics = min(mem.catchup_tics + 2, 70)
		if (mem.catchup_tics >= 70)
			mem.thinkstate = AI_CATCHUP
		end
	else
	
		mem.catchup_tics = max(mem.catchup_tics - 1, 0)
		if (mem.thinkstate == AI_CATCHUP)
			mem.thinkstate = AI_FOLLOW
		end
	end
	-- Orientation
	-- cmd.angleturn won't be relative to player angle, since we're not going through G_BuildTiccmd.
	if (bot.pflags & (PF_SPINNING|PF_STARTDASH))
	
		cmd.angleturn = (sonic.angle) >> 16 -- NOT FRACBITS DAMNIT
	
	elseif (mem.thinkstate == AI_FLYCARRY)
	
		cmd.angleturn = sonic.player.cmd.angleturn
	
	else
	
		cmd.angleturn = (ang) >> 16 -- NOT FRACBITS DAMNIT
	end

	-- ********
	-- FLY MODE
	-- exiting check
	if (player.exiting and mem.thinkstate == AI_THINKFLY)
		mem.thinkstate = AI_FOLLOW
	else
		-- Activate co-op flight
		if (mem.thinkstate == AI_THINKFLY and player.pflags & PF_JUMPED)
		
			if (not jump_last)
			
				jump = true
				mem.thinkstate = AI_FLYSTANDBY
				bot.pflags = $ |  PF_CANCARRY
			end
		end
		-- Check positioning
		-- Thinker for co-op flight
		if (not (water or pmom or bmom)
			and (dist < touchdist and not samepos)
			and not (pcmd.forwardmove or pcmd.sidemove or player.dashspeed)
			and P_IsObjectOnGround(sonic) and P_IsObjectOnGround(tails)
			and not (player.pflags & PF_STASIS)
			and bot.charability == CA_FLY)
				mem.thinkstate = AI_THINKFLY
		elseif (mem.thinkstate == AI_THINKFLY)
			mem.thinkstate = AI_FOLLOW
		end
		-- Set carried state
		if (player.powers[pw_carry] == CR_PLAYER and sonic.tracer == tails)
		
			mem.thinkstate = AI_FLYCARRY
		end
		-- Ready for takeoff
		if (mem.thinkstate == AI_FLYSTANDBY)
			if (zdist < -64*scale or (flip * tails.momz) > scale) -- Make sure we're not too high up
				spin = true
			elseif (not jump_last)
				jump = true
			end
			-- Abort if the player moves away or spins
			if (dist > followthres or player.dashspeed)
				mem.thinkstate = AI_FOLLOW
			end
		-- Read player inputs while carrying
		elseif (mem.thinkstate == AI_FLYCARRY)
			cmd.forwardmove = pcmd.forwardmove
			cmd.sidemove = pcmd.sidemove
			if (pcmd.buttons & BT_SPIN)
				spin = true
				jump = false
			elseif (not jump_last)
				jump = true
			end
			-- End flymode
			if (player.powers[pw_carry] != CR_PLAYER)
				mem.thinkstate = AI_FOLLOW
			end
		end
	end	
	if (P_IsObjectOnGround(tails) and not (pcmd.buttons & BT_JUMP) and (mem.thinkstate == AI_FLYSTANDBY or mem.thinkstate == AI_FLYCARRY))
		mem.thinkstate = AI_FOLLOW
	end
	-- ********
	-- SPINNING
	if (not (player.pflags & (PF_SPINNING|PF_STARTDASH)) and mem.thinkstate == AI_SPINFOLLOW)
		mem.thinkstate = AI_FOLLOW
	elseif (mem.thinkstate == AI_FOLLOW or mem.thinkstate == AI_SPINFOLLOW)
		if (not _2d)
			-- Spindash
			if (player.dashspeed)
				if (dist < followthres and dist > touchdist) -- Do positioning
					cmd.angleturn = (ang) >> 16 -- NOT FRACBITS DAMNIT
					cmd.forwardmove = 50
					mem.thinkstate = AI_SPINFOLLOW
				elseif (dist < touchdist)
					if (not bmom and (not (bot.pflags & PF_SPINNING) or (bot.dashspeed and bot.pflags & PF_SPINNING)))
						cmd.angleturn = (sonic.angle) >> 16 -- NOT FRACBITS DAMNIT
						spin = true
					end
					mem.thinkstate = AI_SPINFOLLOW
				else
					mem.thinkstate = AI_FOLLOW
				end
			-- Spin
			elseif (player.dashspeed == bot.dashspeed and player.pflags & PF_SPINNING)
				if (bot.pflags & PF_SPINNING or not spin_last)
					spin = true
					cmd.angleturn = (sonic.angle) >> 16 -- NOT FRACBITS DAMNIT
					cmd.forwardmove = 50
					mem.thinkstate = AI_SPINFOLLOW
				else
					mem.thinkstate = AI_FOLLOW
				end
			end
		-- 2D mode
		else
			if (((player.dashspeed and not bmom) or (player.dashspeed == bot.dashspeed and (player.pflags & PF_SPINNING)))
				and ((bot.pflags & PF_SPINNING) or not spin_last))
				spin = true
				mem.thinkstate = AI_SPINFOLLOW
			else
				mem.thinkstate = AI_FOLLOW
			end
		end
	end
	-- ********
	-- FOLLOW
	if (mem.thinkstate == AI_FOLLOW or mem.thinkstate == AI_CATCHUP)
		-- Too far
		if (mem.thinkstate == AI_CATCHUP or dist > followthres)
			if (not _2d)
				cmd.forwardmove = 50
			elseif (sonic.x > tails.x)
				cmd.sidemove = 50
			else
				cmd.sidemove = -50
			end
		-- Within threshold
		elseif (dist > followmin and abs(zdist) < 192*scale)
			if (not _2d)
				cmd.forwardmove = FixedHypot(pcmd.forwardmove, pcmd.sidemove)
			else
				cmd.sidemove = pcmd.sidemove
			end
		-- Below min
		elseif (dist < followmin)
			-- Copy inputs
			cmd.angleturn = (sonic.angle) >> 16 -- NOT FRACBITS DAMNIT
			cmd.forwardmove = 8 * pcmd.forwardmove / 10
			cmd.sidemove = 8 * pcmd.sidemove / 10
		end
	end
	-- ********
	-- JUMP
	if (mem.thinkstate == AI_FOLLOW or mem.thinkstate == AI_CATCHUP or (mem.thinkstate == AI_SPINFOLLOW and player.pflags & PF_JUMPED))
		-- Flying catch-up
		if (bot.pflags & PF_THOKKED)
			cmd.forwardmove = min(50, (dist/scale)>>3)
			if (zdist < -64*scale)
				spin = true
			elseif (zdist > 0 and not jump_last)
				jump = true
			end
		end
		-- Just landed
		if (tails.eflags & MFE_JUSTHITFLOOR)
			jump = false
		-- Start jump
		elseif (not jump_last and not (bot.pflags & PF_JUMPED) --and not (player.pflags & PF_SPINNING)
			and ((zdist > 32*scale and player.pflags & PF_JUMPED) -- Following
				or (zdist > 64*scale and mem.thinkstate == AI_CATCHUP) -- Vertical catch-up
				or (stalled and mem.catchup_tics > 20 and bot.powers[pw_carry] == CR_NONE)
				--or (bmom < scale>>3 and dist > followthres and not (bot.powers[pw_carry])) -- Stopped & not in carry state
				or (bot.pflags & PF_SPINNING and not (bot.pflags & PF_JUMPED)))) -- Spinning
					jump = true
		-- Hold jump
		elseif (bot.pflags & PF_JUMPED and jump_last and tails.momz*flip > 0 and (zdist > 0 or mem.thinkstate == AI_CATCHUP))
			jump = true
		-- Start flying
		elseif (bot.pflags & PF_JUMPED and mem.thinkstate == AI_CATCHUP and not jump_last and bot.charability == CA_FLY)
			jump = true
		end
	end
	-- ********
	-- HISTORY
	--jump_last = jump
	--spin_last = spin

	-- Turn the virtual keypresses into ticcmd_t.
	cmd.buttons = $ | (jump and BT_JUMP or 0) | (spin and BT_SPIN or 0)

	-- Update our status
	mem.lastForward = forward
	mem.lastBlocked = blocked
end)

-- Ordinarily, the flight indicator would be handled automatically. However, it relies on player.botmem (which we don't have access to) so we'll also have to reproduce this function in order for the script to be complete.
rawset(_G, "B_HandleFlightIndicator", function(player)
	makeBotMemory(player) -- bot.botmemory needs to be set, for this function to work.

	local tails = player.mo
	local mem = player.botmemory
	local shouldExist = false

	if (not tails)
		return
	end
	
	shouldExist = (mem.thinkstate == AI_THINKFLY) and player.botleader
		and player.bot == BOT_2PAI and player.playerstate == PST_LIVE

	-- check whether the indicator doesn't exist
	if not(tails._hnext and tails._hnext.valid)
		-- if it shouldn't exist, everything is fine
		if (not shouldExist)
			return
		end
		-- otherwise, spawn it
		tails._hnext = P_SpawnMobjFromMobj(tails, 0, 0, 0, MT_OVERLAY)
		tails._hnext.target = tails
		tails._hnext._hprev = tails
		tails._hnext.state = S_FLIGHTINDICATOR
	end

	-- if the mobj isn't a flight indicator, let's not mess with it
	if (tails._hnext.type != MT_OVERLAY or (tails._hnext.state != S_FLIGHTINDICATOR))
		return
	end
	-- if it shouldn't exist, remove it
	if (not shouldExist)
		P_RemoveMobj(tails._hnext)
		tails._hnext = nil
		return
	end

	-- otherwise, update its visibility
	if player.botleader == displayplayer
		tails._hnext.flags2 = $ & ~MF2_DONTDRAW
	else
		tails._hnext.flags2 = $ | MF2_DONTDRAW
	end
end)