//! EAGLE3 sanity check: load weights, run one draft step, print top-5 predictions.
//!
//! This verifies that:
//! - Eagle3Head weights load without shape mismatches
//! - Target hidden states can be captured via decode_core_with_hidden
//! - Eagle3Head::step produces a valid token id (in target vocab)
//!
//! Does NOT measure speedup — that requires a full γ≥2 speculative loop, which
//! is more complex integration work.

use std::path::PathBuf;

use xserv_model::eagle3::{EAGLE_HOOK_LAYERS, Eagle3Head};
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 3 {
        eprintln!("Usage: check-eagle3 <target-model-dir> <eagle3-model-dir> [prompt]");
        std::process::exit(1);
    }
    let target_dir = PathBuf::from(&args[1]);
    let eagle_dir = PathBuf::from(&args[2]);
    let prompt = args
        .get(3)
        .cloned()
        .unwrap_or_else(|| "The capital of France is".to_string());
    let device: u32 = 0;

    xserv_cuda::device::set_device(device).unwrap();

    let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
    eprintln!("Loading target Qwen3-8B...");
    let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
    let target = Qwen3::from_weights(target_config.clone(), target_weights);
    xserv_cuda::allocator::cached_trim();

    eprintln!("Loading EAGLE3 head from {}", eagle_dir.display());
    let mut eagle = Eagle3Head::load(&eagle_dir, device);
    xserv_cuda::allocator::cached_trim();

    let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
    let embed_tokens = target.embed_tokens_tensor();

    let ids = tokenizer.encode(&prompt);
    let max_seq_len = 512;

    let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 2;
    let mut cache = PagedKVCache::new(
        &target_config,
        num_blocks,
        0,
        1,
        num_blocks,
        DType::BF16,
        device,
    );
    cache.register_sequence(0).unwrap();

    // Prefill target.
    let logits = target.forward_prefill_paged(&ids, 0, &mut cache);
    let target_first = *xserv_kernels::argmax_bf16_to_host(&logits).last().unwrap();
    let target_first_text = tokenizer.decode(&[target_first]);
    println!("Prompt: {:?}", prompt);
    println!(
        "Target argmax after prefill: {} ({:?})",
        target_first, target_first_text
    );

    // Now run one target decode step with target_first to get hidden states at the
    // hook layers.
    let pos = cache.seq_len(0);
    target.decode_prepare(&[pos], &[0], &mut cache);
    let ids_gpu = upload_u32(&[target_first]);
    let pos_gpu = upload_u32(&[pos as u32]);
    let (target_next_logits, hooks) = target.decode_core_with_hidden(
        ids_gpu.as_ptr() as *const std::ffi::c_void,
        pos_gpu.as_ptr() as *const std::ffi::c_void,
        1,
        &[0],
        &mut cache,
        &EAGLE_HOOK_LAYERS,
    );
    let target_next = xserv_kernels::argmax_bf16_single(&target_next_logits);
    let target_next_text = tokenizer.decode(&[target_next]);
    println!(
        "Target argmax after 1 decode step: {} ({:?})",
        target_next, target_next_text
    );

    for (i, h) in hooks.iter().enumerate() {
        println!(
            "hook[{}] (layer {}): shape={:?} dtype={:?}",
            i,
            EAGLE_HOOK_LAYERS[i],
            h.shape(),
            h.dtype()
        );
    }

    // Ask EAGLE what it thinks the NEXT token is (given target_first as prev_token
    // and the hidden states from the position where target_first lives).
    // EAGLE should predict target_next (or close to it) to be useful.
    eagle.reset();
    let (eagle_pred, eagle_logits) = eagle.step(&hooks, embed_tokens, target_first, pos);
    let eagle_pred_text = tokenizer.decode(&[eagle_pred]);
    println!(
        "EAGLE draft prediction (pairing A: prev=target_first): {} ({:?})",
        eagle_pred, eagle_pred_text
    );

    if eagle_pred == target_next {
        println!("MATCH: EAGLE agrees with target on next token.");
    } else {
        println!(
            "MISMATCH: EAGLE draft={} vs target={} (this is fine per-step; check top-5 below)",
            eagle_pred, target_next
        );
    }

    // Show top-5 from eagle logits (in draft vocab space, mapped to target).
    print_top5(
        &eagle_logits,
        "EAGLE draft top-5 (pairing A)",
        &eagle,
        &tokenizer,
    );

    // Alternative pairing B: pair hooks with target_next (the token those hooks produced
    // via lm_head), predict token after target_next. Position advances by 1.
    eagle.reset();
    let (eagle_pred_b, eagle_logits_b) = eagle.step(&hooks, embed_tokens, target_next, pos + 1);
    let eagle_pred_b_text = tokenizer.decode(&[eagle_pred_b]);
    println!(
        "\nEAGLE draft prediction (pairing B: prev=target_next): {} ({:?})",
        eagle_pred_b, eagle_pred_b_text
    );
    print_top5(
        &eagle_logits_b,
        "EAGLE draft top-5 (pairing B)",
        &eagle,
        &tokenizer,
    );
}

fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
    let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).unwrap();
    buf.copy_from_host(bytes).unwrap();
    buf
}

fn print_top5(logits: &Tensor, label: &str, eagle: &Eagle3Head, tokenizer: &Tokenizer) {
    use half::bf16;
    let cpu = logits.to_device(Device::Cpu);
    let data = cpu.as_slice::<bf16>();
    let mut vals: Vec<(usize, f32)> = data
        .iter()
        .enumerate()
        .map(|(i, v)| (i, v.to_f32()))
        .collect();
    vals.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
    println!("{label}:");
    for (i, val) in vals.iter().take(5) {
        let target_id = eagle.map_draft_to_target(*i as u32);
        let text = tokenizer.decode(&[target_id]);
        println!(
            "  draft_id={} target_id={} val={:.3} text={:?}",
            i, target_id, val, text
        );
    }
}