// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
mod player;
pub mod forcefield;

pub use player::LocalPlayer;

use std::{num::NonZeroUsize, collections::{BTreeMap, VecDeque, btree_map::Entry}};

use glam::Vec2;

use player::DPlayer;

use super::{forcefield::RenderedForcefield, super::{Camera, HealthBar, HealthBarRenderer, WorldLabel, Sounds, Sfx, LoopingSounds}};

use crate::app::{config::action::{Action, ActionEvent, ActionState}, playing::{model::{Model, client::ClientError, changes::ClientChanges}}};
use crate::world::{
	World, DELTA_TIME, UPDATES_PER_SECOND, MAX_STEPS, player::{PlayerId,
	update::{PlayerUpdate, PlayerBulkUpdate}}, bullet::{Bullet, BulletManager},
	effects::Effect, team::TeamId,
};
use crate::protocol::message::Message;
use crate::utils::{RollingAverage, maths::{Circle, decay::Decay}};

pub struct SmoothWorld {
	current_id: PlayerId,

	/**
	 * Initialised to None and is set to Some(team_id) when the current player
	 * joins a team.
	 *
	 * This is used to remember the team the player was in while spectating so
	 * that team's scores can be indented.
	 */
	last_current_team: Option<TeamId>,
	followed_id: Option<PlayerId>, // INVARIANT: If followed_id == Some(x), then x ∈ world
	prev_followed_id: PlayerId, // Used when restoring to a followed player
	players: BTreeMap<PlayerId, LocalPlayer>,
	bullets: BulletManager,

	/**
	 * Updates that have been performed locally that the server hasn't yet
	 * acknowledged.
	 *
	 * Storing this allows the player to receive immediate feedback on moving
	 * around even if there's large ping.
	 */
	staging_updates: VecDeque<PlayerUpdate>,
	past_staging_updates: RollingAverage,

	/**
	 * The number of updates that have been sent out. This is the sequence number
	 * of the update at the back of the staging update queue.
	 */
	update_count: u64,

	/**
	 * When the client is too far ahead or behind the server, don't immediately
	 * jump the smooth world in time as that looks bad. Instead, keep the error
	 * but exponentially decay it towards zero so things are smoother.
	 */
	time_shift: f32,
}

impl SmoothWorld {
	pub fn new() -> SmoothWorld {
		SmoothWorld {
			current_id: PlayerId::default(),
			last_current_team: None,
			followed_id: None,
			prev_followed_id: PlayerId::default(),
			players: BTreeMap::new(),
			bullets: BulletManager::new(),
			staging_updates: VecDeque::new(),
			past_staging_updates: RollingAverage::new(const { NonZeroUsize::new(256).unwrap() }),
			update_count: 0,
			time_shift: 0.0,
		}
	}

	pub fn reset(&mut self, model: &Model) {
		let world = model.get_world();
		self.current_id = model.get_id();
		self.last_current_team = world.get_player_team(self.current_id);
		self.followed_id = world.has_player(self.current_id).then_some(self.current_id);
		self.prev_followed_id = PlayerId::default();
		self.players = world.get_players().iter().map(|(&id, player)| (id, LocalPlayer::new(id, player.clone(), world.get_effects()))).collect();
		self.bullets = world.get_bullets().clone();
		self.staging_updates.clear();
		self.past_staging_updates.reset();
		self.update_count = 0;
		self.time_shift = 0.0;
	}

	pub fn get_total_time_shift(&self) -> f32 {
		self.staging_updates.len() as f32 * DELTA_TIME + self.time_shift
	}

	pub fn get_player(&self, id: PlayerId) -> Option<&LocalPlayer> {
		self.players.get(&id)
	}

	pub fn player_iter(&self) -> impl Iterator<Item = &LocalPlayer> {
		self.players.values()
	}

	pub fn player_iter_mut(&mut self) -> impl Iterator<Item = (PlayerId, &mut LocalPlayer)> {
		self.players.iter_mut().map(|(&id, player)| (id, player))
	}

	pub fn followed_player_id(&self) -> Option<PlayerId> {
		self.followed_id
	}

