//! Weight export: converts trained models into standalone Rust `const` arrays
//! and optionally Lean 4 definitions.
//!
//! The generated Rust source is meant to be placed in
//! `src/ensemble/gen/{scanner,ddos}_weights.rs` so the inference side can use
//! compile-time weight constants with zero runtime cost.

use anyhow::Result;
use std::fmt::Write as FmtWrite;
use std::path::Path;

/// All data needed to emit a standalone inference source file.
pub struct ExportedModel {
    /// Module name used in generated code comments and Lean defs.
    pub model_name: String,
    /// Number of input features.
    pub input_dim: usize,
    /// Hidden layer width (always 32 in the current ensemble).
    pub hidden_dim: usize,
    /// Weight matrix for layer 1: `hidden_dim x input_dim`.
    pub w1: Vec<Vec<f32>>,
    /// Bias vector for layer 1: length `hidden_dim`.
    pub b1: Vec<f32>,
    /// Weight vector for layer 2: length `hidden_dim`.
    pub w2: Vec<f32>,
    /// Bias scalar for layer 2.
    pub b2: f32,
    /// Packed decision tree nodes.
    pub tree_nodes: Vec<(u8, f32, u16, u16)>,
    /// MLP classification threshold.
    pub threshold: f32,
    /// Per-feature normalization minimums.
    pub norm_mins: Vec<f32>,
    /// Per-feature normalization maximums.
    pub norm_maxs: Vec<f32>,
}

/// Generate a Rust source file with `const` arrays for all model weights.
pub fn generate_rust_source(model: &ExportedModel) -> String {
    let mut s = String::with_capacity(8192);

    writeln!(
        s,
        "//! Auto-generated weights for the {} ensemble.",
        model.model_name
    )
    .unwrap();
    writeln!(
        s,
        "//! DO NOT EDIT — regenerate with `cargo run --features training -- train-{}-mlp`.",
        model.model_name.to_ascii_lowercase()
    )
    .unwrap();
    writeln!(s).unwrap();

    // Threshold.
    writeln!(s, "pub const THRESHOLD: f32 = {:.8};", model.threshold).unwrap();
    writeln!(s).unwrap();

    // Normalization params.
    write_f32_array(&mut s, "NORM_MINS", &model.norm_mins);
    write_f32_array(&mut s, "NORM_MAXS", &model.norm_maxs);

    // W1: hidden_dim x input_dim.
    writeln!(
        s,
        "pub const W1: [[f32; {}]; {}] = [",
        model.input_dim, model.hidden_dim
    )
    .unwrap();
    for row in &model.w1 {
        write!(s, "    [").unwrap();
        for (i, v) in row.iter().enumerate() {
            if i > 0 {
                write!(s, ", ").unwrap();
            }
            write!(s, "{:.8}", v).unwrap();
        }
        writeln!(s, "],").unwrap();
    }
    writeln!(s, "];").unwrap();
    writeln!(s).unwrap();

    // B1.
    write_f32_array(&mut s, "B1", &model.b1);

    // W2.
    write_f32_array(&mut s, "W2", &model.w2);

    // B2.
    writeln!(s, "pub const B2: f32 = {:.8};", model.b2).unwrap();
    writeln!(s).unwrap();

    // Tree nodes.
    writeln!(
        s,
        "pub const TREE_NODES: [(u8, f32, u16, u16); {}] = [",
        model.tree_nodes.len()
    )
    .unwrap();
    for &(feat, thresh, left, right) in &model.tree_nodes {
        writeln!(
            s,
            "    ({}, {:.8}, {}, {}),",
            feat, thresh, left, right
        )
        .unwrap();
    }
    writeln!(s, "];").unwrap();

    s
}

/// Generate Lean 4 definitions for formal verification.
pub fn generate_lean_source(model: &ExportedModel) -> String {
    let mut s = String::with_capacity(8192);

    writeln!(
        s,
        "-- Auto-generated Lean 4 definitions for {} ensemble.",
        model.model_name
    )
    .unwrap();
    writeln!(
        s,
        "-- DO NOT EDIT — regenerate with the training pipeline."
    )
    .unwrap();
    writeln!(s).unwrap();

    writeln!(
        s,
        "namespace Sunbeam.Ensemble.{}",
        capitalize(&model.model_name)
    )
    .unwrap();
    writeln!(s).unwrap();

    writeln!(s, "def inputDim : Nat := {}", model.input_dim).unwrap();
    writeln!(s, "def hiddenDim : Nat := {}", model.hidden_dim).unwrap();
    writeln!(s, "def threshold : Float := {:.8}", model.threshold).unwrap();
    writeln!(s, "def b2 : Float := {:.8}", model.b2).unwrap();
    writeln!(s).unwrap();

    write_lean_float_list(&mut s, "normMins", &model.norm_mins);
    write_lean_float_list(&mut s, "normMaxs", &model.norm_maxs);
    write_lean_float_list(&mut s, "b1", &model.b1);
    write_lean_float_list(&mut s, "w2", &model.w2);

    // W1 as list of lists.
    writeln!(s, "def w1 : List (List Float) := [").unwrap();
    for (i, row) in model.w1.iter().enumerate() {
        let comma = if i + 1 < model.w1.len() { "," } else { "" };
        write!(s, "  [").unwrap();
        for (j, v) in row.iter().enumerate() {
            if j > 0 {
                write!(s, ", ").unwrap();
            }
            write!(s, "{:.8}", v).unwrap();
        }
        writeln!(s, "]{}", comma).unwrap();
    }
    writeln!(s, "]").unwrap();
    writeln!(s).unwrap();

    // Tree nodes as list of tuples.
    writeln!(
        s,
        "def treeNodes : List (Nat × Float × Nat × Nat) := ["
    )
    .unwrap();
    for (i, &(feat, thresh, left, right)) in model.tree_nodes.iter().enumerate() {
        let comma = if i + 1 < model.tree_nodes.len() {
            ","
        } else {
            ""
        };
        writeln!(
            s,
            "  ({}, {:.8}, {}, {}){}",
            feat, thresh, left, right, comma
        )
        .unwrap();
    }
    writeln!(s, "]").unwrap();
    writeln!(s).unwrap();

    writeln!(
        s,
        "end Sunbeam.Ensemble.{}",
        capitalize(&model.model_name)
    )
    .unwrap();

    s
}

