// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
use std::{sync::Arc, collections::BTreeMap, time::Duration};

use tokio::{sync::{mpsc::{self, Sender as MpscSender, Receiver as MpscReceiver}, oneshot::{self, Sender as OneshotSender}}, time};

use super::{SubscribeMessage, Key, status::{PublishingStatus, ConnectionStatus}, super::Server};

use crate::utils::ToStr;
use crate::net::{serp::{InitClient, Sender, request::{Request, PublishRequest}}, utils::{IpCategory, quic::{self, IdleConfig}, hostname::Hostname, security::Security}};
use crate::protocol::discovery::{DiscoveryServer, GameId};

pub struct PublishingServer(MpscSender<Message>);

pub struct PublishingArgs {
	pub host: Box<str>,
	pub port: u16,
	pub return_port: u16,
	pub hostname: Option<Hostname>,
}

enum Message {
	/// The expected key of any Publish Subscribe requests.
	ExpectedKey(Key),

	/// Publish Init response has arrived.
	InitResponse(Result<(), String>),

	/// Publish Subscribe request received, the main task can either accept the
	/// connection or reject it if there already is an active connection.
	SubscribeInit(OneshotSender<Result<(), SubscribeError>>, Key),

	/// Publish-subscribe connection ends, the main task is notified so it can
	/// attempt to reconnect.
	SubscribeFinish,

	/// When the server shuts down, the endpoint closes all connections which
	/// results in any current publish-subscribe connection finishing. To prevent
	/// this from triggering another Publish Init request, and hence slowing down
	/// termination while waiting for that connection to finish, I created this
	/// variant to stop the main task.
	Shutdown,

	/// Gets the status for the `status` command.
	Status(OneshotSender<PublishingStatus>),
}

#[derive(Clone, Copy)]
pub enum SubscribeError {
	ConnectionAlreadyExists,
	IncorrectKey,
	TaskFinished, // Don't think this will ever occur
}

enum State {
	Connecting(ConnectingState),
	Connected,
}

struct ConnectingState {
	timeout: u64,
	waiting_for_pub_init: bool,
}

impl State {
	fn new() -> State {
		State::Connecting(ConnectingState {
			timeout: 0,
			waiting_for_pub_init: false,
		})
	}

	fn connected(&self) -> bool {
		matches!(self, State::Connected)
	}

	fn pub_init(&mut self, init_factory: &Arc<InitFactory>) {
		match self {
			State::Connecting(state) => {
				if !state.waiting_for_pub_init {
					Arc::clone(init_factory).send(state.timeout);
					state.timeout = if state.timeout == 0 { 30 } else { state.timeout * 2 }; // Exponential backoff
					state.waiting_for_pub_init = true;
				} else {
					log::warn!("attempted to send Publish Init while waiting for Publish Init");
				}
			},
			State::Connected => log::warn!("attempted to send Publish Init while connection already exists"),
		}
	}

	fn pub_init_received(&mut self) {
		match self {
			State::Connecting(state) => state.waiting_for_pub_init = false,
			State::Connected => (),
		}
	}
}

impl PublishingServer {
	pub fn new(args: PublishingArgs) -> PublishingServer {
		let (sender, receiver) = mpsc::channel(64);
		tokio::spawn(Self::main_task(receiver, Arc::new(InitFactory { sender: sender.clone(), args })));
		PublishingServer(sender)
	}

	async fn main_task(mut receiver: MpscReceiver<Message>, init_factory: Arc<InitFactory>) {
		let mut state = State::new();
		let mut expected_key = None;

		state.pub_init(&init_factory);

		while let Some(msg) = receiver.recv().await {
			match msg {
				Message::ExpectedKey(key) => expected_key = Some(key),
				Message::InitResponse(res) => {
					state.pub_init_received();

					match (res, &state) {
						(Ok(()), State::Connected) => log::info!("Successfully initiated publication!"),
						(Err(err), State::Connecting(_)) => {
							log::info!("failed initiating publication, retrying: {err}");
							state.pub_init(&init_factory);
						},

						// Unlikely for these to happen
						(Ok(()), State::Connecting(_)) => {
							log::warn!("publish-init was successful despite no connection, retrying (maybe it closed)");
							state.pub_init(&init_factory);
						},
						(Err(err), State::Connected) => log::warn!("publication succeeded but publish-init had error: {err}"),
					}
				},
				Message::SubscribeInit(sender, key) => {
					let response =
						if Some(key) != expected_key { Err(SubscribeError::IncorrectKey) }
						else if state.connected() { Err(SubscribeError::ConnectionAlreadyExists) }
						else { Ok(()) };

					let _ = sender.send(response);
					if response.is_ok() {
						state = State::Connected;
					}
				},
				Message::SubscribeFinish => {
					if state.connected() {
						state = State::new();
						state.pub_init(&init_factory);
					} else {
						debug_assert!(false); // Should never happen
					}
				},
				Message::Shutdown => break,
				Message::Status(sender) => {
					let _ = sender.send(PublishingStatus {
						host: init_factory.args.host.clone(),
						hostname: init_factory.args.hostname.clone(),
						connection: match &state {
							State::Connecting(_) => ConnectionStatus::Connecting,
							State::Connected => ConnectionStatus::Connected,
						},
					});
				},
			}
		}
	}