	pub fn followed_player(&self) -> Option<(PlayerId, &LocalPlayer)> {
		self.followed_id.map(|id| (id, &self.players[&id]))
	}

	pub fn followed_team(&self) -> Option<TeamId> {
		self.followed_id.map_or(self.last_current_team, |id| self.players[&id].inner().get_team())
	}

	pub fn current_player_mut(&mut self) -> Option<&mut LocalPlayer> {
		self.players.get_mut(&self.current_id)
	}

	pub fn get_bullets(&self) -> &[Bullet] { &self.bullets }

	pub fn get_forcefields(&self, time_diff: f32) -> Vec<RenderedForcefield> {
		let time_shift = self.get_total_time_shift() + time_diff;

		self.players
			.iter()
			.filter_map(|(id, player)| player.ffield.clone().map(|ffield| RenderedForcefield {
				inner: ffield.shift(time_shift),
				pos: player.get_pos(time_diff),
				player_id: *id,
				team_id: player.player.get_team(),
			}))
			.collect()
	}

	pub fn create_bulk_update(&self) -> Message {
		// update_count is the sequence number of the end update
		// Instead want that for the beginning update
		let sequence_number = self.oldest_sequence_number();
		let mut bulk_update = PlayerBulkUpdate::new(sequence_number);
		for &update in &self.staging_updates {
			bulk_update.push(update);
		}
		Message::PlayerBulkUpdate(bulk_update)
	}

	fn oldest_sequence_number(&self) -> u64 {
		self.update_count - (self.staging_updates.len() - 1) as u64
	}

	/**
	 * Returns the next valid player id that is greater than or equal to the
	 * input id, wrapping around if nothing's found or returning None if there
	 * are no ids.
	 */
	fn valid_id_at_or_after(&self, id: PlayerId) -> Option<PlayerId> {
		self.players.range(id..).next().map(|(&id, _)| id).or_else(|| self.players.first_key_value().map(|e| *e.0))
	}

	/**
	 * Performs an action related to changing the followed player while
	 * spectating. Returns the position of the newly followed player, if any.
	 */
	pub fn action(&mut self, action: ActionEvent) -> Option<Vec2> {
		// Only when spectating
		if self.players.contains_key(&self.current_id) { return None; }

		let prev_followed = self.followed_id;

		match action {
			ActionEvent::Action(Action::FollowPlayer, ActionState::Pressed) => {
				self.followed_id = if self.followed_id.is_some() {
					None
				} else {
					self.valid_id_at_or_after(self.prev_followed_id)
				};
			},
			ActionEvent::Action(Action::PrevPlayer, ActionState::Pressed) => {
				if let Some(id) = self.followed_id {
					self.followed_id = self.players.range(..id).next_back().map(|(&id, _)| id).or_else(|| self.players.last_key_value().map(|e| *e.0));
				}
			},
			ActionEvent::Action(Action::NextPlayer, ActionState::Pressed) => {
				if let Some(id) = self.followed_id {
					self.followed_id = self.valid_id_at_or_after(id.next());
				}
			},
			ActionEvent::Action(..) | ActionEvent::AllReleased => (),
		}

		if let Some(id) = self.followed_id {
			self.prev_followed_id = id;
		}

		self.followed_player().and_then(|(id, player)| (prev_followed != Some(id)).then_some(player.get_pos(0.0 /* Not perfect but I'm lazy */)))
	}

	pub fn remove_old_players(&mut self, world: &World) {
		self.players.retain(|&id, _| world.has_player(id));
	}

	/**
	 * Returns whether the camera position should be updated if the followed id
	 * changes.
	 */
	#[must_use]
	pub fn update_followed_player(&mut self) -> bool {
		if self.players.contains_key(&self.current_id) {
			if self.followed_id != Some(self.current_id) {
				self.followed_id = Some(self.current_id); // Follows the current player when they join
				return true; // Only moves the camera if this has changed
			}
		} else if let Some(id) = self.followed_id {
			if id == self.current_id {
				self.followed_id = None; // When a player enters spectating mode, don't have them follow any player
			} else if !self.players.contains_key(&id) {
				self.followed_id = self.valid_id_at_or_after(id.next()); // Satisfies the invariant if the player suddenly leaves
				if let Some(id) = self.followed_id {
					self.prev_followed_id = id;
				}
				return true;
			}
		}

		false
	}

