//! Code generation for AST weight tables.
//!
//! Reads an `AstWeightTable` from JSON and generates a compilable `ast_weights.rs`
//! for `rskim-search`, containing per-language bigram and trigram weight arrays
//! with binary-search lookup functions.

use std::io::Write;
use std::path::Path;

use anyhow::Context;

use crate::ast_types::AstWeightTable;
use crate::codegen::find_workspace_root;

/// Read `json_path`, validate the AST weight table, and write a compilable
/// `ast_weights.rs` to `output_path`.
///
/// The generated file contains:
/// - `NODE_KIND_VOCABULARY: &[&str]` — all node kind strings indexed by ID.
/// - Per-language `{LANG}_AST_BIGRAM_WEIGHTS: &[(u32, f32)]` sorted ascending.
/// - Per-language `{LANG}_AST_TRIGRAM_WEIGHTS: &[(u64, f32)]` sorted ascending.
/// - `pub fn ast_bigram_weight(lang: &str, bigram: u32) -> Option<f32>`
/// - `pub fn ast_trigram_weight(lang: &str, trigram: u64) -> Option<f32>`
///
/// # Errors
///
/// Returns an error if:
/// - `json_path` cannot be read or parsed as `AstWeightTable`
/// - The weight table version is 0
/// - Any IDF value is non-positive or non-finite
/// - Writing to `output_path` fails
pub fn generate_ast_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: AstWeightTable = serde_json::from_str(&raw)
        .with_context(|| format!("parsing AST weight table from {}", json_path.display()))?;

    validate_ast_table(&table)?;

    let source = build_ast_weights_rs(&table)?;

    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(())
}

/// Validate the AST weight table contents.
fn validate_ast_table(table: &AstWeightTable) -> anyhow::Result<()> {
    if table.version == 0 {
        anyhow::bail!("AST weight table version must be > 0");
    }
    if table.vocabulary.is_empty() {
        anyhow::bail!("AST weight table vocabulary is empty");
    }
    if table.bigram_weights.is_empty() {
        anyhow::bail!("AST weight table has no bigram weights");
    }
    if table.trigram_weights.is_empty() {
        anyhow::bail!("AST weight table has no trigram weights");
    }

    for (lang, bigrams) in &table.bigram_weights {
        for w in bigrams {
            if !w.idf.is_finite() || w.idf <= 0.0 {
                anyhow::bail!(
                    "invalid IDF {:.6} for bigram 0x{:08X} in language '{}' — all IDF values must be finite and positive",
                    w.idf,
                    w.bigram,
                    lang
                );
            }
        }
    }

    for (lang, trigrams) in &table.trigram_weights {
        for w in trigrams {
            if !w.idf.is_finite() || w.idf <= 0.0 {
                anyhow::bail!(
                    "invalid IDF {:.6} for trigram 0x{:016X} in language '{}' — all IDF values must be finite and positive",
                    w.idf,
                    w.trigram,
                    lang
                );
            }
        }
    }

    Ok(())
}

/// Build the Rust source string for `ast_weights.rs`.
fn build_ast_weights_rs(table: &AstWeightTable) -> anyhow::Result<String> {
    let mut buf: Vec<u8> = Vec::with_capacity(256 * 1024);
    write_ast_file_header(&mut buf, table)?;
    write_vocabulary(&mut buf, &table.vocabulary)?;
    write_language_bigram_arrays(&mut buf, table)?;
    write_language_trigram_arrays(&mut buf, table)?;
    write_bigram_lookup_fn(&mut buf, table)?;
    write_trigram_lookup_fn(&mut buf, table)?;
    write_ast_generated_tests(&mut buf)?;
    String::from_utf8(buf).context("building ast_weights.rs source (non-UTF8 output)")
}

fn write_ast_file_header(buf: &mut Vec<u8>, table: &AstWeightTable) -> anyhow::Result<()> {
    writeln!(buf, "//! Auto-generated by rskim-research ast-codegen.")?;
    writeln!(buf, "//!")?;
    writeln!(
        buf,
        "//! Version: {}  |  Generated: {}",
        table.version, table.generated_at
    )?;
    writeln!(
        buf,
        "//! Vocabulary: {} node kinds  |  Languages: {}",
        table.vocabulary.len(),
        table.bigram_weights.len()
    )?;
    writeln!(buf, "//!")?;
    writeln!(
        buf,
        "//! DO NOT EDIT MANUALLY — re-run `rskim-research ast-codegen` to regenerate."
    )?;
    writeln!(buf)?;
    writeln!(buf, "#![allow(clippy::excessive_precision)]")?;
    writeln!(buf)?;
    Ok(())
}

