mod api;
mod engine;

use axum::{routing::{get, post}, Extension, Router};
use std::path::PathBuf;
use std::sync::{mpsc, Arc, Mutex};
use engine::GenerateRequest;

pub struct AppState {
    pub model_name: String,
    pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
    pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
    pub max_seq_len: usize,
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N]");
        std::process::exit(1);
    }

    let model_dir = PathBuf::from(&args[1]);
    let port: u16 = args.iter()
        .position(|a| a == "--port")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(8080);
    let max_batch: usize = args.iter()
        .position(|a| a == "--max-batch")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(4);
    let max_seq_len: usize = args.iter()
        .position(|a| a == "--max-seq-len")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(2048);

    let model_name = model_dir.file_name()
        .map(|n| n.to_string_lossy().to_string())
        .unwrap_or_else(|| "unknown".to_string());

    let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));

    // Unbounded channel: allows multiple requests to queue up
    let (tx, rx) = mpsc::channel::<GenerateRequest>();

    let model_dir_clone = model_dir.clone();
    std::thread::spawn(move || {
        let engine = engine::Engine::load(&model_dir_clone, max_batch, max_seq_len);
        engine.run(rx);
    });

    let state = Arc::new(AppState {
        model_name,
        engine_sender: Mutex::new(tx),
        engine_tokenizer: Mutex::new(tokenizer),
        max_seq_len,
    });

    let app = Router::new()
        .route("/health", get(api::health))
        .route("/v1/models", get(api::list_models))
        .route("/v1/chat/completions", post(api::chat_completions))
        .layer(Extension(state));

    let addr = format!("0.0.0.0:{port}");
    eprintln!("[server] Listening on {addr} (max_batch={max_batch}, max_seq_len={max_seq_len})");
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    axum::serve(listener, app).await.unwrap();
}