	pub fn sync(&mut self, world: &World, changes: &ClientChanges, ack: u64, recent_update_count: usize, sounds: &mut Sounds) {
		let to_keep = self.update_count.checked_sub(ack).unwrap_or_else(|| {
			log::warn!("server acknowledged update {ack} which hasn't yet been sent (sequence number of most recent update = {})", self.update_count);
			0
		});

		let to_remove = self.staging_updates.len().checked_sub(to_keep as usize).unwrap_or_else(|| {
			log::warn!(
				"server acknowledged update {ack} which is earlier than the oldest staging update ({}), staging update count = {}, most recent update's sequence number = {}",
				self.oldest_sequence_number(),
				self.staging_updates.len(),
				self.update_count
			);

			0
		});

		/*
		 * When the game starts, there can be a pretty big lag spike and that can
		 * result in a pretty large negative time shift, which looks much worse
		 * than any initial jitter that might occur.
		 */
		if self.update_count >= UPDATES_PER_SECOND as u64 {
			self.time_shift += (to_remove as f32 - recent_update_count as f32) * DELTA_TIME;
		}

		for _ in 0..to_remove { self.staging_updates.pop_front(); }

		self.bullets = world.get_bullets().clone();

		/*
		 * Extrapolates the current world's state by the length of the staging
		 * update queue.
		 *
		 * For the current player this is applying the staging updates, and for
		 * all other players this is using the most recent update, as an
		 * estimation that likely will be the next one.
		 *
		 * Also updating the bullets into the future while doing this.
		 */
		let mut new_players: BTreeMap<PlayerId, LocalPlayer> = world.get_players().iter().map(|(id, player)| (*id, LocalPlayer::new(*id, player.clone(), world.get_effects()))).collect();

		let mut sync_time = 0.0;
		for &update in &self.staging_updates {
			for (&id, player) in &mut new_players {
				if id == self.current_id { // Use the most recent update (which was already set) as a prediction if not the current player
					player.player.set_update(update);
				}

				if let Some(bullet) = player.update(world, id) {
					self.bullets.add(bullet);
				}
			}

			SmoothWorld::update_bullets_sync(&mut self.bullets, DELTA_TIME, world, new_players.iter_mut(), &mut sync_time, changes, sounds);
		}
		SmoothWorld::update_bullets_sync(&mut self.bullets, self.time_shift, world, new_players.iter_mut(), &mut sync_time, changes, sounds);

		for (id, player) in new_players {
			match self.players.entry(id) {
				Entry::Vacant(e) => _ = e.insert(player),
				Entry::Occupied(mut e) => {
					let prev_player = e.get_mut();

					if !changes.moved.contains(&id) {
						let delta = &prev_player.player - &player.player;
						prev_player.delta += &delta;
						if id == self.current_id {
							prev_player.delta.reset_angle();
						}
					} else {
						prev_player.delta = DPlayer::zero();
					}

					prev_player.player = player.player;
				},
			}
		}

		for hit in &changes.hits {
			let Some(id) = hit.player_id else { continue; };
			if let Some(player) = self.players.get_mut(&id) { // Shouldn't fail but being safe
				player.last_hurt = if hit.killed { f32::INFINITY } else { 0.0 };
				if hit.killed {
					player.hbar.reset();
				}
			}
		}

		for id in &changes.reset {
			if let Some(player) = self.players.get_mut(id) {
				player.hbar.reset();
			}
		}

		self.last_current_team = changes.current_team_change.or(self.last_current_team);
	}