fn write_vocabulary(buf: &mut Vec<u8>, vocabulary: &[String]) -> anyhow::Result<()> {
    writeln!(buf, "/// All node kind strings, indexed by `NodeKindId`.")?;
    writeln!(buf, "///")?;
    writeln!(
        buf,
        "/// `NODE_KIND_VOCABULARY[id]` resolves a compact ID to its kind string."
    )?;
    writeln!(buf, "pub const NODE_KIND_VOCABULARY: &[&str] = &[")?;
    for kind in vocabulary {
        writeln!(buf, "    {:?},", kind)?;
    }
    writeln!(buf, "];")?;
    writeln!(buf)?;
    Ok(())
}

/// Produce a stable Rust identifier fragment from a language name.
///
/// Examples: `"Rust"` → `"RUST"`, `"TypeScript"` → `"TYPESCRIPT"`,
/// `"Cpp"` → `"CPP"`, `"CSharp"` → `"CSHARP"`, `"C"` → `"C"`.
///
/// Special characters (`+`, `#`, `-`, space) are replaced with `_` and
/// consecutive underscores (runs of any length) are collapsed to a single `_`.
///
/// # Errors
///
/// Returns an error if the resulting identifier is empty, starts with a digit,
/// or contains a character that is not ASCII alphanumeric or `_`. These are
/// hard production errors because the output is written verbatim into generated
/// Rust source files.
fn lang_to_ident(lang: &str) -> anyhow::Result<String> {
    // Semantic mappings for languages whose punctuation-based names would
    // collide after naive character substitution (e.g. "C++" and "C#" both
    // become "C_" with simple replace-and-collapse).
    let result = match lang {
        "C++" => "CPP".to_string(),
        "C#" => "CSHARP".to_string(),
        _ => {
            let mapped: String = lang
                .chars()
                .map(|c| match c {
                    '+' | '#' | '-' | ' ' => '_',
                    _ => c.to_ascii_uppercase(),
                })
                .collect();

            // Collapse runs of consecutive underscores (any length) to a single `_`.
            let mut collapsed = String::with_capacity(mapped.len());
            let mut prev_underscore = false;
            for c in mapped.chars() {
                if c == '_' {
                    if !prev_underscore {
                        collapsed.push(c);
                    }
                    prev_underscore = true;
                } else {
                    collapsed.push(c);
                    prev_underscore = false;
                }
            }
            collapsed
        }
    };

    // Production guard: reject identifiers that would produce invalid Rust code.
    // These are hard errors (not debug_assert!) because lang_to_ident output is
    // written verbatim into generated source files.
    if !result
        .chars()
        .next()
        .map(|c| c.is_ascii_alphabetic() || c == '_')
        .unwrap_or(false)
    {
        anyhow::bail!(
            "lang_to_ident produced an empty or non-identifier-starting string for language \
             {lang:?} (result: {result:?}) — check the language name in the weight table"
        );
    }
    if !result
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '_')
    {
        anyhow::bail!(
            "lang_to_ident produced a non-identifier character for language {lang:?} \
             (result: {result:?}) — only ASCII alphanumeric and '_' are allowed in Rust identifiers"
        );
    }

    Ok(result)
}

/// Collect the language keys from a HashMap, sort them, and resolve each to a
/// Rust identifier via `lang_to_ident`.
///
/// Returns a vector of `(lang_name, ident)` pairs in ascending language-name order,
/// or an error if any name fails identifier validation.
fn sorted_lang_idents<V>(
    map: &std::collections::HashMap<String, V>,
) -> anyhow::Result<Vec<(String, String)>> {
    let mut langs: Vec<&String> = map.keys().collect();
    langs.sort();
    let pairs: Vec<(String, String)> = langs
        .into_iter()
        .map(|lang| lang_to_ident(lang).map(|ident| (lang.clone(), ident)))
        .collect::<anyhow::Result<_>>()?;

    let mut seen = std::collections::HashSet::new();
    for (lang, ident) in &pairs {
        if !seen.insert(ident.as_str()) {
            anyhow::bail!(
                "lang_to_ident collision: language {lang:?} produces ident {ident:?} \
                 which was already used by another language — add a semantic mapping"
            );
        }
    }
    Ok(pairs)
}

