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

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

use super::{SubscribeMessage, Key, status::{PublicationStatus, PublishingServerStatus}, super::{ip_block::IpBlock, rate_limiter::bucket::Bucket}};

use crate::net::{serp::{InitClient, Receiver, request::{Request, PublishRequest}}, utils::{quic::{self, IdleConfig}, hostname::Hostname}};
use crate::protocol::discovery::{DiscoveryServer, UnofficialSource, GameId};

pub struct PublicationServer {
	servers: Arc<Mutex<ServerMap>>,
	sender: MpscSender<(SocketAddr, Message)>,
}

struct ServerMap {
	map: BTreeMap<SocketAddr, Server>,
	blocks: BTreeMap<IpBlock, BTreeSet<SocketAddr>>,
}

impl ServerMap {
	const MAX_CONNECTIONS_PER_BLOCK: usize = 2;

	fn remove_from_block(&mut self, addr: SocketAddr) {
		match self.blocks.entry(addr.into()) {
			Entry::Occupied(mut e) => {
				let block = e.get_mut();
				let found = block.remove(&addr);
				debug_assert!(found);

				if block.is_empty() {
					e.remove();
				}
			},
			Entry::Vacant(_) => debug_assert!(false),
		}
	}
}

struct Server {
	hostname: Option<Hostname>,
	servers: BTreeMap<GameId, DiscoveryServer>,
}

impl Server {
	const MAX_GAME_SERVERS_PER_CONNECTION: usize = 16;

	fn add(&mut self, server: DiscoveryServer) {
		let len = self.servers.len();
		match self.servers.entry(server.game_id.clone()) {
			Entry::Vacant(e) => {
				if len < Server::MAX_GAME_SERVERS_PER_CONNECTION {
					e.insert(server);
				}
			},
			Entry::Occupied(mut e) => _ = e.insert(server),
		}
	}
}

enum Message {
	Init(Option<Hostname>),
	Messages(Vec<SubscribeMessage>),
	Finish,
}

pub enum InitError {
	CannotConnectBack,
	TooManyConnections,
	ConnectionAlreadyExists,
}

impl PublicationServer {
	pub fn new() -> PublicationServer {
		let servers = Arc::new(Mutex::new(ServerMap { map: BTreeMap::new(), blocks: BTreeMap::new() }));
		let (sender, receiver) = mpsc::channel(64);
		tokio::spawn(Self::main_task(Arc::clone(&servers), receiver));
		PublicationServer { servers, sender }
	}

	async fn main_task(servers: Arc<Mutex<ServerMap>>, mut receiver: MpscReceiver<(SocketAddr, Message)>) {
		while let Some((addr, msg)) = receiver.recv().await {
			let mut servers = servers.lock().unwrap();
			match (msg, servers.map.entry(addr)) {
				(Message::Init(hostname), Entry::Vacant(e)) => _ = e.insert(Server { hostname, servers: BTreeMap::new() }),
				(Message::Messages(messages), Entry::Occupied(mut e)) => {
					let pub_server = e.get_mut();
					for msg in messages {
						log::info!("subscribe message from {addr}: {msg:?}");

						match msg {
							SubscribeMessage::Init(new_servers) => {
								for server in new_servers {
									pub_server.add(server);
								}
							},
							SubscribeMessage::Add(server) => pub_server.add(server),
							SubscribeMessage::Remove(id) => _ = pub_server.servers.remove(&id),
							SubscribeMessage::UpdateClientCount(id, players, spectators) => {
								if let Some(server) = pub_server.servers.get_mut(&id) {
									server.player_count = players;
									server.spectator_count = spectators;
								}
							},
						}
					}
				},
				(Message::Finish, Entry::Occupied(e)) => {
					e.remove();
					servers.remove_from_block(addr);
				},
				_ => debug_assert!(false),
			}
			drop(servers);
		}
	}

