//! Test utilities: synthetic GGUF model generator.
//!
//! Creates minimal but valid GGUF files for testing the full inference pipeline
//! without requiring multi-GB model downloads.

use byteorder::{LittleEndian, WriteBytesExt};
use std::io::Write;
use std::path::Path;

/// Tiny test model config: 2 layers, 64-dim, 4 heads, 128 vocab.
pub struct TestModelConfig {
    pub hidden_size: u32,
    pub num_heads: u32,
    pub num_kv_heads: u32,
    pub num_layers: u32,
    pub vocab_size: u32,
    pub intermediate_size: u32,
    pub max_seq_len: u32,
}

impl Default for TestModelConfig {
    fn default() -> Self {
        Self {
            hidden_size: 64,
            num_heads: 4,
            num_kv_heads: 2,
            vocab_size: 128,
            intermediate_size: 128,
            num_layers: 2,
            max_seq_len: 256,
        }
    }
}

/// Generate a minimal valid GGUF v3 file with F32 weights.
///
/// All weights are random f32 in [-0.1, 0.1] (Xavier-like).
/// Returns the path to the generated file.
pub fn generate_test_gguf(dir: &Path, config: &TestModelConfig) -> std::path::PathBuf {
    let path = dir.join("test_model.gguf");
    let mut buf: Vec<u8> = Vec::new();

    let magic: u32 = 0x46554747; // "GGUF" as little-endian u32
    buf.write_u32::<LittleEndian>(magic).unwrap();
    buf.write_u32::<LittleEndian>(3).unwrap(); // version 3

    // Count tensors: token_embd + per-layer (9 tensors) + output_norm + output
    let tensors_per_layer: u32 = 9;
    let tensor_count = 1 + (config.num_layers * tensors_per_layer) + 2;
    buf.write_u64::<LittleEndian>(tensor_count as u64).unwrap();

    // Metadata KV count
    let metadata_count: u64 = 8;
    buf.write_u64::<LittleEndian>(metadata_count).unwrap();

    // --- Metadata ---
    write_string_kv(&mut buf, "general.architecture", "llama");
    write_string_kv(&mut buf, "general.name", "llmhosts-test-model");
    write_u32_kv(&mut buf, "llama.embedding_length", config.hidden_size);
    write_u32_kv(&mut buf, "llama.attention.head_count", config.num_heads);
    write_u32_kv(&mut buf, "llama.attention.head_count_kv", config.num_kv_heads);
    write_u32_kv(&mut buf, "llama.block_count", config.num_layers);
    write_u32_kv(&mut buf, "llama.feed_forward_length", config.intermediate_size);
    write_u32_kv(&mut buf, "llama.context_length", config.max_seq_len);

    // --- Tensor descriptors ---
    let h = config.hidden_size as u64;
    let v = config.vocab_size as u64;
    let ff = config.intermediate_size as u64;
    let head_dim = h / config.num_heads as u64;
    let kv_dim = head_dim * config.num_kv_heads as u64;

    let mut offset: u64 = 0;
    let mut tensor_specs: Vec<(String, Vec<u64>, u64)> = Vec::new(); // (name, dims, byte_size)

    // token_embd.weight: [hidden_size, vocab_size] (GGUF column-major convention)
    let sz = v * h * 4;
    write_tensor_info(&mut buf, "token_embd.weight", &[h, v], 0, offset);
    tensor_specs.push(("token_embd.weight".into(), vec![v, h], sz));
    offset += sz;

    for layer in 0..config.num_layers {
        // GGUF column-major convention: dims are [in_features, out_features]
        let names_dims_sizes: Vec<(String, Vec<u64>, u64)> = vec![
            (format!("blk.{layer}.attn_norm.weight"), vec![h], h * 4),
            (format!("blk.{layer}.attn_q.weight"), vec![h, h], h * h * 4),
            (format!("blk.{layer}.attn_k.weight"), vec![h, kv_dim], kv_dim * h * 4),
            (format!("blk.{layer}.attn_v.weight"), vec![h, kv_dim], kv_dim * h * 4),
            (format!("blk.{layer}.attn_output.weight"), vec![h, h], h * h * 4),
            (format!("blk.{layer}.ffn_norm.weight"), vec![h], h * 4),
            (format!("blk.{layer}.ffn_gate.weight"), vec![h, ff], ff * h * 4),
            (format!("blk.{layer}.ffn_up.weight"), vec![h, ff], ff * h * 4),
            (format!("blk.{layer}.ffn_down.weight"), vec![ff, h], h * ff * 4),
        ];
        for (name, dims, byte_size) in &names_dims_sizes {
            write_tensor_info(&mut buf, name, dims, 0, offset);
            tensor_specs.push((name.clone(), dims.clone(), *byte_size));
            offset += byte_size;
        }
    }

    // output_norm + output
    write_tensor_info(&mut buf, "output_norm.weight", &[h], 0, offset);
    tensor_specs.push(("output_norm.weight".into(), vec![h], h * 4));
    offset += h * 4;
    write_tensor_info(&mut buf, "output.weight", &[h, v], 0, offset);
    tensor_specs.push(("output.weight".into(), vec![h, v], v * h * 4));

    // --- Align to 32 bytes ---
    let current_len = buf.len();
    let aligned = current_len.div_ceil(32) * 32;
    buf.resize(aligned, 0);

    // --- Write weight data (random f32) ---
    use rand::Rng;
    let mut rng = rand::thread_rng();
    for (_name, _dims, byte_size) in &tensor_specs {
        let n_floats = *byte_size as usize / 4;
        for _ in 0..n_floats {
            let v: f32 = rng.gen_range(-0.1..0.1);
            buf.write_f32::<LittleEndian>(v).unwrap();
        }
    }

    std::fs::write(&path, &buf).unwrap();
    path
}

