use half::bf16;
use std::path::PathBuf;
use xserv_model::{KVCache, ModelConfig, Qwen3, loader};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let model_dir = PathBuf::from(&args[1]);
    let prompt = &args[2];

    xserv_cuda::device::set_device(0).unwrap();
    let config = ModelConfig::from_file(&model_dir.join("config.json"));
    let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
    let model = Qwen3::from_weights(config.clone(), weights);
    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));

    let token_ids = tokenizer.encode(prompt);
    eprintln!("Prompt: {prompt}");
    eprintln!("Token IDs: {token_ids:?}");

    let mut cache = KVCache::new(
        config.num_layers(),
        config.num_kv_heads(),
        config.head_dim(),
        DType::BF16,
        Device::Cuda(0),
    );
    let logits = model.forward_with_cache(&token_ids, &mut cache);
    let logits_cpu = logits.to_device(Device::Cpu);
    let data = logits_cpu.as_slice::<bf16>();
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];

    // Print top-20 logits for the last position
    let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
    let mut indexed: Vec<(usize, f32)> = last_row
        .iter()
        .enumerate()
        .map(|(i, v)| (i, v.to_f32()))
        .collect();
    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    println!("Top-20 logits (last position):");
    for (rank, (id, val)) in indexed.iter().take(20).enumerate() {
        let tok = tokenizer.decode(&[*id as u32]);
        println!("  [{rank:>2}] id={id:>6} logit={val:>10.4} token={tok:?}");
    }
}