	pub async fn respond_publish_init(&self, endpoint: Endpoint, addr: SocketAddr, hostname: Option<Hostname>, key: Key) -> Result<(), (InitError, String)> {
		{
			let mut servers = self.servers.lock().unwrap();
			let block = servers.blocks.entry(addr.into()).or_default();
			if block.len() < ServerMap::MAX_CONNECTIONS_PER_BLOCK {
				let not_found = block.insert(addr);
				drop(servers);
				if !not_found {
					return Err((InitError::ConnectionAlreadyExists, String::from("connection already exists")));
				}
			} else {
				drop(servers);
				return Err((InitError::TooManyConnections, String::from("too many connections")));
			}
		}

		let res = time::timeout(Duration::from_secs(30), Self::request_publish_subscribe(endpoint, addr, hostname.as_ref(), key)).await;
		let mut receiver = res
			.unwrap_or_else(|_| Err((InitError::CannotConnectBack, String::from("publish-subscribe connection timed out"))))
			.inspect_err(|_| self.servers.lock().unwrap().remove_from_block(addr))?;

		let sender = self.sender.clone();
		tokio::spawn(async move {
			const BUCKET_AMOUNT_PER_MESSAGE: f32 = 0.5;
			const BUCKET_CAPACITY: f32 = 8.0;

			log::info!("new publish-subscribe connection, addr = {addr}, hostname = {hostname:?}");

			let mut bucket = Bucket::new(Instant::now());

			let _ = sender.send((addr, Message::Init(hostname))).await;
			loop {
				match receiver.recv::<Vec<SubscribeMessage>>().await {
					Ok(messages) => {
						let mut size = 1 + messages.len();
						for msg in &messages {
							size += match msg {
								SubscribeMessage::Init(servers) => servers.len(),
								SubscribeMessage::Add(..) | SubscribeMessage::Remove(..) | SubscribeMessage::UpdateClientCount(..) => 0,
							};
						}
						let _ = sender.send((addr, Message::Messages(messages))).await;
						bucket.always_add(Instant::now(), BUCKET_AMOUNT_PER_MESSAGE, BUCKET_CAPACITY * size as f32).await;
					},
					Err(err) => {
						log::info!("publish-subscribe connection failed: {err}, addr = {addr}");
						break;
					},
				}
			}
			let _ = sender.send((addr, Message::Finish)).await;
		});

		Ok(())
	}

	async fn request_publish_subscribe(endpoint: Endpoint, addr: SocketAddr, hostname: Option<&Hostname>, key: Key) -> Result<Receiver, (InitError, String)> {
		let conn = if let Some(host) = hostname {
			quic::connect_to_addr_secure(endpoint, addr, host, IdleConfig::Larger).await
		} else {
			quic::connect_to_addr(endpoint, addr, IdleConfig::Larger).await
		}.map_err(|err| (InitError::CannotConnectBack, err))?;

		let conn = InitClient::send(conn, Request::Publish(PublishRequest::Subscribe(key))).await.map_err(|err| (InitError::CannotConnectBack, format!("failed sending publish-subscribe request: {err}")))?;
		Ok(conn.main_stream.receiver)
	}

	pub fn get_publications(&self) -> BTreeMap<UnofficialSource, Vec<DiscoveryServer>> {
		let mut publications = BTreeMap::new();
		for (addr, pub_server) in &self.servers.lock().unwrap().map {
			let src = match &pub_server.hostname {
				Some(host) => UnofficialSource::Hostname(host.clone()),
				None => UnofficialSource::Ip(addr.ip()),
			};

			let mut servers = pub_server.servers.values().cloned().collect();
			publications.entry(src).and_modify(|s: &mut Vec<DiscoveryServer>| s.append(&mut servers)).or_insert(servers);
		}
		publications
	}

	pub fn get_status(&self) -> PublicationStatus {
		let all_servers = self.servers.lock().unwrap();
		PublicationStatus(all_servers.map
			.iter()
			.map(|(addr, pub_server)| PublishingServerStatus {
				addr: addr.ip(),
				hostname: pub_server.hostname.clone(),
				game_server_count: pub_server.servers.len(),
			})
			.collect())
	}
}
