button/src/main.rs

use std::{
	collections::HashMap,
	fs,
	io::Write,
	ops::{Deref, DerefMut},
	sync::Arc,
};

use axum::{
	Router,
	extract::{
		State,
		ws::{Message, WebSocket, WebSocketUpgrade},
	},
	response::{Html, IntoResponse},
	routing::get,
};
use futures::StreamExt as _;
use rand::random_bool;
use serde::{Deserialize, Serialize, de};
use serde_json::json;
use tokio::{
	sync::{Mutex, mpsc},
	time::{Duration, sleep},
};

#[derive(Deserialize, Serialize, Debug, Ord, Eq, PartialEq, PartialOrd, Clone)]
struct Entry {
	score: u32,
	person: String,
}

#[derive(Clone)]
struct AppState {
	tx: mpsc::Sender<LeaderboardUpdate>,
	hiscores: Arc<Mutex<Vec<Entry>>>,
	loscores: Arc<Mutex<Vec<Entry>>>,
	pingscores: Arc<Mutex<HashMap<String, (u64, u32)>>>, // u64 is reset count and u32 is PB
}

struct LeaderboardUpdate {
	name: Arc<str>,
	update: LeaderboardUpdateType,
}

enum LeaderboardUpdateType {
	Reset { hiscore_pingscore: u32 },
	Increment { loscore: u32 },
}

const CHANCE: f64 = 1.0 / 3.0;

const PATH_HISCORES: &str = "hiscores.json";
const PATH_LOSCORES: &str = "loscores.json";
const PATH_PINGSCORES: &str = "pingscores.json";

const MAX_LEADERBOARD: usize = 20;

#[tokio::main]
async fn main() {
	fn read_file<T: for<'de> de::Deserialize<'de>>(file_path: &str) -> Arc<Mutex<T>> {
		let file_contents: String = fs::read_to_string(file_path).unwrap();
		Arc::new(Mutex::new(serde_json::from_str(&file_contents).unwrap()))
	}

	/// Makes the vector at `vec` one with a capacity of exactly [`MAX_LEADERBOARD`] if `vec` is
	/// smaller or equal.
	fn exact_leaderboard<T>(mut vec: impl DerefMut<Target = Vec<T>>) {
		let mut old_vec = std::mem::replace(&mut *vec, Vec::with_capacity(MAX_LEADERBOARD));
		old_vec.drain(..).for_each(|e| vec.push(e));
	}

	let hiscores: Arc<Mutex<Vec<Entry>>> = read_file(PATH_HISCORES);
	exact_leaderboard(hiscores.lock().await);
	let loscores: Arc<Mutex<Vec<Entry>>> = read_file(PATH_LOSCORES);
	exact_leaderboard(loscores.lock().await);
	let pingscores: Arc<Mutex<HashMap<String, (u64, u32)>>> = read_file(PATH_PINGSCORES);

	let (tx, rx) = mpsc::channel::<LeaderboardUpdate>(1024);

	{
		let (hiscores, loscores, pingscores) =
			(hiscores.clone(), loscores.clone(), pingscores.clone());
		tokio::spawn(async move {
			handle_hiscores(rx, &hiscores, &loscores, &pingscores).await;
		});
	}

	{
		let pingscores = pingscores.clone();
		tokio::spawn(async move {
			// write pingscores every 30s
			loop {
				sleep(Duration::from_millis(30000)).await;
				let pingscores = pingscores.lock().await;
				let file_contents: String = serde_json::to_string(&pingscores.clone()).unwrap();
				drop(pingscores);
				let mut file = fs::OpenOptions::new()
					.write(true)
					.truncate(true)
					.open(PATH_PINGSCORES)
					.unwrap();
				file.write_all(file_contents.as_bytes()).unwrap();
				file.flush().unwrap();
			}
		});
	}

	let app = Router::new()
		.route("/", get(index))
		.route("/ws", get(ws_handler))
		.route("/leaderboard", get(leaderboard))
		.with_state(AppState {
			tx,
			hiscores: Arc::clone(&hiscores),
			loscores: Arc::clone(&loscores),
			pingscores: Arc::clone(&pingscores),
		});

	let listener = tokio::net::TcpListener::bind("0.0.0.0:8084").await.unwrap();

	println!("http://0.0.0.0:8084");

	axum::serve(listener, app).await.unwrap();
}

async fn index() -> Html<&'static str> {
	Html(include_str!("../index.html"))
}

