openws-server 0.1.5

mod broadcast;
mod config;
mod handler;
mod message;
mod room;
mod state;

use std::time::Instant;

use axum::routing::{get, post};
use axum::{Json, Router};
use clap::Parser;
use tower_http::cors::CorsLayer;
use tracing::info;

use crate::config::mask_api_key;
use crate::state::AppState;

#[derive(Parser)]
#[command(name = "openws-server", version, about = "WebSocket server with room-based messaging")]
struct Cli {
    #[arg(long, help = "Host to bind to")]
    host: Option<String>,

    #[arg(short, long, help = "Port to bind to")]
    port: Option<u16>,

    #[arg(long, help = "Max clients allowed")]
    max_clients: Option<usize>,

    #[arg(long, help = "Max rooms allowed")]
    max_rooms: Option<usize>,

    #[arg(long, help = "Max message size in bytes")]
    max_message_size: Option<usize>,

    #[arg(long, help = "Max username length")]
    max_username_len: Option<usize>,

    #[arg(long, help = "Ping interval in seconds")]
    ping_interval: Option<u64>,

    #[arg(long, help = "Pong timeout in seconds")]
    pong_timeout: Option<u64>,

    #[arg(long, help = "Broadcast API key")]
    api_key: Option<String>,
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "openws_server=info".into()),
        )
        .init();

    let cli = Cli::parse();

    let mut cfg = config::load_config().expect("failed to load configuration");

    if let Some(host) = cli.host {
        cfg.server.host = host;
    }
    if let Some(port) = cli.port {
        cfg.server.port = port;
    }
    if let Some(v) = cli.max_clients {
        cfg.limits.max_clients = v;
    }
    if let Some(v) = cli.max_rooms {
        cfg.limits.max_rooms = v;
    }
    if let Some(v) = cli.max_message_size {
        cfg.limits.max_message_size = v;
    }
    if let Some(v) = cli.max_username_len {
        cfg.limits.max_username_len = v;
    }
    if let Some(v) = cli.ping_interval {
        cfg.heartbeat.ping_interval_secs = v;
    }
    if let Some(v) = cli.pong_timeout {
        cfg.heartbeat.pong_timeout_secs = v;
    }
    if let Some(v) = cli.api_key {
        cfg.broadcast.api_key = v;
    }
    let state = AppState::new(cfg.clone());
    let started_at: Instant = state.started_at;

    let app = Router::new()
        .route("/ws", get(handler::ws_handler))
        .route("/api/broadcast", post(broadcast::broadcast_handler))
        .route(
            "/health",
            {
                let state = state.clone();
                get(move || {
                    let state = state.clone();
                    async move {
                        let rooms = state.rooms.read().await;
                        let total_clients: usize = rooms.values().map(|r| r.clients.len()).sum();
                        Json(serde_json::json!({
                            "status": "ok",
                            "rooms": rooms.len(),
                            "clients": total_clients,
                            "uptime_secs": started_at.elapsed().as_secs(),
                        }))
                    }
                })
            },
        )
        .route(
            "/stats",
            {
                let state = state.clone();
                get(move || {
                    let state = state.clone();
                    async move { Json(state.stats().await) }
                })
            },
        )
        .layer(CorsLayer::permissive())
        .with_state(state);

    let bind_addr = format!("{}:{}", cfg.server.host, cfg.server.port);

    info!(
        "openws-server v{} starting on {}",
        env!("CARGO_PKG_VERSION"),
        bind_addr,
    );
    info!(
        "limits: max_clients={}, max_rooms={}, max_message_size={}, max_username_len={}",
        cfg.limits.max_clients,
        cfg.limits.max_rooms,
        cfg.limits.max_message_size,
        cfg.limits.max_username_len,
    );
    info!(
        "heartbeat: ping_interval={}s, pong_timeout={}s",
        cfg.heartbeat.ping_interval_secs, cfg.heartbeat.pong_timeout_secs,
    );
    info!(
        "broadcast api key: {}",
        mask_api_key(&cfg.broadcast.api_key)
    );

    let listener = tokio::net::TcpListener::bind(&bind_addr)
        .await
        .expect("failed to bind to address");
    axum::serve(listener, app)
        .await
        .expect("server error");
}