//! Load any GGUF (incl. quantized: Q8_0 / Q4_K / …) and generate,
//! reporting decode tok/s. Usage:
//!   cargo run --release -p jouleclaw-runtime --example gguf_gen -- \
//!       /path/model.gguf "prompt" [n_tokens]
use jouleclaw_loader_gguf::read_gguf_file;
use jouleclaw_loader_gguf::sample::SamplingConfig;
use jouleclaw_loader_gguf::tokenizer::Vocab;
use jouleclaw_runtime::generate::{GenerateConfig, KvCacheKind, TokenizerKind};
use jouleclaw_runtime::streaming::Conversation;
use jouleclaw_runtime::Runtime;
use std::time::Instant;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: gguf_gen <model.gguf> [prompt] [n_tokens]");
        std::process::exit(1);
    }
    let path = &args[1];
    let prompt = args.get(2).map(|s| s.as_str()).unwrap_or("The capital of France is");
    let n: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(32);

    let t0 = Instant::now();
    let model = read_gguf_file(path).expect("read_gguf_file");
    let vocab = Vocab::from_gguf(&model).expect("Vocab::from_gguf");
    let arch = model.metadata.get("general.architecture")
        .and_then(|v| v.as_string()).unwrap_or("?").to_string();
    let ftype = model.metadata.get("general.file_type")
        .map(|v| format!("{v:?}")).unwrap_or_default();
    eprintln!("arch={arch} file_type={ftype} tensors={} | loaded in {:?}",
        model.tensors.len(), t0.elapsed());

    let rt = if std::env::var("REF").is_ok() { Runtime::reference_only() } else { Runtime::boot() };
    let mut conv = Conversation::with_runtime(&model, &vocab, (n + 64).max(128), rt)
        .expect("Conversation::with_runtime");
    let cfg = GenerateConfig {
        max_new_tokens: n,
        add_bos: true,
        tokenizer_kind: TokenizerKind::Auto,
        cache_kind: KvCacheKind::InPlace,
        sampling: SamplingConfig::greedy(),
        max_seq: Some((n + 64).max(128)),
        stop_strings: vec![],
    };
    let t = Instant::now();
    let mut tokens: Vec<u32> = Vec::new();
    let mut step_ms: Vec<f64> = Vec::new();
    let mut last = Instant::now();
    for st in conv.extend(prompt, &cfg).expect("extend") {
        tokens.push(st.expect("decode step").id);
        step_ms.push(last.elapsed().as_secs_f64() * 1000.0);
        last = Instant::now();
    }
    let dt = t.elapsed();
    // Warm steady-state = median of the last half (page faults + cold
    // first steps live in the head).
    if step_ms.len() >= 4 {
        let mut tail: Vec<f64> = step_ms[step_ms.len() / 2..].to_vec();
        tail.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let med = tail[tail.len() / 2];
        println!("per-step ms: {:?}", step_ms.iter().map(|m| *m as u64).collect::<Vec<_>>());
        println!("WARM median: {:.1} ms/tok = {:.1} tok/s", med, 1000.0 / med);
    }
    jouleclaw_runtime::print_op_profile();
    println!("\n>>> {}{}", prompt, vocab.decode_bpe(&tokens));
    let per = dt.as_secs_f64() * 1000.0 / tokens.len().max(1) as f64;
    println!("\n{} tokens in {:?} = {:.1} ms/tok = {:.1} tok/s",
        tokens.len(), dt, per, 1000.0 / per);
}