//! Smoke + coherence check for the device-resident llama-family
//! decode driver against a real GGUF (Devstral).
//!
//!   cargo run --release --features fabric \
//!     --example devstral_resident_check -- <model.gguf> [n_gen] [prompt]
//!
//! Coherent continuation = the whole pipeline (embed, interleaved
//! rope, fused attention, fused MLP, lm_head) is correct, since any
//! wiring bug (rope convention, residual, norm) yields garbage.

#[cfg(all(target_os = "macos", feature = "fabric"))]
fn main() {
    use jouleclaw_loader_gguf::devstral_resident::DevstralResident;
    use jouleclaw_loader_gguf::read_gguf_file;
    use jouleclaw_loader_gguf::tokenizer::Vocab;

    let path = std::env::args().nth(1).expect("usage: <model.gguf> [n_gen] [prompt]");
    let n_gen: usize = std::env::args().nth(2).and_then(|s| s.parse().ok()).unwrap_or(24);
    let prompt =
        std::env::args().nth(3).unwrap_or_else(|| "The capital of France is".to_string());

    eprintln!("loading {path} ...");
    let model = read_gguf_file(&path).expect("read gguf");
    let vocab = Vocab::from_gguf(&model).expect("vocab");
    let driver = DevstralResident::bind(&model).expect("bind resident driver");
    let cfg = driver.config();
    eprintln!(
        "arch={} layers={} hidden={} n_q={} n_kv={} hd={} ffn={} vocab={} rope_base={}",
        cfg.arch,
        cfg.block_count,
        cfg.embedding_length,
        cfg.head_count,
        cfg.head_count_kv,
        cfg.head_dim,
        cfg.feed_forward_length,
        cfg.vocab_size,
        cfg.rope_base,
    );

    let mut ids = vocab.encode_bpe_regex(&prompt, true);
    eprintln!("prompt {} tokens", ids.len());
    let cap = ids.len() + n_gen + 8;
    let mut state = driver.new_state(cap).expect("new_state");

    // Prefill A/B with WARM weights (isolates compute from the
    // one-time 14GB upload): warm once, then time both paths fresh.
    if std::env::var("JC_PREFILL_BENCH").is_ok() {
        let mut warm = driver.new_state(cap).expect("state");
        let _ = driver.forward_prompt_batched(&ids, &mut warm).expect("warm");
        let mut s1 = driver.new_state(cap).expect("state");
        let a = std::time::Instant::now();
        driver.forward_prompt_sequential(&ids, &mut s1).expect("seq");
        let seq_ms = a.elapsed().as_secs_f64() * 1e3;
        let mut s2 = driver.new_state(cap).expect("state");
        let b = std::time::Instant::now();
        driver.forward_prompt_batched(&ids, &mut s2).expect("batched");
        let bat_ms = b.elapsed().as_secs_f64() * 1e3;
        let mut s3 = driver.new_state(cap).expect("state");
        let c = std::time::Instant::now();
        driver.forward_prompt_prefill_fused(&ids, &mut s3).expect("fused");
        let fused_ms = c.elapsed().as_secs_f64() * 1e3;
        eprintln!(
            "PREFILL (warm, {} tok): sequential {seq_ms:.0} ms ({:.1} ms/tok), batched {bat_ms:.0} ms, FUSED {fused_ms:.0} ms ({:.1} ms/tok) = {:.2}x vs seq",
            ids.len(), seq_ms / ids.len() as f64, fused_ms / ids.len() as f64, seq_ms / fused_ms.max(1e-3),
        );
        return;
    }

    let mode = std::env::var("JC_PREFILL").unwrap_or_default();
    let t0 = std::time::Instant::now();
    let mut logits = match mode.as_str() {
        "fused" => driver.forward_prompt_prefill_fused(&ids, &mut state).expect("fused prefill"),
        "batched" => driver.forward_prompt_batched(&ids, &mut state).expect("batched prefill"),
        _ => driver.forward_prompt(&ids, &mut state).expect("prefill"),
    };
    let prefill_ms = t0.elapsed().as_secs_f64() * 1e3;
    eprintln!("prefill mode: {}", if mode.is_empty() { "default (fused-fast)" } else { &mode });

    let td = std::time::Instant::now();
    let mut steps = 0usize;
    let mut step_ms: Vec<f64> = Vec::new();
    for _ in 0..n_gen {
        let next = argmax(&logits);
        ids.push(next);
        steps += 1;
        if Some(next) == vocab.eos_id {
            break;
        }
        let s = std::time::Instant::now();
        logits = driver.forward_token(next, &mut state).expect("decode");
        step_ms.push(s.elapsed().as_secs_f64() * 1e3);
    }
    let dec_ms = td.elapsed().as_secs_f64() * 1e3;

    println!("\nOUT: {}", vocab.decode_bpe(&ids));
    let mut sorted = step_ms.clone();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let median = sorted.get(sorted.len() / 2).copied().unwrap_or(0.0);
    let first = step_ms.first().copied().unwrap_or(0.0);
    let last = step_ms.last().copied().unwrap_or(0.0);
    eprintln!(
        "\nprefill {prefill_ms:.0} ms ({} tok), decode avg {:.1} ms/tok over {steps} tok",
        ids.len() - steps,
        dec_ms / steps.max(1) as f64,
    );
    eprintln!(
        "decode WARM median {median:.1} ms/tok = {:.1} tok/s  (first {first:.1}, last {last:.1} — growth = attention scan)",
        1000.0 / median.max(0.001),
    );
}

#[cfg(all(target_os = "macos", feature = "fabric"))]
fn argmax(v: &[f32]) -> u32 {
    let mut bi = 0usize;
    let mut bv = f32::NEG_INFINITY;
    for (i, &x) in v.iter().enumerate() {
        if x > bv {
            bv = x;
            bi = i;
        }
    }
    bi as u32
}

#[cfg(not(all(target_os = "macos", feature = "fabric")))]
fn main() {
    eprintln!("requires macOS + --features fabric");
}