//! Code generation — produce `weights.rs` for `rskim-search` from the JSON weight table.

use std::io::Write;
use std::path::{Path, PathBuf};

use anyhow::Context;

use crate::extract::bigram_to_display;
use crate::types::WeightTable;

/// Walk up from the current directory to find the workspace root
/// (the directory containing a `Cargo.toml` with `[workspace]`).
pub fn find_workspace_root() -> anyhow::Result<PathBuf> {
    let mut dir = std::env::current_dir().context("getting current directory")?;
    let limit = 32usize; // Reliability: fixed upper bound on directory traversal.
    for _ in 0..limit {
        let cargo_toml = dir.join("Cargo.toml");
        if cargo_toml.exists() {
            let contents = std::fs::read_to_string(&cargo_toml)
                .with_context(|| format!("reading {}", cargo_toml.display()))?;
            if contents.contains("[workspace]") {
                return Ok(dir);
            }
        }
        match dir.parent() {
            Some(parent) => dir = parent.to_path_buf(),
            None => break,
        }
    }
    anyhow::bail!("could not find workspace root (no Cargo.toml with [workspace] found)")
}

/// Read `json_path`, validate the weight table, and write a compilable `weights.rs`
/// to `output_path`.
///
/// # Errors
///
/// Returns an error if:
/// - `json_path` cannot be read or parsed
/// - The weight table is empty or version is 0
/// - Any IDF value is non-positive
/// - Writing to `output_path` fails
pub fn generate_weights_rs(json_path: &Path, output_path: &Path) -> anyhow::Result<()> {
    let raw = std::fs::read_to_string(json_path)
        .with_context(|| format!("reading {}", json_path.display()))?;

    let table: WeightTable = serde_json::from_str(&raw)
        .with_context(|| format!("parsing weight table from {}", json_path.display()))?;

    // Validate table contents.
    if table.version == 0 {
        anyhow::bail!("weight table version must be > 0");
    }
    if table.weights.is_empty() {
        anyhow::bail!("weight table is empty");
    }
    for w in &table.weights {
        if !w.idf.is_finite() || w.idf <= 0.0 {
            anyhow::bail!(
                "invalid IDF {} for bigram 0x{:04X} — all IDF values must be finite and positive",
                w.idf,
                w.bigram
            );
        }
    }

    // Generate the Rust source.
    let source = build_weights_rs(&table)?;

    // Create output directory if needed.
    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("creating directory {}", parent.display()))?;
    }

    std::fs::write(output_path, source)
        .with_context(|| format!("writing {}", output_path.display()))?;

    Ok(())
}

fn build_weights_rs(table: &WeightTable) -> anyhow::Result<String> {
    let mut buf = Vec::with_capacity(64 * 1024);
    write_file_header(&mut buf, table)?;
    write_entries(&mut buf, &table.weights)?;
    write_lookup_fn(&mut buf)?;
    write_generated_tests(&mut buf)?;
    String::from_utf8(buf).context("building weights.rs source (non-UTF8 output)")
}

fn write_file_header(buf: &mut Vec<u8>, table: &WeightTable) -> anyhow::Result<()> {
    writeln!(
        buf,
        "//! Empirical character bigram IDF weight table for rskim-search sparse index."
    )?;
    writeln!(
        buf,
        "//! Generated by rskim-research from data/bigram_weights.json"
    )?;
    writeln!(buf, "//!")?;
    writeln!(
        buf,
        "//! Version: {}  |  Generated: {}",
        table.version, table.generated_at
    )?;
    writeln!(
        buf,
        "//! Corpus: {} files, {} unique bigrams",
        table.corpus_stats.total_files, table.corpus_stats.unique_bigrams
    )?;
    writeln!(buf, "//!")?;
    writeln!(
        buf,
        "//! DO NOT EDIT MANUALLY — re-run `rskim-research codegen` to regenerate."
    )?;
    writeln!(buf)?;
    // Suppress excessive_precision lint — the exact float representation is intentional
    // in machine-generated const data.
    writeln!(buf, "#![allow(clippy::excessive_precision)]")?;
    writeln!(buf)?;
    Ok(())
}