	pub fn update(&mut self, update: PlayerUpdate, world: &World, sounds: &mut Sounds) -> Result<(), ClientError> {
		if let Some(player) = self.players.get_mut(&self.current_id) {
			player.player.set_update(update);
		}

		self.staging_updates.push_back(update);
		self.update_count += 1;

		if self.staging_updates.len() >= MAX_STEPS {
			return Err(ClientError::Limit(String::from("too many staging updates accumulated by client")));
		}

		self.past_staging_updates.add(self.staging_updates.len() as f32);
		let mean_len = self.past_staging_updates.mean().unwrap();
		let decay_rate = (100.0 / (mean_len + 1.0)) + 1.0;
		let prev_time_shift = self.time_shift;
		self.time_shift.decay(2.0, DELTA_TIME);

		for (&id, player) in &mut self.players {
			player.delta.decay(decay_rate);
			if let Some(mut bullet) = player.update(world, id) {
				bullet.update(prev_time_shift, &world.get_forcefields());
				self.bullets.add(bullet);

				/*
				 * Doesn't include any fake bullets (bullets the SmoothWorld
				 * predicts would be fired). This is acceptable to avoid the
				 * combined sound of multiple bullets being fired from being too
				 * loud if the ping is large.
				 */
				sounds.play_spatial(Sfx::Shot(world.get_effects().get_mul(id, Effect::Reload)), player.player.get_pos(), player.player.get_vel());
			}
		}

		/*
		 * NOTE: In the future I can investigate adding the time shift to the
		 * players similar to what I do with the bullets. If done properly it
		 * could make things a bit smoother on bad network conditions.
		 */
		for bullet in SmoothWorld::update_bullets(&mut self.bullets, DELTA_TIME - prev_time_shift + self.time_shift, world, self.players.iter_mut()) {
			sounds.play_spatial(Sfx::BlockBullet, bullet.get_pos(), Vec2::ZERO);
		}

		Ok(())
	}

	#[allow(clippy::too_many_arguments)]
	fn update_bullets_sync<'a>(
		bullets: &mut BulletManager, dt: f32, world: &World, players: impl Iterator<Item = (&'a PlayerId, &'a mut LocalPlayer)>,
		sync_time: &mut f32, changes: &ClientChanges, sounds: &mut Sounds,
	) {
		*sync_time += DELTA_TIME;
		for bullet in SmoothWorld::update_bullets(bullets, dt, world, players) {
			let bullet_time = bullet.get_time();
			if changes.shot_not_predicted(&bullet) && bullet_time > *sync_time && bullet_time <= *sync_time + changes.update_time {
				sounds.play_spatial(Sfx::BlockBullet, bullet.get_pos(), Vec2::ZERO);
			}
		}
	}

	fn update_bullets<'a>(bullets: &mut BulletManager, dt: f32, world: &World, players: impl Iterator<Item = (&'a PlayerId, &'a mut LocalPlayer)>) -> Vec<Bullet> {
		bullets.update(dt, &world.get_forcefields());
		for (&id, player) in players {
			for bullet in bullets.collide_player(&player.player, id, world.get_config().friendly_fire) {
				player.player.on_bullet_hit(&bullet);
			}
		}
		bullets.collide_blocks(world.get_blocks())
	}

	pub fn push_health_bars_and_labels(&self, hbar_renderer: &mut HealthBarRenderer, lwinsize: Vec2, camera: &Camera, time_diff: f32, labels: &mut Vec<WorldLabel>) {
		// Render this player's health bar on top of all others
		labels.reserve(self.players.len());
		let builder = HealthBar::above_player_builder(lwinsize, camera);
		for (&id, player) in &self.players {
			if Some(id) != self.followed_id {
				let (hbar, name_pos) = builder.make(player, &player.hbar, time_diff);
				hbar_renderer.push(hbar);
				if let Some(label) = WorldLabel::new_player(name_pos, player.player.get_style()) {
					labels.push(label);
				}
			}
		}

		if let Some(id) = self.followed_id {
			hbar_renderer.push(self.players[&id].hbar.make_hud(lwinsize));
		}
	}

	pub fn update_sounds(&mut self, dt: f32, time_diff: f32, looping_sounds: &mut LoopingSounds) {
		for player in self.players.values_mut() {
			player.play_looping_sounds(dt, time_diff, looping_sounds);
		}
	}
}
