curltastic/backend/src/main.rs

mod protocol;
mod physics;
mod game;

use axum::{
    extract::{Query, State, WebSocketUpgrade},
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Router,
};
use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
use futures_util::{sink::SinkExt, stream::StreamExt};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{info, warn};
use uuid::Uuid;

use crate::game::{Room, ThrowOutcome};
use crate::protocol::*;

type WsSender = futures_util::stream::SplitSink<WebSocket, Message>;
type WsReceiver = futures_util::stream::SplitStream<WebSocket>;

#[derive(Clone)]
struct AppState {
    rooms: Arc<Mutex<HashMap<String, Arc<Mutex<Room>>>>>,
}

impl AppState {
    fn new() -> Self {
        Self {
            rooms: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    async fn get_or_create_room(&self, room_id: &str) -> Arc<Mutex<Room>> {
        let mut rooms = self.rooms.lock().await;
        if let Some(room) = rooms.get(room_id).cloned() {
            return room;
        }
        let room = Arc::new(Mutex::new(Room::new(room_id)));
        rooms.insert(room_id.to_string(), room.clone());
        room
    }
}

#[derive(Deserialize)]
struct RoomQuery {
    room: String,
    team: Option<Team>,
}

#[derive(Serialize)]
struct NewRoomResponse {
    room: String,
}

fn generate_room_code() -> String {
    let s = Uuid::new_v4().to_string().replace("-", "");
    s.chars().take(6).collect::<String>().to_uppercase()
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();

    let state = Arc::new(AppState::new());

    let app = Router::new()
        .route("/", get(health))
        .route("/room", post(new_room))
        .route("/ws", get(ws_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    info!("Server listening on 0.0.0.0:3000");
    axum::serve(listener, app).await.unwrap();
}

async fn health() -> impl IntoResponse {
    "curltastic ok"
}

async fn new_room(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let room_id = generate_room_code();
    state.get_or_create_room(&room_id).await;
    (StatusCode::OK, axum::Json(NewRoomResponse { room: room_id }))
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    Query(query): Query<RoomQuery>,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, state, query.room, query.team))
}

async fn handle_socket(
    socket: WebSocket,
    state: Arc<AppState>,
    room_id: String,
    preferred_team: Option<Team>,
) {
    let room = state.get_or_create_room(&room_id).await;
    let (mut sender, receiver) = socket.split();

    if try_join_room(&room).await.is_err() {
        let err = serde_json::to_string(&ServerMessage::Error {
            message: "Room is full".to_string(),
        })
        .unwrap();
        let _ = sender.send(Message::Text(Utf8Bytes::from(err))).await;
        return;
    }

    let player_id = Uuid::new_v4().to_string();
    let team = register_player(&room, &player_id, preferred_team).await;

    let joined = serde_json::to_string(&ServerMessage::Joined {
        room: room_id.clone(),
        team,
    })
    .unwrap();
    let _ = sender.send(Message::Text(Utf8Bytes::from(joined))).await;

    let tx = { room.lock().await.tx.clone() };

    let send_task = spawn_forwarder(sender, tx.subscribe());
    broadcast_room_state(&room, &tx).await;

    let recv_task = spawn_message_handler(room.clone(), player_id.clone(), tx, receiver);

    tokio::select! {
        _ = send_task => {}
        _ = recv_task => {}
    }

    remove_player(&room, &player_id).await;
}

async fn try_join_room(room: &Arc<Mutex<Room>>) -> Result<(), ()> {
    let room_guard = room.lock().await;
    if room_guard.game.can_start() {
        return Err(());
    }
    Ok(())
}

async fn register_player(
    room: &Arc<Mutex<Room>>,
    player_id: &str,
    preferred_team: Option<Team>,
) -> Team {
    let mut room_guard = room.lock().await;
    let team = room_guard
        .game
        .add_player(player_id.to_string(), preferred_team)
        .unwrap_or(Team::Red);
    if room_guard.game.can_start() {
        room_guard.game.start();
    }
    team
}

fn spawn_forwarder(
    mut sender: WsSender,
    mut rx: tokio::sync::broadcast::Receiver<ServerMessage>,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(msg) => {
                    let text = match serde_json::to_string(&msg) {
                        Ok(t) => t,
                        Err(_) => continue,
                    };
                    if sender.send(Message::Text(Utf8Bytes::from(text))).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    })
}

async fn broadcast_room_state(
    room: &Arc<Mutex<Room>>,
    tx: &tokio::sync::broadcast::Sender<ServerMessage>,
) {
    let room_guard = room.lock().await;
    let msg = if room_guard.game.can_start() {
        room_guard.game.game_state_message()
    } else {
        ServerMessage::Waiting {
            message: "Waiting for other player".to_string(),
        }
    };
    let _ = tx.send(msg);
}

fn spawn_message_handler(
    room: Arc<Mutex<Room>>,
    player_id: String,
    tx: tokio::sync::broadcast::Sender<ServerMessage>,
    mut receiver: WsReceiver,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            let Message::Text(text) = msg else { continue; };
            let text_ref = text.as_str();
            let parsed: Result<ClientMessage, _> = serde_json::from_str(text_ref);
            match parsed {
                Ok(ClientMessage::Throw { broom_x, broom_y, weight, curl, friction }) => {
                    let mut room_guard = room.lock().await;
                    match room_guard
                        .game
                        .process_throw(&player_id, broom_x, broom_y, weight, curl, friction)
                    {
                        Ok(ThrowOutcome {
                            trajectory,
                            end_scored,
                            state_message,
                            game_over,
                        }) => {
                            let _ = tx.send(ServerMessage::Trajectory { path: trajectory });
                            if let Some(scored) = end_scored {
                                let _ = tx.send(scored);
                            }
                            let _ = tx.send(state_message);
                            if let Some(over) = game_over {
                                let _ = tx.send(over);
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(ServerMessage::Error { message: e });
                        }
                    }
                }
                Err(e) => {
                    warn!("Invalid message: {}", e);
                }
            }
        }
    })
}

async fn remove_player(room: &Arc<Mutex<Room>>, player_id: &str) {
    let mut room_guard = room.lock().await;
    room_guard.game.remove_player(player_id);
}