/// Write a `binary_search_by_key` lookup function into `buf`.
///
/// Parameters:
/// - `fn_doc`       — first doc-comment line (e.g. `"/// Look up the IDF weight …"`).
/// - `fn_name`      — bare function name (e.g. `"ast_bigram_weight"`).
/// - `param_name`   — key parameter name in the generated signature (`"bigram"` or `"trigram"`).
/// - `key_type`     — Rust type for the key (`"u32"` or `"u64"`).
/// - `const_suffix` — suffix for the const array name (e.g. `"BIGRAM_WEIGHTS"`).
/// - `lang_idents`  — sorted `(lang_name, ident)` pairs.
fn write_lookup_fn(
    buf: &mut Vec<u8>,
    fn_doc: &str,
    fn_name: &str,
    param_name: &str,
    key_type: &str,
    const_suffix: &str,
    lang_idents: &[(String, String)],
) -> anyhow::Result<()> {
    writeln!(buf, "{fn_doc}")?;
    writeln!(buf, "///")?;
    writeln!(
        buf,
        "/// Returns `None` if the language or {param_name} is not in the table."
    )?;
    writeln!(buf, "#[must_use]")?;
    writeln!(
        buf,
        "pub fn {fn_name}(lang: &str, {param_name}: {key_type}) -> Option<f32> {{"
    )?;
    writeln!(buf, "    let table: &[({key_type}, f32)] = match lang {{")?;
    for (lang, ident) in lang_idents {
        writeln!(
            buf,
            "        {:?} => {ident}_AST_{const_suffix},",
            lang.as_str()
        )?;
    }
    writeln!(buf, "        _ => return None,")?;
    writeln!(buf, "    }};")?;
    writeln!(buf, "    table")?;
    writeln!(
        buf,
        "        .binary_search_by_key(&{param_name}, |&(k, _)| k)"
    )?;
    writeln!(buf, "        .ok()")?;
    writeln!(buf, "        .map(|idx| table[idx].1)")?;
    writeln!(buf, "}}")?;
    writeln!(buf)?;
    Ok(())
}

fn write_language_bigram_arrays(buf: &mut Vec<u8>, table: &AstWeightTable) -> anyhow::Result<()> {
    for (lang, ident) in sorted_lang_idents(&table.bigram_weights)? {
        let weights = &table.bigram_weights[&lang];

        // Sort by bigram key ascending for binary search.
        let mut sorted: Vec<_> = weights.iter().collect();
        sorted.sort_by_key(|w| w.bigram);

        writeln!(
            buf,
            "/// {lang} AST bigram weights — sorted ascending by bigram key for O(log n) lookup."
        )?;
        writeln!(
            buf,
            "pub const {ident}_AST_BIGRAM_WEIGHTS: &[(u32, f32)] = &["
        )?;
        for w in &sorted {
            writeln!(
                buf,
                "    (0x{:08X}, {:.6}_f32), // {:?} -> {:?}",
                w.bigram, w.idf, w.parent_kind, w.child_kind
            )?;
        }
        writeln!(buf, "];")?;
        writeln!(buf)?;
    }

    Ok(())
}

fn write_language_trigram_arrays(buf: &mut Vec<u8>, table: &AstWeightTable) -> anyhow::Result<()> {
    for (lang, ident) in sorted_lang_idents(&table.trigram_weights)? {
        let weights = &table.trigram_weights[&lang];

        // Sort by trigram key ascending for binary search.
        let mut sorted: Vec<_> = weights.iter().collect();
        sorted.sort_by_key(|w| w.trigram);

        writeln!(
            buf,
            "/// {lang} AST trigram weights — sorted ascending by trigram key for O(log n) lookup."
        )?;
        writeln!(
            buf,
            "pub const {ident}_AST_TRIGRAM_WEIGHTS: &[(u64, f32)] = &["
        )?;
        for w in &sorted {
            writeln!(
                buf,
                "    (0x{:016X}, {:.6}_f32), // {:?} -> {:?} -> {:?}",
                w.trigram, w.idf, w.grandparent_kind, w.parent_kind, w.child_kind
            )?;
        }
        writeln!(buf, "];")?;
        writeln!(buf)?;
    }

    Ok(())
}

fn write_bigram_lookup_fn(buf: &mut Vec<u8>, table: &AstWeightTable) -> anyhow::Result<()> {
    write_lookup_fn(
        buf,
        "/// Look up the IDF weight of an AST bigram for a given language.",
        "ast_bigram_weight",
        "bigram",
        "u32",
        "BIGRAM_WEIGHTS",
        &sorted_lang_idents(&table.bigram_weights)?,
    )
}