fn write_string_kv(buf: &mut Vec<u8>, key: &str, value: &str) {
    buf.write_u64::<LittleEndian>(key.len() as u64).unwrap();
    buf.write_all(key.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(8).unwrap(); // String type
    buf.write_u64::<LittleEndian>(value.len() as u64).unwrap();
    buf.write_all(value.as_bytes()).unwrap();
}

fn write_u32_kv(buf: &mut Vec<u8>, key: &str, value: u32) {
    buf.write_u64::<LittleEndian>(key.len() as u64).unwrap();
    buf.write_all(key.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(4).unwrap(); // U32 type
    buf.write_u32::<LittleEndian>(value).unwrap();
}

fn write_tensor_info(buf: &mut Vec<u8>, name: &str, dims: &[u64], dtype: u32, offset: u64) {
    buf.write_u64::<LittleEndian>(name.len() as u64).unwrap();
    buf.write_all(name.as_bytes()).unwrap();
    buf.write_u32::<LittleEndian>(dims.len() as u32).unwrap();
    for d in dims {
        buf.write_u64::<LittleEndian>(*d).unwrap();
    }
    buf.write_u32::<LittleEndian>(dtype).unwrap(); // GgmlType
    buf.write_u64::<LittleEndian>(offset).unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::gguf;

    #[test]
    fn test_generate_test_gguf_is_parseable() {
        let dir = tempfile::tempdir().unwrap();
        let config = TestModelConfig::default();
        let path = generate_test_gguf(dir.path(), &config);

        let header = gguf::parse_header(&path).unwrap();
        assert_eq!(header.version, 3);
        assert_eq!(gguf::get_string(&header, "general.architecture"), Some("llama".into()));
        assert_eq!(gguf::get_u32(&header, "llama.embedding_length"), Some(64));
        assert_eq!(gguf::get_u32(&header, "llama.block_count"), Some(2));
        assert!(!header.tensors.is_empty());

        let embd = header.tensors.iter().find(|t| t.name == "token_embd.weight");
        assert!(embd.is_some());
        let embd = embd.unwrap();
        assert_eq!(embd.dimensions, vec![64, 128]); // GGUF col-major: [hidden, vocab]
    }

    #[test]
    fn test_generate_test_gguf_weights_are_loadable() {
        let dir = tempfile::tempdir().unwrap();
        let config = TestModelConfig::default();
        let path = generate_test_gguf(dir.path(), &config);

        let header = gguf::parse_header(&path).unwrap();
        let mapped = gguf::MappedWeights::open(&path, header.data_offset).unwrap();

        let embd = header.tensors.iter().find(|t| t.name == "token_embd.weight").unwrap();
        let n_elements = embd.dimensions.iter().product::<u64>() as usize;
        let byte_len = n_elements * 4;
        let bytes = mapped.tensor_bytes(embd.offset, byte_len).unwrap();
        assert_eq!(bytes.len(), byte_len);

        let weights = crate::quant::dequantize(bytes, crate::gguf::GgmlType::F32, n_elements).unwrap();
        assert_eq!(weights.len(), n_elements);
        assert!(weights.iter().all(|w| w.abs() < 0.2));
    }
}