fn write_entries(buf: &mut Vec<u8>, weights: &[crate::types::BigramWeight]) -> anyhow::Result<()> {
    writeln!(buf, "/// Sorted (bigram_key, idf_weight) pairs.")?;
    writeln!(buf, "///")?;
    writeln!(
        buf,
        "/// The table is sorted ascending by bigram key, enabling O(log n) binary search."
    )?;
    writeln!(
        buf,
        "/// Each bigram key encodes two bytes: `key = (byte1 << 8) | byte2`."
    )?;
    writeln!(buf, "pub const BIGRAM_WEIGHTS: &[(u16, f32)] = &[")?;
    for w in weights {
        let display = bigram_to_display(w.bigram);
        writeln!(
            buf,
            "    (0x{:04X}, {:.6}_f32), // {:?}",
            w.bigram, w.idf, display
        )?;
    }
    writeln!(buf, "];")?;
    writeln!(buf)?;
    writeln!(buf, "/// Default weight for bigrams not in the table.")?;
    writeln!(buf, "pub const DEFAULT_WEIGHT: f32 = 1.0;")?;
    writeln!(buf)?;
    Ok(())
}

fn write_lookup_fn(buf: &mut Vec<u8>) -> anyhow::Result<()> {
    writeln!(
        buf,
        "/// Look up the IDF weight of a bigram by binary search."
    )?;
    writeln!(buf, "///")?;
    writeln!(
        buf,
        "/// Returns `None` if the bigram is absent from the table."
    )?;
    writeln!(buf, "#[must_use]")?;
    writeln!(buf, "pub fn bigram_weight(bigram: u16) -> Option<f32> {{")?;
    writeln!(buf, "    BIGRAM_WEIGHTS")?;
    writeln!(buf, "        .binary_search_by_key(&bigram, |&(k, _)| k)")?;
    writeln!(buf, "        .ok()")?;
    writeln!(buf, "        .map(|idx| BIGRAM_WEIGHTS[idx].1)")?;
    writeln!(buf, "}}")?;
    writeln!(buf)?;
    Ok(())
}

