//! Tensor-parallel E2E benchmark for Qwen3.
//!
//! Spawns one thread per TP rank (each bound to one GPU), loads the sharded
//! model, and runs greedy autoregressive generation. Because lm_head is
//! replicated and the post-AllReduce hidden state is identical on every rank,
//! all ranks compute identical logits and pick the same greedy token — so the
//! rank threads stay in lockstep via the per-layer AllReduces without any
//! token broadcast. Rank 0 records output + timings.
//!
//! Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]
//!
//! Run with --tp 1 / 2 / 4 and compare the printed text (correctness) and
//! tok/s (performance).

use std::path::PathBuf;
use std::sync::Arc;
use std::thread;
use std::time::Instant;

use xserv_model::qwen3::sample_greedy;
use xserv_model::{loader, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

struct PromptResult {
    gen_ids: Vec<u32>,
    ttft_ms: f64,
    decode_tok_s: f64,
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: bench-tp <model-dir> [--tp N] [--gen-tokens N] [--devices 0,1,2,3]");
        std::process::exit(1);
    }
    let model_dir = PathBuf::from(&args[1]);
    let world: usize = arg(&args, "--tp").and_then(|s| s.parse().ok()).unwrap_or(1).max(1);
    let gen_tokens: usize = arg(&args, "--gen-tokens").and_then(|s| s.parse().ok()).unwrap_or(64);
    let devices: Vec<u32> = match arg(&args, "--devices") {
        Some(s) => s.split(',').filter_map(|d| d.trim().parse().ok()).collect(),
        None => (0..world as u32).collect(),
    };
    assert_eq!(devices.len(), world, "--devices count must equal --tp");

    let config = ModelConfig::from_file(&model_dir.join("config.json"));
    assert!(
        config.num_kv_heads() % world == 0,
        "num_kv_heads {} not divisible by tp {world}",
        config.num_kv_heads()
    );
    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
    let eos = tokenizer.eos_token_id();

    let prompts: Vec<&str> = vec![
        "The capital of France is",
        "Explain photosynthesis in one sentence.",
        "Write a haiku about the ocean.",
        "List three uses of a hammer.",
        "What is the speed of light?",
        "Describe the water cycle briefly.",
        "Who wrote Romeo and Juliet?",
        "Translate 'good morning' into Spanish.",
    ];
    let prompt_ids: Vec<Vec<u32>> = prompts.iter().map(|p| tokenizer.encode(p)).collect();

    // Tensors are not Send (their Storage holds a raw GPU pointer), so each rank
    // thread loads its own CPU copy of the weights and shards in-thread. Loading
    // is not part of the timed region.
    let id = if world > 1 { Some(xserv_distributed::get_unique_id()) } else { None };

    let handles: Vec<_> = (0..world)
        .map(|rank| {
            let model_dir = model_dir.clone();
            let config = config.clone();
            let prompt_ids = prompt_ids.clone();
            let device = devices[rank];
            thread::spawn(move || {
                run_rank(rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos)
            })
        })
        .collect();

    let mut rank0: Option<Vec<PromptResult>> = None;
    for (rank, h) in handles.into_iter().enumerate() {
        let r = h.join().expect("rank thread panicked");
        if rank == 0 {
            rank0 = r;
        }
    }

    let results = rank0.expect("rank 0 produced no results");
    println!("\n=== TP={world} (devices {devices:?}) — Qwen3 E2E benchmark ===");
    println!("{:<45} {:>10} {:>12} {:>8}", "prompt", "TTFT(ms)", "decode tok/s", "gen");
    let mut tps_sum = 0.0;
    for (i, r) in results.iter().enumerate() {
        let text = tokenizer.decode(&r.gen_ids).replace('\n', " ");
        let short: String = text.chars().take(50).collect();
        let p: String = prompts[i].chars().take(43).collect();
        println!(
            "{:<45} {:>10.1} {:>12.1} {:>8} | {}",
            p, r.ttft_ms, r.decode_tok_s, r.gen_ids.len(), short
        );
        tps_sum += r.decode_tok_s;
    }
    println!("--- mean decode throughput: {:.1} tok/s ---", tps_sum / results.len() as f64);

    // Machine-readable line for cross-TP correctness diffing (rank 0 token ids).
    let all_ids: Vec<String> = results
        .iter()
        .map(|r| r.gen_ids.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","))
        .collect();
    println!("CORRECTNESS_IDS tp={world} {}", all_ids.join(" | "));
}

fn run_rank(
    rank: usize,
    world: usize,
    device: u32,
    id: Option<xserv_distributed::UniqueId>,
    config: ModelConfig,
    model_dir: PathBuf,
    prompt_ids: Vec<Vec<u32>>,
    gen_tokens: usize,
    eos: Option<u32>,
) -> Option<Vec<PromptResult>> {
    // Bind this thread to its GPU and set up the TP communicator.
    let tp = if world > 1 {
        Some(Arc::new(xserv_distributed::TpContext::init(rank, world, id.unwrap(), device)))
    } else {
        xserv_cuda::device::set_device(device).unwrap();
        None
    };

    // Load this rank's own CPU copy of the weights and shard in-thread.
    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp.clone());

    // Per-rank paged KV cache holds only this rank's local KV heads.
    let local_kv = config.num_kv_heads() / world;
    let max_seq = 2048usize;
    let max_blocks_per_seq = max_seq.div_ceil(BLOCK_SIZE);
    let total_blocks = max_blocks_per_seq + 8;
    let mut cache = PagedKVCache::new_tp(
        &config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, device,
    );

    // Warmup (init kernels / allocator / NCCL channels) — not timed.
    cache.register_sequence(0).unwrap();
    let _ = model.forward_prefill_paged(&[1u32, 2, 3], 0, &mut cache);
    cache.free_sequence(0);

    let mut out = Vec::new();
    for ids in &prompt_ids {
        cache.register_sequence(0).unwrap();

        // Prefill (TTFT).
        let t0 = Instant::now();
        let logits = model.forward_prefill_paged(ids, 0, &mut cache);
        let first = sample_greedy(&logits);
        let ttft_ms = t0.elapsed().as_secs_f64() * 1000.0;

        let mut generated = vec![first];

        // Decode.
        let t1 = Instant::now();
        let mut steps = 0usize;
        for _ in 1..gen_tokens {
            let last = *generated.last().unwrap();
            if eos == Some(last) {
                break;
            }
            let pos = cache.seq_len(0);
            let logits = model.forward_decode_paged(&[last], &[pos], &[0], &mut cache);
            let next = sample_greedy(&logits);
            generated.push(next);
            steps += 1;
        }
        let decode_s = t1.elapsed().as_secs_f64();
        let decode_tok_s = if steps > 0 && decode_s > 0.0 { steps as f64 / decode_s } else { 0.0 };

        cache.free_sequence(0);

        if rank == 0 {
            out.push(PromptResult { gen_ids: generated, ttft_ms, decode_tok_s });
        }
    }

    if rank == 0 { Some(out) } else { None }
}

fn arg<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(|s| s.as_str())
}