//! Speed and identity harness for Cohere2-MoE (North Mini Code)
//! decode: one prompt from a file, batched prefill, greedy
//! decode, and per step the token id, a hash of every logit's
//! bits (so two builds can be compared bit for bit, not only
//! token for token) and the step's wall-clock.
//!
//! ```bash
//! cargo run --release --features rayon --example north_bench -- \
//!     /path/to/snapshot-dir prompt.txt [n_tokens] [--chat]
//! ```
//!
//! `--chat` wraps the file's text as one user turn of the
//! model's chat template, the shape the program-writing loop
//! feeds. `--then <file>` afterwards feeds that file's text as a
//! further user turn into the same context with
//! `forward_extend` (`--then-seq <file>`: one `forward_token` per
//! token, the same logits) and greedily decodes 8 more tokens.
//! With `JOULECLAW_PROF=1` it also prints where the time went, op
//! by op, for the prefill and per decoded token.

use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::cohere2_moe::{prof, Cohere2MoeConfig, Cohere2MoeModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use std::path::PathBuf;
use std::time::Instant;

/// FNV-1a over the bit patterns of every logit.
fn logits_hash(logits: &[f32]) -> u64 {
    let mut h = 0xcbf2_9ce4_8422_2325u64;
    for v in logits {
        for b in v.to_bits().to_le_bytes() {
            h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3);
        }
    }
    h
}

fn print_prof(label: &str, per: f64) {
    if !prof::enabled() {
        return;
    }
    let rows = prof::take();
    let total: f64 = rows.iter().map(|r| r.1).sum();
    println!("profile {label} (ms{}):", if per > 1.0 { " per token" } else { "" });
    for (name, ms) in rows.iter().filter(|r| r.1 > 0.0) {
        println!("  {:<34} {:>10.2}  {:>5.1}%", name, ms / per, 100.0 * ms / total.max(1e-9));
    }
    println!("  {:<34} {:>10.2}", "total", total / per);
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 3 {
        eprintln!("usage: north_bench <snapshot-dir> <prompt-file> [n_tokens] [--chat]");
        std::process::exit(1);
    }
    let dir = PathBuf::from(&args[1]);
    let text = std::fs::read_to_string(&args[2])?;
    let n_tokens: usize = args.get(3).map(|s| s.parse()).transpose()?.unwrap_or(64);
    let chat = args.iter().any(|a| a == "--chat");
    let then = |flag: &str| args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).cloned();
    let (then_file, then_seq) = match (then("--then"), then("--then-seq")) {
        (Some(f), _) => (Some(f), false),
        (None, Some(f)) => (Some(f), true),
        (None, None) => (None, false),
    };

    let config: Cohere2MoeConfig = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
    let mut shard_paths: Vec<PathBuf> = std::fs::read_dir(&dir)?
        .filter_map(|e| e.ok().map(|e| e.path()))
        .filter(|p| {
            p.extension().is_some_and(|e| e == "safetensors")
                && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("model"))
        })
        .collect();
    shard_paths.sort();
    let shards = shard_paths
        .iter()
        .map(|p| read_safetensors_file(p).map_err(|e| format!("{p:?}: {e:?}")))
        .collect::<Result<Vec<_>, _>>()?;
    let model = Cohere2MoeModel::bind_sharded(config, shards)?;
    let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json"))?;

    let body = if chat {
        format!(
            "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
            text.trim_end()
        )
    } else {
        text
    };
    let mut prompt = vec![2u32]; // bos_token_id
    prompt.extend(tokenizer.encode(&body));
    println!("prompt_tokens {}", prompt.len());

    let mut state = model.new_state();
    let t_pre = Instant::now();
    let mut logits = model.forward_prompt(&prompt, &mut state);
    let pre = t_pre.elapsed().as_secs_f64();
    println!(
        "prefill_ms {:.0} tok_s {:.2}",
        pre * 1e3,
        prompt.len() as f64 / pre
    );
    print_prof("prefill", 1.0);

    let mut ids = Vec::with_capacity(n_tokens);
    let mut hashes = Vec::with_capacity(n_tokens);
    let mut ms = Vec::with_capacity(n_tokens);
    for step in 0..n_tokens {
        let next = logits
            .iter()
            .enumerate()
            .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| if v > bv { (i, v) } else { (bi, bv) })
            .0 as u32;
        let hash = logits_hash(&logits);
        ids.push(next);
        hashes.push(hash);
        if step + 1 == n_tokens {
            println!("step {step:>3} id {next:>6} hash {hash:016x}");
            break;
        }
        let t = Instant::now();
        logits = model.forward_token(next, &mut state);
        let dt = t.elapsed().as_secs_f64() * 1e3;
        ms.push(dt);
        println!("step {step:>3} id {next:>6} hash {hash:016x} ms {dt:.1}");
    }
    let mut sorted = ms.clone();
    sorted.sort_by(f64::total_cmp);
    let median = sorted.get(sorted.len() / 2).copied().unwrap_or(0.0);
    let mean = ms.iter().sum::<f64>() / ms.len().max(1) as f64;
    println!(
        "decode_ms median {median:.1} mean {mean:.1} over {} steps, context {}..{}",
        ms.len(),
        prompt.len(),
        state.position
    );
    print_prof("decode", ms.len().max(1) as f64);
    println!("ids {}", ids.iter().map(|i| i.to_string()).collect::<Vec<_>>().join(","));
    println!("hashes {}", hashes.iter().map(|h| format!("{h:016x}")).collect::<Vec<_>>().join(","));
    println!("text {:?}", tokenizer.decode(&ids));

    if let Some(file) = then_file {
        let more = std::fs::read_to_string(&file)?;
        let mut feed = vec![*ids.last().unwrap()];
        feed.extend(tokenizer.encode(&format!(
            "<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>{}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
            more.trim_end()
        )));
        let t = Instant::now();
        logits = if then_seq {
            let mut l = Vec::new();
            for &x in &feed {
                l = model.forward_token(x, &mut state);
            }
            l
        } else {
            model.forward_extend(&feed, &mut state)
        };
        let dt = t.elapsed().as_secs_f64();
        println!(
            "then_{} {} tokens in {:.0} ms ({:.2} tok/s), context {}, hash {:016x}",
            if then_seq { "seq" } else { "extend" },
            feed.len(),
            dt * 1e3,
            feed.len() as f64 / dt,
            state.position,
            logits_hash(&logits)
        );
        let mut more_ids = Vec::new();
        for _ in 0..8 {
            let next = logits
                .iter()
                .enumerate()
                .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| if v > bv { (i, v) } else { (bi, bv) })
                .0 as u32;
            more_ids.push(next);
            logits = model.forward_token(next, &mut state);
        }
        println!("then_ids {} hash {:016x}", more_ids.iter().map(|i| i.to_string()).collect::<Vec<_>>().join(","), logits_hash(&logits));
    }
    Ok(())
}