fn write_generated_tests(buf: &mut Vec<u8>) -> anyhow::Result<()> {
    writeln!(buf, "#[cfg(test)]")?;
    writeln!(buf, "mod tests {{")?;
    writeln!(buf, "    use super::*;")?;
    writeln!(buf)?;
    writeln!(buf, "    #[test]")?;
    writeln!(buf, "    fn table_is_sorted_ascending() {{")?;
    writeln!(
        buf,
        "        let keys: Vec<u16> = BIGRAM_WEIGHTS.iter().map(|&(k, _)| k).collect();"
    )?;
    writeln!(buf, "        let mut sorted = keys.clone();")?;
    writeln!(buf, "        sorted.sort();")?;
    writeln!(
        buf,
        "        assert_eq!(keys, sorted, \"BIGRAM_WEIGHTS must be sorted ascending\");"
    )?;
    writeln!(buf, "    }}")?;
    writeln!(buf)?;
    writeln!(buf, "    #[test]")?;
    writeln!(buf, "    fn no_duplicate_keys() {{")?;
    writeln!(
        buf,
        "        let keys: Vec<u16> = BIGRAM_WEIGHTS.iter().map(|&(k, _)| k).collect();"
    )?;
    writeln!(
        buf,
        "        let unique: std::collections::HashSet<u16> = keys.iter().copied().collect();"
    )?;
    writeln!(
        buf,
        "        assert_eq!(keys.len(), unique.len(), \"BIGRAM_WEIGHTS must not have duplicate keys\");"
    )?;
    writeln!(buf, "    }}")?;
    writeln!(buf)?;
    writeln!(buf, "    #[test]")?;
    writeln!(buf, "    fn bigram_weight_lookup_works() {{")?;
    writeln!(
        buf,
        "        // If the table is non-empty, the first entry must be findable."
    )?;
    writeln!(
        buf,
        "        if let Some(&(key, idf)) = BIGRAM_WEIGHTS.first() {{"
    )?;
    writeln!(
        buf,
        "            assert_eq!(bigram_weight(key), Some(idf));"
    )?;
    writeln!(buf, "        }}")?;
    writeln!(buf, "    }}")?;
    writeln!(buf)?;
    writeln!(buf, "    #[test]")?;
    writeln!(buf, "    fn default_weight_is_positive() {{")?;
    writeln!(buf, "        const {{ assert!(DEFAULT_WEIGHT > 0.0) }};")?;
    writeln!(buf, "    }}")?;
    writeln!(buf, "}}")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used)]

    use super::*;
    use crate::types::{BigramWeight, CorpusStats, WeightTable};

    fn sample_table() -> WeightTable {
        WeightTable {
            version: 1,
            generated_at: "unix:0".to_string(),
            corpus_stats: CorpusStats {
                total_files: 10,
                total_bigrams: 500,
                unique_bigrams: 5,
                deduplicated_files: 0,
                language_breakdown: vec![],
            },
            weights: vec![
                BigramWeight {
                    bigram: 0x666E, // "fn"
                    idf: 8.5,
                },
                BigramWeight {
                    bigram: 0x696D, // "im"
                    idf: 6.2,
                },
            ],
        }
    }

    #[test]
    fn generate_valid_rust_source() {
        let table = sample_table();
        let source = build_weights_rs(&table).unwrap();

        // Must contain the const definition
        assert!(source.contains("pub const BIGRAM_WEIGHTS: &[(u16, f32)] = &["));
        assert!(source.contains("pub const DEFAULT_WEIGHT: f32 = 1.0;"));
        assert!(source.contains("pub fn bigram_weight(bigram: u16) -> Option<f32>"));
        // Must contain both bigrams
        assert!(source.contains("0x666E"));
        assert!(source.contains("0x696D"));
    }

    #[test]
    fn empty_weights_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let json_path = dir.path().join("weights.json");
        let out_path = dir.path().join("weights.rs");

        let empty_table = WeightTable {
            version: 1,
            generated_at: "unix:0".to_string(),
            corpus_stats: CorpusStats {
                total_files: 0,
                total_bigrams: 0,
                unique_bigrams: 0,
                deduplicated_files: 0,
                language_breakdown: vec![],
            },
            weights: vec![],
        };

        let json = serde_json::to_string(&empty_table).unwrap();
        std::fs::write(&json_path, json).unwrap();

        let err = generate_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(err.to_string().contains("empty"));
    }

    #[test]
    fn invalid_json_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let json_path = dir.path().join("bad.json");
        let out_path = dir.path().join("weights.rs");

        std::fs::write(&json_path, "not json at all").unwrap();

        let err = generate_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(err.to_string().to_lowercase().contains("parsing"));
    }

    #[test]
    fn version_zero_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let json_path = dir.path().join("weights.json");
        let out_path = dir.path().join("weights.rs");

        let table = WeightTable {
            version: 0,
            generated_at: "unix:0".to_string(),
            corpus_stats: CorpusStats {
                total_files: 10,
                total_bigrams: 500,
                unique_bigrams: 1,
                deduplicated_files: 0,
                language_breakdown: vec![],
            },
            weights: vec![BigramWeight {
                bigram: 0x666E,
                idf: 8.5,
            }],
        };

        let json = serde_json::to_string(&table).unwrap();
        std::fs::write(&json_path, json).unwrap();

        let err = generate_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(
            err.to_string().contains("version must be > 0"),
            "expected version error, got: {}",
            err
        );
    }

    #[test]
    fn negative_idf_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let json_path = dir.path().join("weights.json");
        let out_path = dir.path().join("weights.rs");

        let table = WeightTable {
            version: 1,
            generated_at: "unix:0".to_string(),
            corpus_stats: CorpusStats {
                total_files: 10,
                total_bigrams: 500,
                unique_bigrams: 1,
                deduplicated_files: 0,
                language_breakdown: vec![],
            },
            weights: vec![BigramWeight {
                bigram: 0x666E,
                idf: -1.5,
            }],
        };

        let json = serde_json::to_string(&table).unwrap();
        std::fs::write(&json_path, json).unwrap();

        let err = generate_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(
            err.to_string().contains("invalid IDF"),
            "expected invalid IDF error, got: {}",
            err
        );
    }

    #[test]
    fn snapshot_generated_format() {
        let table = sample_table();
        let source = build_weights_rs(&table).unwrap();
        insta::assert_snapshot!(source);
    }
}