async fn leaderboard() -> Html<&'static str> {
	Html(include_str!("../leaderboard.html"))
}
// receiver: 0 for hiscore, 1 for loscore, 2 for pingscore
async fn handle_hiscores(
	mut rx: mpsc::Receiver<LeaderboardUpdate>,
	hiscores: &Mutex<Vec<Entry>>,
	loscores: &Mutex<Vec<Entry>>,
	pingscores: &Mutex<HashMap<String, (u64, u32)>>,
) {
	fn update_scoretable<G: Deref<Target = Vec<Entry>> + DerefMut>(
		score_name: &str,
		mut scoretable_lock: G,
		name: &str,
		score: u32,
		file_path: &str,
	) {
		let scoretable = &mut *scoretable_lock;

		if let Some(index_to_insert_at) = scoretable.iter().position(|e| score > e.score) {
			println!("New {score_name} {score} by {name}");
			scoretable[index_to_insert_at..].rotate_right(1);
			let push_out = std::mem::replace(
				&mut scoretable[index_to_insert_at],
				Entry {
					score,
					person: name.to_string(),
				},
			);
			if scoretable.len() < MAX_LEADERBOARD {
				scoretable.push(push_out);
			}

			let file_contents: String = serde_json::to_string(&*scoretable_lock).unwrap();
			drop(scoretable_lock);
			let mut file = fs::OpenOptions::new()
				.write(true)
				.truncate(true)
				.open(file_path)
				.unwrap();
			file.write_all(file_contents.as_bytes()).unwrap();
			file.flush().unwrap();
		} else if scoretable.len() < MAX_LEADERBOARD {
			println!("New {score_name} {score} by {name}");
			scoretable.push(Entry {
				score,
				person: name.to_string(),
			});
		}
	}

	// Panic galore
	let mut hiscores_lock = hiscores.lock().await;
	hiscores_lock.sort();
	hiscores_lock.reverse();
	let file_contents: String = serde_json::to_string(&hiscores_lock.clone()).unwrap();
	drop(hiscores_lock);
	let mut file = fs::OpenOptions::new()
		.write(true)
		.truncate(true)
		.open(PATH_HISCORES)
		.unwrap();
	file.write_all(file_contents.as_bytes()).unwrap();
	drop(file);

	loop {
		let LeaderboardUpdate { name, update } = rx.recv().await.expect("channel error");

		match update {
			LeaderboardUpdateType::Reset { hiscore_pingscore } => {
				// Hiscore
				update_scoretable(
					"hiscore",
					hiscores.lock().await,
					&name,
					hiscore_pingscore,
					PATH_HISCORES,
				);

				// Pingscore
				let mut pingscores = pingscores.lock().await;
				// pb
				if hiscore_pingscore > pingscores.get(&*name).unwrap_or(&(0, 0)).1 {
					pingscores.entry(name.to_string()).or_insert((0, 0)).1 = hiscore_pingscore;
					println!("{name} new PB: {hiscore_pingscore}");
				};
				pingscores.entry(name.to_string()).or_insert((0, 0)).0 += 1; // reset count
				drop(pingscores);
			}

			LeaderboardUpdateType::Increment { loscore } => {
				update_scoretable(
					"loscore",
					loscores.lock().await,
					&name,
					loscore,
					PATH_LOSCORES,
				);
			}
		}
	}
}

async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
	ws.on_upgrade(|socket| async move {
		handle_socket(
			socket,
			&state.tx,
			&state.hiscores,
			&state.loscores,
			&state.pingscores,
		)
		.await;
	})
}

async fn handle_socket(
	mut socket: WebSocket,
	tx: &mpsc::Sender<LeaderboardUpdate>,
	hiscores: &Mutex<Vec<Entry>>,
	loscores: &Mutex<Vec<Entry>>,
	pingscores: &Mutex<HashMap<String, (u64, u32)>>,
) {
	let mut value: u32 = 0;

	let msg = {
		let hiscores = hiscores.lock().await;
		let loscores = loscores.lock().await;
		let pingscores = pingscores.lock().await;
		json!({ "hiscores": &*hiscores, "loscores": &*loscores, "pingscores": &*pingscores})
			.to_string()
	};

	let name: Arc<str> = match socket.next().await.unwrap().unwrap() {
		Message::Text(text)
			if let text = text.to_string()
				&& validate_name(&text) =>
		{
			Arc::from(text.into_boxed_str())
		}
		_ => Arc::from("anon"),
	};

	println!("Client connected: {name}");

	let mut resets: u32 = 0;
	let mut prev: u32 = 0;

	let _ = socket.send(Message::Text(msg.into())).await;

	while let Some(msg) = socket.next().await {
		match msg {
			Ok(Message::Text(_)) => {
				if random_bool(CHANCE) {
					// reset
					let _ = tx
						.send(LeaderboardUpdate {
							name: name.clone(),
							update: LeaderboardUpdateType::Reset {
								hiscore_pingscore: value,
							},
						})
						.await;
					resets += 1;
					value = 0
				}
				// 1/3 chance of failing
				else {
					value += 1;
					if prev == 0 {
						let _ = tx
							.send(LeaderboardUpdate {
								name: name.clone(),
								update: LeaderboardUpdateType::Increment { loscore: resets },
							})
							.await;
						resets = 0;
					}
				}
				let _ = socket.send(Message::Text(value.to_string().into())).await;
				prev = value;
			}

			Ok(Message::Close(_)) => {
				break;
			}

			_ => {}
		}
	}
}

fn validate_name(input: &str) -> bool {
	let input = input.trim();

	// Length check
	if input.is_empty() || input.len() > 32 {
		return false;
	}

	input.chars().all(|c| c.is_ascii_alphanumeric())
}