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, Mutex};
use tokio::sync::Mutex as TokioMutex;
use tracing::{info, warn};
use uuid::Uuid;

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

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

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

    fn get_or_create_room(&self, room_id: &str) -> Arc<TokioMutex<Room>> {
        let mut rooms = self.rooms.lock().unwrap();
        if let Some(room) = rooms.get(room_id).cloned() {
            return room;
        }
        let room = Arc::new(TokioMutex::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);
    (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_arc = state.get_or_create_room(&room_id);
    let player_id = Uuid::new_v4().to_string();
    let player_id_for_recv = player_id.clone();

    let (mut sender, mut receiver) = socket.split();

    // Reject room-full before subscribing so the error goes only to the joining socket.
    {
        let room = room_arc.lock().await;
        if room.game.can_start() {
            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;
        }
    }

    // Add the player to the room and tell only this socket its assigned team.
    let team = {
        let mut room = room_arc.lock().await;
        let team = room.game.add_player(player_id.clone(), preferred_team)
            .unwrap_or(Team::Red);
        if room.game.can_start() {
            room.game.start();
        }
        team
    };

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

    // Subscribe to broadcast and spawn the forwarding task.
    let tx = {
        let room = room_arc.lock().await;
        room.tx.clone()
    };
    let mut rx = tx.subscribe();

    let send_task = 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,
            }
        }
    });

    // Broadcast waiting/game state to everyone in the room.
    {
        let room = room_arc.lock().await;
        if room.game.can_start() {
            let state_msg = room.game.game_state_message();
            let _ = room.tx.send(state_msg);
        } else {
            let _ = room.tx.send(ServerMessage::Waiting { message: "Waiting for other player".to_string() });
        }
    }

    let recv_room = room_arc.clone();
    let recv_id = player_id_for_recv;
    let recv_task = 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 room = recv_room.clone();
                    let mut room = room.lock().await;
                    if room.game.current_team_for_player(&recv_id) != Some(team) {
                        let _ = room.tx.send(ServerMessage::Error { message: "Not your turn".to_string() });
                        continue;
                    }
                    match room.game.handle_throw(&recv_id, broom_x, broom_y, weight, curl, friction) {
                        Ok(path) => {
                            room.tx.send(ServerMessage::Trajectory { path }).ok();
                            room.game.finish_simulation();
                            if let Some(scored) = room.game.take_last_end_scored() {
                                room.tx.send(scored).ok();
                            }
                            let after = room.game.game_state_message();
                            room.tx.send(after).ok();
                            if room.game.phase() == GamePhase::GameComplete {
                                room.tx.send(room.game.game_over_message()).ok();
                            }
                        }
                        Err(e) => {
                            room.tx.send(ServerMessage::Error { message: e }).ok();
                        }
                    }
                }
                Err(e) => {
                    warn!("Invalid message: {}", e);
                }
            }
        }
    });

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

    // On disconnect, free the team slot so refreshes and new tabs can rejoin.
    {
        let mut room = room_arc.lock().await;
        let _ = room.game.remove_player(&player_id);
    }
}