mod api;
mod engine;
mod pp_engine;
mod tp_engine;

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

pub struct AppState {
    pub model_name: String,
    pub chat_template: api::ChatTemplate,
    pub engine_sender: Mutex<mpsc::SyncSender<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] [--swap-space-gb N] [--tp N] [--pp 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)
        .max(1);
    let requested_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)
        .max(1);
    let swap_space_gb: usize = args
        .iter()
        .position(|a| a == "--swap-space-gb")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(8);
    let tp: usize = args
        .iter()
        .position(|a| a == "--tp")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(1)
        .max(1);
    let pp: usize = args
        .iter()
        .position(|a| a == "--pp")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(1)
        .max(1);
    if tp > 1 && pp > 1 {
        eprintln!("--tp and --pp cannot be combined yet (2D TP×PP is future work)");
        std::process::exit(1);
    }
    let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
    // gpt-oss is only implemented in the TP engine; route it there even at
    // tp=1 (single-rank world) so quantized models can serve on one GPU.
    let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss");
    if pp > 1 && is_gpt_oss {
        eprintln!(
            "gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead"
        );
        std::process::exit(1);
    }
    let model_max_seq_len = model_config.max_seq_len();
    if model_max_seq_len == 0 {
        eprintln!("model config has invalid max_seq_len=0");
        std::process::exit(1);
    }
    let max_seq_len = requested_max_seq_len.min(model_max_seq_len);
    if max_seq_len != requested_max_seq_len {
        eprintln!(
            "[server] --max-seq-len {requested_max_seq_len} exceeds model limit {model_max_seq_len}; using {max_seq_len}"
        );
    }

    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"));

    // Bounded channel to backpressure incoming requests when the engine falls
    // behind, instead of letting them pile up in RAM. try_send in the API
    // handler surfaces this as 503 to the client.
    let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(256);

    let model_dir_clone = model_dir.clone();
    std::thread::spawn(move || {
        if pp > 1 {
            // Pipeline-parallel path: stage-0 coordinator + worker stage threads.
            pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
        } else if tp <= 1 && !is_gpt_oss {
            let mut engine = engine::Engine::load_with_swap(
                &model_dir_clone,
                max_batch,
                max_seq_len,
                swap_space_gb,
            );
            engine.run(rx);
        } else {
            // Tensor-parallel path: rank-0 coordinator + worker rank threads.
            tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx);
        }
    });

    let model_type = model_config.model_type.clone().unwrap_or_default();
    let chat_template = api::ChatTemplate::load(&model_dir, &model_type);
    let state = Arc::new(AppState {
        model_name,
        chat_template,
        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(DefaultBodyLimit::max(4 * 1024 * 1024))
        .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();
}