//! Real-model bring-up for safetensors files. Mirror of
//! `examples/load_model.rs` for the safetensors format.
//!
//! ```bash
//! cargo run --release --example load_safetensors -- path/to/model.safetensors
//! ```

use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: load_safetensors <path.safetensors>");
        std::process::exit(1);
    }
    let path = PathBuf::from(&args[1]);
    println!("model:  {}", path.display());
    let t0 = Instant::now();
    let m = read_safetensors_file(&path).map_err(|e| format!("parse: {:?}", e))?;
    let load_secs = t0.elapsed().as_secs_f64();

    println!();
    println!("=== load ===");
    println!("  load_wall_clock_secs: {:.3}", load_secs);
    println!("  tensor_count: {}", m.tensors.len());
    println!("  metadata_keys: {}", m.metadata.len());
    println!("  mmap_len: {} ({})", m.mmap.len(), humanize_bytes(m.mmap.len() as u64));
    println!("  data_offset: {}", m.data_offset);

    let total_data: u64 = m.tensors.iter().map(|t| t.data_len).sum();
    println!();
    println!("=== tensor data ===");
    println!("  total_bytes: {} ({})", total_data, humanize_bytes(total_data));

    // Dtype distribution.
    let mut counts: BTreeMap<String, (u32, u64)> = BTreeMap::new();
    for t in &m.tensors {
        let key = format!("{:?}", t.dtype);
        let e = counts.entry(key).or_insert((0, 0));
        e.0 += 1;
        e.1 += t.data_len;
    }
    println!();
    println!("=== dtype distribution ===");
    for (k, (c, b)) in &counts {
        println!("  {:<10} count={:>5}  size={}", k, c, humanize_bytes(*b));
    }

    // Top tensors by size.
    let mut sorted: Vec<&_> = m.tensors.iter().collect();
    sorted.sort_by_key(|a| std::cmp::Reverse(a.data_len));
    println!();
    println!("=== top 10 tensors by byte size ===");
    for t in sorted.iter().take(10) {
        println!(
            "  {:<60} {:>14} bytes ({})",
            truncate(&t.name, 60),
            t.data_len,
            humanize_bytes(t.data_len),
        );
    }

    // Metadata.
    if !m.metadata.is_empty() {
        println!();
        println!("=== metadata ===");
        for (k, v) in &m.metadata {
            let v_display = if v.len() > 80 { format!("{}…", &v[..80]) } else { v.clone() };
            println!("  {}: {}", k, v_display);
        }
    }
    Ok(())
}

fn truncate(s: &str, max: usize) -> String {
    if s.len() <= max {
        s.to_string()
    } else {
        format!("…{}", &s[s.len() - max + 1..])
    }
}

fn humanize_bytes(n: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;
    if n >= GB {
        format!("{:.2} GiB", n as f64 / GB as f64)
    } else if n >= MB {
        format!("{:.2} MiB", n as f64 / MB as f64)
    } else if n >= KB {
        format!("{:.2} KiB", n as f64 / KB as f64)
    } else {
        format!("{} B", n)
    }
}