	pub async fn respond_publish_subscribe(&self, key: Key) -> Result<Subscriber, SubscribeError> {
		let (sender, receiver) = oneshot::channel();
		let _ = self.0.send(Message::SubscribeInit(sender, key)).await;
		let can_subscribe = receiver.await.map_err(|_| SubscribeError::TaskFinished)?;
		can_subscribe?;
		Ok(Subscriber(self.0.clone()))
	}

	pub fn shutdown(&self) {
		let _ = self.0.blocking_send(Message::Shutdown);
	}

	pub fn get_status(&self) -> Option<PublishingStatus> {
		let (sender, receiver) = oneshot::channel();
		self.0.blocking_send(Message::Status(sender)).ok()?;
		receiver.blocking_recv().ok()
	}
}

struct InitFactory {
	sender: MpscSender<Message>,
	args: PublishingArgs,
}

impl InitFactory {
	fn send(self: Arc<Self>, timeout: u64) {
		tokio::spawn(async move {
			if timeout > 0 {
				log::info!("Waiting for {timeout} seconds before initiating publication...");
				time::sleep(Duration::from_secs(timeout)).await;
			}

			log::info!("Initiating publication!");
			let res = self.make_request().await;
			let _ = self.sender.send(Message::InitResponse(res)).await;
		});
	}

	async fn make_request(&self) -> Result<(), String> {
		/*
		 * Using a larger idle timeout to allow the publication server to wait a
		 * reasonable amount of time (30 seconds) for Publish Subscribe responses
		 * without this connection timing out.
		 */
		let (conn, security) = quic::connect(quic::client_endpoint()?, &self.args.host, self.args.port, IdleConfig::Larger).await?;
		match security {
			Security::Secure => (),
			Security::Local | Security::Insecure => log::warn!("insecure connection with publication server (security = {})", security.to_str()),
		}

		let key = Key::try_new()?;
		let _ = self.sender.send(Message::ExpectedKey(key.clone())).await;
		InitClient::send(conn, Request::Publish(PublishRequest::Init(self.args.return_port, self.args.hostname.clone(), key))).await.map_err(|err| err.to_string())?;
		Ok(())
	}
}

pub struct Subscriber(MpscSender<Message>);

impl Subscriber {
	pub async fn run(self, mut sender: Sender, server: Arc<Server>, ip_category: IpCategory) -> Result<(), String> {
		let (notify, init_servers) = {
			let router = server.router.read().unwrap();
			let notify = router.get_notify();
			let init_servers: Vec<DiscoveryServer> = router.discovery_servers(ip_category).collect();
			drop(router);
			(notify, init_servers)
		};

		let mut servers: BTreeMap<GameId, DiscoveryServer> = init_servers.iter().map(|server| (server.game_id.clone(), server.clone())).collect();

		sender.send(&vec![SubscribeMessage::Init(init_servers)]).await?;

		loop {
			tokio::select! {
				() = notify.notified() => {
					let new_servers: BTreeMap<GameId, DiscoveryServer> = server.router.read().unwrap()
						.discovery_servers(ip_category)
						.map(|server| (server.game_id.clone(), server))
						.collect();

					let msg = Subscriber::diff_messages(&new_servers, &mut servers);
					if !msg.is_empty() {
						sender.send(&msg).await?;
					}
					servers = new_servers;
				},
				() = sender.closed() => return Ok(()),
			}
		}
	}

	fn diff_messages(new: &BTreeMap<GameId, DiscoveryServer>, old: &mut BTreeMap<GameId, DiscoveryServer>) -> Vec<SubscribeMessage> {
		let mut messages = Vec::new();

		// Removed
		for id in old.keys() {
			if !new.contains_key(id) {
				messages.push(SubscribeMessage::Remove(id.clone()));
			}
		}

		// Added or changed
		for (id, new_server) in new {
			if let Some(old_server) = old.get_mut(id) {
				let (old_players, old_spectators) = (old_server.player_count, old_server.spectator_count);
				(old_server.player_count, old_server.spectator_count) = (new_server.player_count, new_server.spectator_count);

				if old_server != new_server {
					// Something changed that isn't the client counts, so send the entire server with an `Add` message
					messages.push(SubscribeMessage::Add(new_server.clone()));
				} else if (old_players, old_spectators) != (new_server.player_count, new_server.spectator_count) {
					// Equal up to client count, but these client counts differ, so send a more efficient message
					messages.push(SubscribeMessage::UpdateClientCount(id.clone(), new_server.player_count, new_server.spectator_count));
				}
			} else {
				messages.push(SubscribeMessage::Add(new_server.clone()));
			}
		}

		messages
	}
}

impl Drop for Subscriber {
	fn drop(&mut self) {
		let sender = self.0.clone();
		tokio::spawn(async move { sender.send(Message::SubscribeFinish).await });
	}
}