fn write_trigram_lookup_fn(buf: &mut Vec<u8>, table: &AstWeightTable) -> anyhow::Result<()> {
    write_lookup_fn(
        buf,
        "/// Look up the IDF weight of an AST trigram for a given language.",
        "ast_trigram_weight",
        "trigram",
        "u64",
        "TRIGRAM_WEIGHTS",
        &sorted_lang_idents(&table.trigram_weights)?,
    )
}

fn write_ast_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 vocabulary_is_non_empty() {{")?;
    writeln!(
        buf,
        "        assert!(!NODE_KIND_VOCABULARY.is_empty(), \"vocabulary must not be empty\");"
    )?;
    writeln!(buf, "    }}")?;
    writeln!(buf, "}}")?;
    Ok(())
}

/// Default output path for `ast_weights.rs` within the workspace.
pub fn default_ast_weights_output_path() -> anyhow::Result<std::path::PathBuf> {
    let root = find_workspace_root()?;
    Ok(root.join("crates/rskim-search/src/ast_weights.rs"))
}

/// Default input path for `ast_weights.json` within the workspace.
pub fn default_ast_weights_json_path() -> anyhow::Result<std::path::PathBuf> {
    let root = find_workspace_root()?;
    Ok(root.join("crates/rskim-search/data/ast_weights.json"))
}

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

    use std::collections::HashMap;

    use super::*;
    use crate::ast_types::{
        AstBigramWeight, AstCorpusStats, AstLanguageStats, AstTrigramWeight, AstWeightTable,
        encode_ast_bigram, encode_ast_trigram,
    };

    fn sample_table() -> AstWeightTable {
        let mut bigram_weights = HashMap::new();
        bigram_weights.insert(
            "Rust".to_string(),
            vec![
                AstBigramWeight {
                    parent_kind: "source_file".to_string(),
                    child_kind: "function_item".to_string(),
                    bigram: encode_ast_bigram(1, 0),
                    idf: 3.5,
                },
                AstBigramWeight {
                    parent_kind: "function_item".to_string(),
                    child_kind: "identifier".to_string(),
                    bigram: encode_ast_bigram(0, 1),
                    idf: 2.5,
                },
            ],
        );

        let mut trigram_weights = HashMap::new();
        trigram_weights.insert(
            "Rust".to_string(),
            vec![AstTrigramWeight {
                grandparent_kind: "source_file".to_string(),
                parent_kind: "function_item".to_string(),
                child_kind: "identifier".to_string(),
                trigram: encode_ast_trigram(2, 1, 0),
                idf: 4.0,
            }],
        );

        AstWeightTable {
            version: 1,
            generated_at: "unix:0".to_string(),
            vocabulary: vec![
                "function_item".to_string(),
                "identifier".to_string(),
                "source_file".to_string(),
            ],
            corpus_stats: AstCorpusStats {
                total_files: 10,
                deduplicated_files: 0,
                language_stats: vec![AstLanguageStats {
                    language: "Rust".to_string(),
                    file_count: 10,
                    unique_bigrams: 2,
                    unique_trigrams: 1,
                    error_node_count: 0,
                    total_node_count: 100,
                }],
            },
            bigram_weights,
            trigram_weights,
        }
    }

    #[test]
    fn generated_source_contains_vocabulary() {
        let table = sample_table();
        let source = build_ast_weights_rs(&table).unwrap();
        // Structural markers — verifies the const declaration is syntactically valid Rust.
        assert!(
            source.contains("pub const NODE_KIND_VOCABULARY: &[&str]"),
            "missing vocabulary const declaration"
        );
        assert!(
            source.contains("\"function_item\""),
            "vocabulary must contain function_item"
        );
        assert!(
            source.contains("\"identifier\""),
            "vocabulary must contain identifier"
        );
    }

    #[test]
    fn generated_source_contains_language_arrays() {
        let table = sample_table();
        let source = build_ast_weights_rs(&table).unwrap();
        // Structural markers — verifies the const declarations are syntactically valid Rust.
        assert!(
            source.contains("pub const RUST_AST_BIGRAM_WEIGHTS: &[(u32, f32)]"),
            "missing Rust bigram weights const"
        );
        assert!(
            source.contains("pub const RUST_AST_TRIGRAM_WEIGHTS: &[(u64, f32)]"),
            "missing Rust trigram weights const"
        );
    }

    #[test]
    fn generated_source_contains_lookup_functions() {
        let table = sample_table();
        let source = build_ast_weights_rs(&table).unwrap();
        // Structural markers — verifies the function signatures are syntactically valid Rust.
        assert!(
            source.contains("pub fn ast_bigram_weight(lang: &str, bigram: u32) -> Option<f32>"),
            "missing ast_bigram_weight function signature"
        );
        assert!(
            source.contains("pub fn ast_trigram_weight(lang: &str, trigram: u64) -> Option<f32>"),
            "missing ast_trigram_weight function signature"
        );
        assert!(
            source.contains("#[cfg(test)]"),
            "generated source must include embedded test module"
        );
    }

    #[test]
    fn bigram_arrays_are_sorted_ascending() {
        let table = sample_table();
        let source = build_ast_weights_rs(&table).unwrap();
        // The first bigram in the sorted array should be the one with the smaller key.
        let key_a = encode_ast_bigram(0, 1); // smaller
        let key_b = encode_ast_bigram(1, 0); // larger
        let pos_a = source.find(&format!("0x{:08X}", key_a)).unwrap();
        let pos_b = source.find(&format!("0x{:08X}", key_b)).unwrap();
        assert!(
            pos_a < pos_b,
            "smaller key (0x{key_a:08X}) should appear before larger key (0x{key_b:08X})"
        );
    }

    #[test]
    fn lang_to_ident_handles_language_names() {
        assert_eq!(lang_to_ident("Rust").unwrap(), "RUST");
        assert_eq!(lang_to_ident("TypeScript").unwrap(), "TYPESCRIPT");
        assert_eq!(lang_to_ident("C").unwrap(), "C");
        assert_eq!(lang_to_ident("JavaScript").unwrap(), "JAVASCRIPT");
        // Semantic mappings for punctuation-heavy names that would collide.
        assert_eq!(lang_to_ident("C++").unwrap(), "CPP");
        assert_eq!(lang_to_ident("C#").unwrap(), "CSHARP");
    }

    #[test]
    fn lang_to_ident_collapses_consecutive_underscores() {
        assert_eq!(lang_to_ident("some--lang").unwrap(), "SOME_LANG");
        assert_eq!(lang_to_ident("Go").unwrap(), "GO");
    }

    #[test]
    fn lang_to_ident_rejects_digit_leading_name() {
        // A language name starting with a digit produces an identifier-start error.
        let err = lang_to_ident("1lang").unwrap_err();
        assert!(
            err.to_string().contains("empty or non-identifier-starting"),
            "expected identifier-start error, got: {err}"
        );
    }

    #[test]
    fn lang_to_ident_rejects_non_ascii_name() {
        // A language name with a non-ASCII character (e.g. a Unicode letter not in
        // [A-Za-z0-9_]) must fail the identifier-character check.
        let err = lang_to_ident("Ré").unwrap_err();
        assert!(
            err.to_string().contains("non-identifier character"),
            "expected non-identifier character error, got: {err}"
        );
    }

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

        let mut table = sample_table();
        table.version = 0;

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

        let err = generate_ast_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(err.to_string().contains("version must be > 0"));
    }

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

        let mut table = sample_table();
        table.bigram_weights.get_mut("Rust").unwrap()[0].idf = -1.0;

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

        let err = generate_ast_weights_rs(&json_path, &out_path).unwrap_err();
        assert!(err.to_string().contains("invalid IDF"));
    }

    #[test]
    fn nan_idf_returns_error() {
        // NaN cannot round-trip through JSON, so call validate_ast_table directly.
        let mut table = sample_table();
        table.bigram_weights.get_mut("Rust").unwrap()[0].idf = f32::NAN;

        let err = validate_ast_table(&table).unwrap_err();
        assert!(
            err.to_string().contains("invalid IDF"),
            "NaN IDF should be rejected: {err}"
        );
    }

    #[test]
    fn infinity_idf_returns_error() {
        // Infinity cannot round-trip through JSON, so call validate_ast_table directly.
        let mut table = sample_table();
        table.bigram_weights.get_mut("Rust").unwrap()[0].idf = f32::INFINITY;

        let err = validate_ast_table(&table).unwrap_err();
        assert!(
            err.to_string().contains("invalid IDF"),
            "Infinity IDF should be rejected: {err}"
        );
    }

    #[test]
    fn empty_bigram_weights_returns_error() {
        let mut table = sample_table();
        table.bigram_weights.clear();
        let err = validate_ast_table(&table).unwrap_err();
        assert!(
            err.to_string().contains("no bigram weights"),
            "empty bigram_weights should be rejected: {err}"
        );
    }

    #[test]
    fn empty_trigram_weights_returns_error() {
        let mut table = sample_table();
        table.trigram_weights.clear();
        let err = validate_ast_table(&table).unwrap_err();
        assert!(
            err.to_string().contains("no trigram weights"),
            "empty trigram_weights should be rejected: {err}"
        );
    }
}
