// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
mod router;
pub mod client_count_notify;
pub(super) mod publish;
mod repl;
mod ip_block;
mod rate_limiter;
mod privacy;
pub mod init;
#[cfg(feature = "client")] pub mod init_gui;

#[cfg(feature = "client")] pub use router::StreamSender;

use std::{net::{IpAddr, SocketAddr, UdpSocket}, sync::{Arc, Mutex, RwLock}, fmt::Write};

use tokio::{fs, task};
use quinn::{Endpoint, Incoming, ServerConfig, EndpointConfig, VarInt};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, pem::PemObject};
use sysinfo::Networks;
#[cfg(feature = "client")] use sysinfo::NetworkData;

use rate_limiter::RateLimiter;
use repl::Repl;
use router::{Router, RouteError};
use publish::{PublishingServer, PublishingArgs, PublicationServer, InitError as PublishInitError, SubscribeError, Key, status::ConnectionStatus};

use crate::ServerArgs as ServerCliArgs;
use crate::net::{
	serp::{InitServer, PendingResponse as RawPendingResponse, Connection, Error as SerpError, request::{Request, PublishRequest}},
	utils::{WaitIdleTimeout, SocketAddrToCanonical, IpClassify, quic::{self, IdleConfig}, hostname::Hostname},
};
use crate::protocol::{info::Info, discovery::DiscoveryMessage, message::stream::networked::IncomingServerMessageStream};
use crate::utils::{INDENTATION, or::Or};

pub const PORT: u16 = 4080; // UDP
pub const SHUTTING_DOWN_MSG: &str = "server shutting down";
pub const SHUTTING_DOWN_STATUS: u32 = 0;

pub struct Server {
	endpoint: Endpoint,
	rate_limiter: Arc<Mutex<RateLimiter>>,
	router: RwLock<Router>,
	info: Info,
	networks: Networks,
	publishing_server: Option<PublishingServer>,
	publication_server: Option<PublicationServer>,
}

pub struct ServerArgs {
	cert_info: Option<CertInfo>,
	publishing: Option<PublishingArgs>,
	publication: bool,
}

struct CertInfo {
	cert_path: Box<str>,
	key_path: Box<str>,
}

impl ServerArgs {
	#[cfg(feature = "client")]
	pub fn new_gui_server() -> ServerArgs {
		ServerArgs { cert_info: None, publishing: None, publication: false }
	}

	pub fn try_from_command_line(args: &ServerCliArgs) -> Result<ServerArgs, String> {
		let hostname = args.server_name.clone().map(Hostname::try_from).transpose().map_err(|err| format!("invalid hostname: {err}"))?;

		Ok(ServerArgs {
			cert_info: args.cert_path.as_ref().zip(args.key_path.as_ref()).map(|c| CertInfo { cert_path: c.0.clone(), key_path: c.1.clone() }),
			publishing: args.publish_addr.clone().map(|host| PublishingArgs { host, port: args.publish_port, return_port: args.port, hostname }),
			publication: args.accept_publications,
		})
	}
}

struct PendingResponse {
	inner: RawPendingResponse,
	addr: SocketAddr,
}

impl PendingResponse {
	async fn success(self) -> Result<Connection, SerpError> {
		log::info!("sending success response, addr = {}", self.addr);
		self.inner.success().await
	}

	async fn success_unflushed(self) -> Result<Connection, SerpError> {
		log::info!("sending success response, addr = {}", self.addr);
		self.inner.success_unflushed().await
	}

	async fn understood(self, msg: &str) {
		log::info!("sending understood response, addr = {}: {msg}", self.addr);
		self.inner.understood(msg).await;
	}
}

impl Server {
	pub async fn try_new(socket: UdpSocket, router: Router, args: ServerArgs) -> Result<Arc<Server>, String> {
		// Loads certificate information from the provided paths, otherwise using a self-signed certificate
		let (certs, key) = if let Some(info) = args.cert_info {
			Server::load_cert_info(info).await.map_err(|err| format!("failed loading certificate info: {err}"))?
		} else {
			let cert = rcgen::generate_simple_self_signed(vec![String::from("localhost")]).map_err(|err| format!("cannot generate self-signed certificate: {err}"))?;
			let key = PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der());
			(vec![cert.cert.into()], key.into())
		};
		let mut config = ServerConfig::with_single_cert(certs, key).map_err(|err| format!("cannot create server config: {err}"))?;
		let idle_config =
			if args.publishing.is_some() { IdleConfig::LargerAndKeepAlive } // Prevents incoming Publish Subscribe connections from timing out when there's no activity
			else if args.publication { IdleConfig::Larger } // Prevents incoming Publish Init requests from timing out too quickly for a SERP response to be sent saying that the server can't connect back
			else { IdleConfig::Default };
		config.transport_config(quic::transport_config(idle_config));