/// Write the generated Rust source to a file.
pub fn export_to_file(model: &ExportedModel, path: &Path) -> Result<()> {
    let source = generate_rust_source(model);
    std::fs::write(path, source.as_bytes())
        .map_err(|e| anyhow::anyhow!("writing export to {}: {}", path.display(), e))?;
    Ok(())
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn write_f32_array(s: &mut String, name: &str, values: &[f32]) {
    writeln!(s, "pub const {}: [f32; {}] = [", name, values.len()).unwrap();
    write!(s, "    ").unwrap();
    for (i, v) in values.iter().enumerate() {
        if i > 0 {
            write!(s, ", ").unwrap();
        }
        // Line-wrap every 8 values for readability.
        if i > 0 && i % 8 == 0 {
            write!(s, "\n    ").unwrap();
        }
        write!(s, "{:.8}", v).unwrap();
    }
    writeln!(s, "\n];").unwrap();
    writeln!(s).unwrap();
}

fn write_lean_float_list(s: &mut String, name: &str, values: &[f32]) {
    write!(s, "def {} : List Float := [", name).unwrap();
    for (i, v) in values.iter().enumerate() {
        if i > 0 {
            write!(s, ", ").unwrap();
        }
        write!(s, "{:.8}", v).unwrap();
    }
    writeln!(s, "]").unwrap();
    writeln!(s).unwrap();
}

fn capitalize(s: &str) -> String {
    let mut c = s.chars();
    match c.next() {
        None => String::new(),
        Some(first) => first.to_uppercase().to_string() + c.as_str(),
    }
}

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

    fn make_test_model() -> ExportedModel {
        ExportedModel {
            model_name: "scanner".to_string(),
            input_dim: 2,
            hidden_dim: 2,
            w1: vec![vec![0.1, 0.2], vec![0.3, 0.4]],
            b1: vec![0.01, 0.02],
            w2: vec![0.5, 0.6],
            b2: -0.1,
            tree_nodes: vec![
                (0, 0.5, 1, 2),
                (255, 0.0, 0, 0),
                (255, 1.0, 0, 0),
            ],
            threshold: 0.5,
            norm_mins: vec![0.0, 0.0],
            norm_maxs: vec![1.0, 10.0],
        }
    }

    #[test]
    fn test_rust_source_contains_consts() {
        let model = make_test_model();
        let src = generate_rust_source(&model);

        assert!(src.contains("pub const THRESHOLD: f32 ="), "missing THRESHOLD");
        assert!(src.contains("pub const NORM_MINS:"), "missing NORM_MINS");
        assert!(src.contains("pub const NORM_MAXS:"), "missing NORM_MAXS");
        assert!(src.contains("pub const W1:"), "missing W1");
        assert!(src.contains("pub const B1:"), "missing B1");
        assert!(src.contains("pub const W2:"), "missing W2");
        assert!(src.contains("pub const B2: f32 ="), "missing B2");
        assert!(src.contains("pub const TREE_NODES:"), "missing TREE_NODES");
    }

    #[test]
    fn test_rust_source_array_dims() {
        let model = make_test_model();
        let src = generate_rust_source(&model);

        // W1 should be [f32; 2]; 2]
        assert!(src.contains("[[f32; 2]; 2]"), "W1 dimensions wrong");
        // B1 should be [f32; 2]
        assert!(src.contains("B1: [f32; 2]"), "B1 dimensions wrong");
        // W2 should be [f32; 2]
        assert!(src.contains("W2: [f32; 2]"), "W2 dimensions wrong");
        // TREE_NODES should have 3 entries
        assert!(
            src.contains("[(u8, f32, u16, u16); 3]"),
            "TREE_NODES count wrong"
        );
    }

    #[test]
    fn test_weight_values_roundtrip() {
        let model = make_test_model();
        let src = generate_rust_source(&model);

        // The threshold should appear with reasonable precision.
        assert!(src.contains("0.50000000"), "threshold value missing");
        // B2 value.
        assert!(src.contains("-0.10000000"), "b2 value missing");
    }

    #[test]
    fn test_lean_source_structure() {
        let model = make_test_model();
        let src = generate_lean_source(&model);

        assert!(src.contains("namespace Sunbeam.Ensemble.Scanner"));
        assert!(src.contains("def inputDim : Nat := 2"));
        assert!(src.contains("def hiddenDim : Nat := 2"));
        assert!(src.contains("def threshold : Float :="));
        assert!(src.contains("def normMins : List Float :="));
        assert!(src.contains("def w1 : List (List Float) :="));
        assert!(src.contains("def treeNodes :"));
        assert!(src.contains("end Sunbeam.Ensemble.Scanner"));
    }

    #[test]
    fn test_export_to_file() {
        let model = make_test_model();
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test_weights.rs");

        export_to_file(&model, &path).unwrap();

        let content = std::fs::read_to_string(&path).unwrap();
        assert!(content.contains("pub const THRESHOLD:"));
        assert!(content.contains("pub const W1:"));
    }
}