		let runtime = quinn::default_runtime().ok_or_else(|| String::from("failed getting tokio runtime"))?;
		let endpoint = Endpoint::new(EndpointConfig::default(), Some(config), socket, runtime).map_err(|err| format!("cannot create QUIC server endpoint: {err}"))?;

		Ok(Arc::new(Server {
			endpoint,
			rate_limiter: Arc::new(Mutex::new(RateLimiter::new())),
			router: RwLock::new(router),
			info: Info::new(),
			networks: task::block_in_place(Networks::new_with_refreshed_list),
			publishing_server: args.publishing.map(PublishingServer::new),
			publication_server: args.publication.then(PublicationServer::new),
		}))
	}

	// Copied and modified from https://github.com/quinn-rs/quinn/blob/main/quinn/examples/server.rs
	async fn load_cert_info(info: CertInfo) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), String> {
		let cert_data = fs::read(info.cert_path.as_ref()).await.map_err(|err| format!("failed loading certificate chain at \"{}\": {err}", info.cert_path))?;
		let certs = CertificateDer::pem_reader_iter(&mut &*cert_data)
			.collect::<Result<_, _>>()
			.map_err(|err| format!("failed parsing certificates, make sure they're in PEM format: {err}"))?;

		let key_data = fs::read(info.key_path.as_ref()).await.map_err(|err| format!("failed loading private key at \"{}\": {err}", info.key_path))?;
		let key = match PrivateKeyDer::from_pem_reader(&mut &*key_data) {
			Ok(key) => key,
			Err(err) => return Err(format!("failed parsing private key, make sure it's in the PEM format: {err}")),
		};

		Ok((certs, key))
	}

	pub async fn run(self: Arc<Self>) {
		while let Some(incoming) = self.endpoint.accept().await {
			let endpoint = Arc::clone(&self);
			tokio::spawn(endpoint.handle_incoming(incoming));
		}

		self.endpoint.wait_idle_timeout().await;
	}

	pub fn close(&self) {
		if let Some(server) = &self.publishing_server {
			server.shutdown();
		}

		self.router.write().unwrap().clear();
		self.endpoint.close(VarInt::from_u32(SHUTTING_DOWN_STATUS), SHUTTING_DOWN_MSG.as_bytes());
	}

	fn get_publish_status(&self) -> String {
		let mut status = String::new();

		if let Some(info) = self.publishing_server.as_ref().and_then(PublishingServer::get_status) {
			let _ = write!(status, "Publishing to {}", info.host);
			match info.connection {
				ConnectionStatus::Connecting => _ = write!(status, " (connecting)"),
				ConnectionStatus::Connected => (),
			}
			if let Some(hostname) = &info.hostname {
				let _ = write!(status, " with hostname of {hostname}");
			}
			let _ = writeln!(status);
		} else {
			let _ = writeln!(status, "Publishing disabled");
		}

		if let Some(info) = self.publication_server.as_ref().map(PublicationServer::get_status) {
			let count = info.0.len();

			if count == 0 {
				let _ = writeln!(status, "Accepting publications from no servers");
			} else {
				let _ = writeln!(status, "Accepting publications from {count} server{}:", if count == 1 { "" } else { "s" });
				for server in info.0 {
					let _ = write!(status, "{INDENTATION}{}", server.addr);
					if let Some(hostname) = server.hostname {
						let _ = write!(status, " ({hostname})");
					}
					let game_count = server.game_server_count;
					let _ = writeln!(status, ": {game_count} game server{}", if game_count == 1 { "" } else { "s" });
				}
			}
		} else {
			let _ = writeln!(status, "Not accepting publications");
		}

		status
	}

	/**
	 * Returns a list of IP addresses that other devices on the local network can
	 * connect to.
	 *
	 * This obviously excludes loopback addresses, but also excludes IPv6 unicast
	 * link-local addresses because they require a scope ID which cannot be
	 * found.
	 */
	#[cfg(feature = "client")]
	pub fn get_locally_connectable_hosts(&self) -> impl Iterator<Item = IpAddr> + use<'_> {
		self.networks
			.values()
			.flat_map(NetworkData::ip_networks)
			.map(|ip_net| ip_net.addr)
			.filter(|addr| !(addr.is_loopback() || match addr {
				IpAddr::V4(_addr) => false,
				IpAddr::V6(addr) => addr.is_unicast_link_local(),
			}))
	}

	async fn handle_incoming(self: Arc<Self>, incoming: Incoming) {
		let addr = incoming.remote_address().to_canonical();
		match incoming.await {
			Ok(connection) => {
				match InitServer::recv(connection).await {
					Ok((request, pr)) => {
						let pr = PendingResponse { inner: pr, addr };

						log::info!("received request, addr = {addr}, request = {request:?}");

						if self.rate_limiter.lock().unwrap().allowed(addr, &request) {
							if let Err(err) = self.respond_request(request, pr, addr).await {
								log::info!("failed responding to request, addr = {addr}: {err}");
							}
						} else {
							log::info!("rate limited, addr = {addr}");
							pr.understood("rate limited").await;
						}
					},
					Err(Or::Both(err, response)) => log::info!("failed receiving request, addr = {addr}, response = {response}: {err}"),
					Err(Or::A(err)) => log::info!("failed receiving request, addr = {addr}, response not sent: {err}"),
					Err(Or::B(response)) => log::info!("failed receiving request, addr = {addr}, response = {response}"),
				}
			},
			Err(err) => log::info!("connection attempt failed, addr = {addr}: {err}"),
		}
	}

	async fn respond_request(self: Arc<Self>, request: Request, pr: PendingResponse, addr: SocketAddr) -> Result<(), String> {
		match request {
			Request::Play(request) => {
				let router_res = self.router.read().unwrap().route(&request.game_id, (&self.networks).classify(addr.ip()));
				match router_res {
					Ok(stream_sender) => {
						let _ = stream_sender.send((Box::new(IncomingServerMessageStream::try_new(pr.inner, addr).await?), request)).await;
						Ok(())
					},
					Err(err) => {
						let response_msg = match err {
							RouteError::NotFound => "no server matches game id",
							RouteError::AccessDenied => "access denied",
						};

						pr.understood(response_msg).await;
						Err(String::from(response_msg))
					},
				}
			},
			Request::Info => self.respond_info_request(pr).await,
			Request::Discover => self.respond_discover_request(pr, addr.ip()).await,
			Request::Publish(PublishRequest::Init(return_port, hostname, key)) => self.respond_publish_init_request(pr, addr, return_port, hostname, key).await,
			Request::Publish(PublishRequest::Subscribe(key)) => Arc::clone(&self).respond_publish_subscribe_request(pr, addr, key).await,
		}
	}

	async fn respond_info_request(&self, pr: PendingResponse) -> Result<(), String> {
		Ok(pr.success_unflushed().await?.main_stream.sender.send(&self.info).await?)
	}

	async fn respond_discover_request(&self, pr: PendingResponse, addr: IpAddr) -> Result<(), String> {
		let official = self.router.read().unwrap().discovery_servers((&self.networks).classify(addr)).collect();
		let publications = self.publication_server.as_ref().map(PublicationServer::get_publications).unwrap_or_default();
		let msg = DiscoveryMessage::new(official, publications);
		Ok(pr.success_unflushed().await?.main_stream.sender.send(&msg).await?)
	}

	async fn respond_publish_init_request(&self, pr: PendingResponse, addr: SocketAddr, return_port: u16, hostname: Option<Hostname>, key: Key) -> Result<(), String> {
		let Some(server) = &self.publication_server else {
			pr.understood("publish-init requests disabled").await;
			return Ok(());
		};

		if return_port == addr.port() {
			pr.understood("invalid return port").await;
			return Ok(());
		}

		let return_addr = SocketAddr::new(addr.ip(), return_port);

		match server.respond_publish_init(self.endpoint.clone(), return_addr, hostname, key).await {
			Ok(()) => {
				pr.success().await?;
				Ok(())
			},
			Err((pub_err, priv_err)) => {
				pr.understood(match pub_err {
					PublishInitError::CannotConnectBack => "cannot connect back (perhaps you're behind a NAT)",
					PublishInitError::TooManyConnections => "too many connections",
					PublishInitError::ConnectionAlreadyExists => "connection already exists",
				}).await;

				Err(format!("failed responding to publish-init: {priv_err}"))
			},
		}
	}

	async fn respond_publish_subscribe_request(self: Arc<Self>, pr: PendingResponse, addr: SocketAddr, key: Key) -> Result<(), String> {
		static DISABLED_MESSAGE: &str = "publish-subscribe requests disabled";

		let Some(server) = &self.publishing_server else {
			pr.understood(DISABLED_MESSAGE).await;
			return Ok(());
		};

		match server.respond_publish_subscribe(key).await {
			Ok(subscriber) => {
				let sender = pr.success_unflushed().await?.main_stream.sender; // Initial message containing servers is sent quickly, so no need to flush
				let ip_category = (&self.networks).classify(addr.ip());
				subscriber.run(sender, self, ip_category).await?;
			},
			Err(err) => {
				pr.understood(match err {
					SubscribeError::ConnectionAlreadyExists => "publish-subscribe connection already exists",
					SubscribeError::IncorrectKey => "incorrect key",
					SubscribeError::TaskFinished => DISABLED_MESSAGE,
				}).await;
			},
		}

		Ok(())
	}
}
