chipi_core/

codegen.rs

//! Rust code generation from validated definitions and decision trees.
//!
//! Generates clean, idiomatic Rust code with an `impl` block containing a `decode()` function
//! and an enum representing all instruction variants with their fields.

use std::collections::HashMap;
use std::fmt::Write;

use crate::tree::DecodeNode;
use crate::types::*;

/// Check if any instruction in the decoder requires multiple units.
fn needs_variable_length_decode(def: &ValidatedDef) -> bool {
    def.instructions.iter().any(|i| i.unit_count() > 1)
}

/// Generate Rust source code from a validated definition and dispatch tree.
///
/// The output includes:
/// - An enum with all instruction variants
/// - A `decode()` function using efficient match statements
/// - Proper type conversions and bit extraction
pub fn generate_code(
    def: &ValidatedDef,
    tree: &DecodeNode,
    type_maps: &HashMap<String, String>,
    dispatch_overrides: &HashMap<String, crate::Dispatch>,
) -> String {
    let mut out = String::new();

    writeln!(out, "// Auto-generated by https://github.com/ioncodes/chipi").unwrap();
    writeln!(out, "// Do not edit.").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "use std::fmt;").unwrap();
    writeln!(out, "use std::marker::PhantomData;").unwrap();
    writeln!(out).unwrap();

    // Type map imports
    let mut import_paths: Vec<&str> = type_maps.values().map(|v| v.as_str()).collect();
    import_paths.sort();
    import_paths.dedup();
    for path in &import_paths {
        // Only emit `use` for paths that look like Rust module paths (contain `::`)
        if path.contains("::") {
            writeln!(out, "use {};", path).unwrap();
        }
    }
    if !import_paths.is_empty() {
        writeln!(out).unwrap();
    }

    let word_type = word_type_for_width(def.config.width);
    let unit_bytes = def.config.width / 8;
    let variable_length = needs_variable_length_decode(def);
    let enum_name = format!("{}Instruction", def.config.name);
    let trait_name = format!("{}Format", def.config.name);
    let default_struct = format!("Default{}Format", def.config.name);
    let display_with = "DisplayWith";
    let endian_suffix = match def.config.endian {
        ByteEndian::Big => "be",
        ByteEndian::Little => "le",
    };

    // Display format helpers
    generate_display_helpers(&mut out, def);

    // Map functions
    generate_map_functions(&mut out, def);

    // Sub-decoder types and dispatch
    for sd in &def.sub_decoders {
        let dispatch = dispatch_overrides
            .get(&sd.name)
            .copied()
            .unwrap_or(crate::Dispatch::FnPtrLut);
        generate_subdecoder(&mut out, sd, dispatch);
    }

    // Enum definition
    writeln!(out, "#[derive(Debug, Clone, Copy, PartialEq, Eq)]").unwrap();
    writeln!(out, "pub enum {} {{", enum_name).unwrap();

    for instr in &def.instructions {
        let variant_name = to_pascal_case(&instr.name);
        if instr.resolved_fields.is_empty() {
            writeln!(out, "    {},", variant_name).unwrap();
        } else {
            let fields: Vec<String> = instr
                .resolved_fields
                .iter()
                .map(|f| {
                    let rust_type = field_rust_type_with_maps(f, &type_maps);
                    format!("{}: {}", f.name, rust_type)
                })
                .collect();
            writeln!(out, "    {} {{ {} }},", variant_name, fields.join(", ")).unwrap();
        }
    }

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // Formatting trait
    generate_format_trait(&mut out, def, &trait_name, type_maps);

    // impl block with decode() and write_asm() and display()
    writeln!(out, "impl {} {{", enum_name).unwrap();
    writeln!(out, "    #[inline]").unwrap();

    // Always generate &[u8] signature, returning (Self, bytes_consumed)
    writeln!(
        out,
        "    pub fn decode(data: &[u8]) -> Option<(Self, usize)> {{"
    )
    .unwrap();
    writeln!(
        out,
        "        if data.len() < {} {{ return None; }}",
        unit_bytes
    )
    .unwrap();
    writeln!(
        out,
        "        let opcode = {}::from_{}_bytes(data[0..{}].try_into().unwrap());",
        word_type, endian_suffix, unit_bytes
    )
    .unwrap();

    emit_tree(
        &mut out,
        tree,
        def,
        &enum_name,
        2,
        variable_length,
        &word_type,
        type_maps,
    );

    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();

    // write_asm method
    generate_write_asm(&mut out, def, &enum_name, &trait_name);

    // display method
    writeln!(out).unwrap();
    writeln!(out, "    #[allow(dead_code)]").unwrap();
    writeln!(
        out,
        "    pub fn display<F: {}>(&self) -> {}<'_, F> {{",
        trait_name, display_with
    )
    .unwrap();
    writeln!(
        out,
        "        {} {{ insn: self, _phantom: PhantomData }}",
        display_with
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // DisplayWith struct
    writeln!(out, "#[allow(dead_code)]").unwrap();
    writeln!(out, "pub struct {}<'a, F: {}> {{", display_with, trait_name).unwrap();
    writeln!(out, "    insn: &'a {},", enum_name).unwrap();
    writeln!(out, "    _phantom: PhantomData<F>,").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    writeln!(
        out,
        "impl<F: {}> fmt::Display for {}<'_, F> {{",
        trait_name, display_with
    )
    .unwrap();
    writeln!(
        out,
        "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{"
    )
    .unwrap();
    writeln!(out, "        self.insn.write_asm::<F>(f)").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // Default format struct
    writeln!(out, "pub struct {};", default_struct).unwrap();
    writeln!(out, "impl {} for {} {{}}", trait_name, default_struct).unwrap();
    writeln!(out).unwrap();

    // Display impl using default format
    writeln!(out, "impl fmt::Display for {} {{", enum_name).unwrap();
    writeln!(
        out,
        "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{"
    )
    .unwrap();
    writeln!(out, "        self.write_asm::<{}>(f)", default_struct).unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();

    out
}

/// Emit a decode tree node in block context (i.e. as a statement/expression, not a match arm).
fn emit_tree(
    out: &mut String,
    node: &DecodeNode,
    def: &ValidatedDef,
    enum_name: &str,
    indent: usize,
    variable_length: bool,
    word_type: &str,
    type_maps: &HashMap<String, String>,
) {
    let unit_bytes = def.config.width / 8;
    let endian_suffix = match def.config.endian {
        ByteEndian::Big => "be",
        ByteEndian::Little => "le",
    };
    let pad = "    ".repeat(indent);
    match node {
        DecodeNode::Leaf { instruction_index } => {
            let instr = &def.instructions[*instruction_index];
            if let Some(guard) = leaf_guard(instr, word_type, unit_bytes, endian_suffix) {
                writeln!(out, "{}if {} {{", pad, guard).unwrap();
                emit_some(
                    out,
                    instr,
                    enum_name,
                    &format!("{}    ", pad),
                    variable_length,
                    word_type,
                    unit_bytes,
                    endian_suffix,
                    &type_maps,
                );
                writeln!(out, "{}}} else {{", pad).unwrap();
                writeln!(out, "{}    None", pad).unwrap();
                writeln!(out, "{}}}", pad).unwrap();
            } else {
                emit_some(
                    out,
                    instr,
                    enum_name,
                    &pad,
                    variable_length,
                    word_type,
                    unit_bytes,
                    endian_suffix,
                    &type_maps,
                );
            }
        }
        DecodeNode::PriorityLeaves { candidates } => {
            // Try each candidate in priority order (most specific first)
            for (i, &idx) in candidates.iter().enumerate() {
                let instr = &def.instructions[idx];
                let guard = leaf_guard(instr, word_type, unit_bytes, endian_suffix);

                if i == 0 {
                    // First candidate
                    if let Some(guard_expr) = guard {
                        writeln!(out, "{}if {} {{", pad, guard_expr).unwrap();
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}    ", pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                    } else {
                        // No guard needed - this should match unconditionally
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &pad,
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                        break; // No need to check further candidates
                    }
                } else if i == candidates.len() - 1 {
                    // Last candidate
                    writeln!(out, "{}}} else {{", pad).unwrap();
                    if let Some(guard_expr) = guard {
                        writeln!(out, "{}    if {} {{", pad, guard_expr).unwrap();
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}        ", pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                        writeln!(out, "{}    }} else {{", pad).unwrap();
                        writeln!(out, "{}        None", pad).unwrap();
                        writeln!(out, "{}    }}", pad).unwrap();
                    } else {
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}    ", pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                    }
                    writeln!(out, "{}}}", pad).unwrap();
                } else {
                    // Middle candidates
                    writeln!(
                        out,
                        "{}}} else if {} {{",
                        pad,
                        guard.unwrap_or_else(|| "true".to_string())
                    )
                    .unwrap();
                    emit_some(
                        out,
                        instr,
                        enum_name,
                        &format!("{}    ", pad),
                        variable_length,
                        word_type,
                        unit_bytes,
                        endian_suffix,
                        &type_maps,
                    );
                }
            }
        }
        DecodeNode::Fail => {
            writeln!(out, "{}None", pad).unwrap();
        }
        DecodeNode::Branch {
            range,
            arms,
            default,
        } => {
            let extract_expr =
                extract_expression("opcode", &[*range], word_type, unit_bytes, endian_suffix);
            writeln!(out, "{}match {} {{", pad, extract_expr).unwrap();

            for (value, child) in arms {
                emit_arm(
                    out,
                    child,
                    def,
                    enum_name,
                    indent + 1,
                    &format!("{:#x}", value),
                    variable_length,
                    word_type,
                    type_maps,
                );
            }

            emit_arm(
                out,
                default,
                def,
                enum_name,
                indent + 1,
                "_",
                variable_length,
                word_type,
                type_maps,
            );

            writeln!(out, "{}}}", pad).unwrap();
        }
    }
}

/// Emit a single match arm for a given pattern and child node.
fn emit_arm(
    out: &mut String,
    node: &DecodeNode,
    def: &ValidatedDef,
    enum_name: &str,
    indent: usize,
    pattern: &str,
    variable_length: bool,
    word_type: &str,
    type_maps: &HashMap<String, String>,
) {
    let unit_bytes = def.config.width / 8;
    let endian_suffix = match def.config.endian {
        ByteEndian::Big => "be",
        ByteEndian::Little => "le",
    };
    let pad = "    ".repeat(indent);
    match node {
        DecodeNode::Fail => {
            writeln!(out, "{}{} => None,", pad, pattern).unwrap();
        }
        DecodeNode::Leaf { instruction_index } => {
            let instr = &def.instructions[*instruction_index];
            if let Some(guard) = leaf_guard(instr, word_type, unit_bytes, endian_suffix) {
                if pattern == "_" {
                    // Default arm with guard: emit guarded arm then fallback
                    write!(out, "{}{} if {} => ", pad, pattern, guard).unwrap();
                    emit_some_inline(
                        out,
                        instr,
                        enum_name,
                        variable_length,
                        word_type,
                        unit_bytes,
                        endian_suffix,
                        &type_maps,
                    );
                    writeln!(out, "{}{} => None,", pad, pattern).unwrap();
                } else {
                    write!(out, "{}{} if {} => ", pad, pattern, guard).unwrap();
                    emit_some_inline(
                        out,
                        instr,
                        enum_name,
                        variable_length,
                        word_type,
                        unit_bytes,
                        endian_suffix,
                        &type_maps,
                    );
                }
            } else {
                write!(out, "{}{} => ", pad, pattern).unwrap();
                emit_some_inline(
                    out,
                    instr,
                    enum_name,
                    variable_length,
                    word_type,
                    unit_bytes,
                    endian_suffix,
                    &type_maps,
                );
            }
        }
        DecodeNode::PriorityLeaves { candidates } => {
            // Emit as a block with if-else chain
            writeln!(out, "{}{} => {{", pad, pattern).unwrap();
            let inner_pad = "    ".repeat(indent + 1);

            for (i, &idx) in candidates.iter().enumerate() {
                let instr = &def.instructions[idx];
                let guard = leaf_guard(instr, word_type, unit_bytes, endian_suffix);

                if i == 0 {
                    // First candidate
                    if let Some(guard_expr) = guard {
                        writeln!(out, "{}if {} {{", inner_pad, guard_expr).unwrap();
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}    ", inner_pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                    } else {
                        // No guard - matches unconditionally
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &inner_pad,
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                        writeln!(out, "{}}}", pad).unwrap();
                        return;
                    }
                } else if i == candidates.len() - 1 {
                    // Last candidate
                    writeln!(out, "{}}} else {{", inner_pad).unwrap();
                    if let Some(guard_expr) = guard {
                        writeln!(out, "{}    if {} {{", inner_pad, guard_expr).unwrap();
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}        ", inner_pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                        writeln!(out, "{}    }} else {{", inner_pad).unwrap();
                        writeln!(out, "{}        None", inner_pad).unwrap();
                        writeln!(out, "{}    }}", inner_pad).unwrap();
                    } else {
                        emit_some(
                            out,
                            instr,
                            enum_name,
                            &format!("{}    ", inner_pad),
                            variable_length,
                            word_type,
                            unit_bytes,
                            endian_suffix,
                            &type_maps,
                        );
                    }
                    writeln!(out, "{}}}", inner_pad).unwrap();
                    writeln!(out, "{}}}", pad).unwrap();
                } else {
                    // Middle candidates
                    writeln!(
                        out,
                        "{}}} else if {} {{",
                        inner_pad,
                        guard.unwrap_or_else(|| "true".to_string())
                    )
                    .unwrap();
                    emit_some(
                        out,
                        instr,
                        enum_name,
                        &format!("{}    ", inner_pad),
                        variable_length,
                        word_type,
                        unit_bytes,
                        endian_suffix,
                        &type_maps,
                    );
                }
            }
        }
        DecodeNode::Branch {
            range,
            arms,
            default,
        } => {
            writeln!(out, "{}{} => {{", pad, pattern).unwrap();
            let extract_expr =
                extract_expression("opcode", &[*range], word_type, unit_bytes, endian_suffix);
            let inner_pad = "    ".repeat(indent + 1);
            writeln!(out, "{}match {} {{", inner_pad, extract_expr).unwrap();

            for (value, child) in arms {
                emit_arm(
                    out,
                    child,
                    def,
                    enum_name,
                    indent + 2,
                    &format!("{:#x}", value),
                    variable_length,
                    word_type,
                    type_maps,
                );
            }

            emit_arm(
                out,
                default,
                def,
                enum_name,
                indent + 2,
                "_",
                variable_length,
                word_type,
                type_maps,
            );

            writeln!(out, "{}}}", inner_pad).unwrap();
            writeln!(out, "{}}}", pad).unwrap();
        }
    }
}

/// Compute the guard condition string for a leaf instruction, if needed.
/// Returns `None` if all fixed bits were already fully dispatched by the tree.
/// For multi-unit instructions, generates guards for ALL units (unit 0 and unit 1+).
fn leaf_guard(
    instr: &ValidatedInstruction,
    word_type: &str,
    unit_bytes: u32,
    endian_suffix: &str,
) -> Option<String> {
    let fixed_bits = instr.fixed_bits();
    if fixed_bits.is_empty() {
        return None;
    }

    // Group fixed bits by unit
    let mut units_map: std::collections::HashMap<u32, Vec<(u32, Bit)>> =
        std::collections::HashMap::new();
    for (unit, hw_bit, bit) in fixed_bits {
        units_map.entry(unit).or_default().push((hw_bit, bit));
    }

    let mut conditions = Vec::new();

    for (unit, bits) in units_map {
        let (mask, value) = compute_mask_value(&bits);
        if mask != 0 {
            let source = unit_read_expr(unit, word_type, unit_bytes, endian_suffix);
            conditions.push(format!("{} & {:#x} == {:#x}", source, mask, value));
        }
    }

    if conditions.is_empty() {
        None
    } else {
        Some(conditions.join(" && "))
    }
}

/// Write `Some((EnumName::Variant { ... }, bytes))` as an inline match arm value (terminated with comma+newline).
fn emit_some_inline(
    out: &mut String,
    instr: &ValidatedInstruction,
    enum_name: &str,
    variable_length: bool,
    word_type: &str,
    unit_bytes: u32,
    endian_suffix: &str,
    type_maps: &HashMap<String, String>,
) {
    let variant_name = to_pascal_case(&instr.name);
    let unit_count = instr.unit_count();
    let bytes_consumed = unit_count * unit_bytes;

    if variable_length && unit_count > 1 {
        // Emit bounds check for multi-unit instruction
        let cond = format!("data.len() >= {}", bytes_consumed);
        write!(out, "if {} {{ ", cond).unwrap();
    }

    if instr.resolved_fields.is_empty() {
        write!(
            out,
            "Some(({}::{}, {}))",
            enum_name, variant_name, bytes_consumed
        )
        .unwrap();
    } else {
        let fields: Vec<String> = instr
            .resolved_fields
            .iter()
            .map(|f| {
                let extract =
                    extract_expression("opcode", &f.ranges, word_type, unit_bytes, endian_suffix);
                let expr = apply_transforms_with_maps(&extract, &f.resolved_type, type_maps);
                format!("{}: {}", f.name, expr)
            })
            .collect();
        write!(
            out,
            "Some(({}::{} {{ {} }}, {}))",
            enum_name,
            variant_name,
            fields.join(", "),
            bytes_consumed
        )
        .unwrap();
    }

    if variable_length && unit_count > 1 {
        write!(out, " }} else {{ None }}").unwrap();
    }
    writeln!(out, ",").unwrap();
}

/// Write `Some((EnumName::Variant { ... }, bytes))` in block context (multi-line, no trailing comma).
fn emit_some(
    out: &mut String,
    instr: &ValidatedInstruction,
    enum_name: &str,
    pad: &str,
    variable_length: bool,
    word_type: &str,
    unit_bytes: u32,
    endian_suffix: &str,
    type_maps: &HashMap<String, String>,
) {
    let unit_count = instr.unit_count();
    let bytes_consumed = unit_count * unit_bytes;

    if variable_length && unit_count > 1 {
        writeln!(out, "{}if data.len() >= {} {{", pad, bytes_consumed).unwrap();
        let inner_pad = format!("{}    ", pad);
        emit_some_inner(
            out,
            instr,
            enum_name,
            &inner_pad,
            word_type,
            unit_bytes,
            endian_suffix,
            bytes_consumed,
            type_maps,
        );
        writeln!(out, "{}}} else {{", pad).unwrap();
        writeln!(out, "{}    None", pad).unwrap();
        writeln!(out, "{}}}", pad).unwrap();
    } else {
        emit_some_inner(
            out,
            instr,
            enum_name,
            pad,
            word_type,
            unit_bytes,
            endian_suffix,
            bytes_consumed,
            type_maps,
        );
    }
}

/// Helper to emit the Some(...) part without bounds checking.
fn emit_some_inner(
    out: &mut String,
    instr: &ValidatedInstruction,
    enum_name: &str,
    pad: &str,
    word_type: &str,
    unit_bytes: u32,
    endian_suffix: &str,
    bytes_consumed: u32,
    type_maps: &HashMap<String, String>,
) {
    let variant_name = to_pascal_case(&instr.name);
    if instr.resolved_fields.is_empty() {
        writeln!(
            out,
            "{}Some(({}::{}, {}))",
            pad, enum_name, variant_name, bytes_consumed
        )
        .unwrap();
    } else {
        writeln!(out, "{}Some(({}::{} {{", pad, enum_name, variant_name).unwrap();
        for field in &instr.resolved_fields {
            let extract = extract_expression(
                "opcode",
                &field.ranges,
                word_type,
                unit_bytes,
                endian_suffix,
            );
            let expr = apply_transforms_with_maps(&extract, &field.resolved_type, type_maps);
            writeln!(out, "{}    {}: {},", pad, field.name, expr).unwrap();
        }
        writeln!(out, "{}}}, {}))", pad, bytes_consumed).unwrap();
    }
}

/// Compute a bitmask and expected value from a list of (hw_bit_position, Bit) pairs.
/// Wildcard bits are skipped (not included in the mask).
fn compute_mask_value(fixed_bits: &[(u32, Bit)]) -> (u64, u64) {
    let mut mask: u64 = 0;
    let mut value: u64 = 0;
    for &(bit_pos, bit_val) in fixed_bits {
        // Skip wildcard bits - they don't contribute to the mask
        if bit_val == Bit::Wildcard {
            continue;
        }
        mask |= 1u64 << bit_pos;
        if bit_val == Bit::One {
            value |= 1u64 << bit_pos;
        }
    }
    (mask, value)
}

/// Generate an expression to read a unit from the byte slice.
/// Unit 0 returns "opcode" (already decoded in preamble).
/// Unit N>0 returns an inline byte read expression.
fn unit_read_expr(unit: u32, word_type: &str, unit_bytes: u32, endian_suffix: &str) -> String {
    if unit == 0 {
        "opcode".to_string()
    } else {
        let start = unit * unit_bytes;
        let end = start + unit_bytes;
        format!(
            "{}::from_{}_bytes(data[{}..{}].try_into().unwrap())",
            word_type, endian_suffix, start, end
        )
    }
}

/// Generate an expression to extract bits from multiple ranges (potentially cross-unit).
fn extract_expression(
    var: &str,
    ranges: &[BitRange],
    word_type: &str,
    unit_bytes: u32,
    endian_suffix: &str,
) -> String {
    if ranges.is_empty() {
        return "0".to_string();
    }

    if ranges.len() == 1 {
        // Single range - simple extraction
        let range = ranges[0];
        let source = if range.unit == 0 {
            var.to_string()
        } else {
            unit_read_expr(range.unit, word_type, unit_bytes, endian_suffix)
        };

        let width = range.width();
        let shift = range.end;
        let mask = (1u64 << width) - 1;

        if shift == 0 {
            format!("{} & {:#x}", source, mask)
        } else {
            format!("({} >> {}) & {:#x}", source, shift, mask)
        }
    } else {
        // Multi-range extraction - combine bits from multiple ranges
        let mut parts = Vec::new();
        let mut accumulated_width = 0u32;

        // Ranges are ordered from low-order to high-order bits
        for range in ranges {
            let source = if range.unit == 0 {
                var.to_string()
            } else {
                unit_read_expr(range.unit, word_type, unit_bytes, endian_suffix)
            };

            let width = range.width();
            let shift = range.end;
            let mask = (1u64 << width) - 1;

            let extracted = if shift == 0 {
                format!("({} & {:#x})", source, mask)
            } else {
                format!("(({} >> {}) & {:#x})", source, shift, mask)
            };

            // Shift this part into its position in the final value
            if accumulated_width > 0 {
                parts.push(format!("({} << {})", extracted, accumulated_width));
            } else {
                parts.push(extracted);
            }

            accumulated_width += width;
        }

        parts.join(" | ")
    }
}

/// Apply transforms and type conversions to extracted field values.
fn apply_transforms(extract_expr: &str, resolved: &ResolvedFieldType) -> String {
    apply_transforms_with_maps(extract_expr, resolved, &HashMap::new())
}

/// Apply transforms with type map support.
fn apply_transforms_with_maps(
    extract_expr: &str,
    resolved: &ResolvedFieldType,
    type_maps: &HashMap<String, String>,
) -> String {
    let mut expr = extract_expr.to_string();

    for transform in &resolved.transforms {
        match transform {
            Transform::SignExtend(n) => {
                let signed_type = signed_type_for(&resolved.base_type);
                let bits = type_bits(&resolved.base_type);
                expr = format!(
                    "(((({}) as {}) << ({} - {})) >> ({} - {}))",
                    expr, signed_type, bits, n, bits, n
                );
            }
            Transform::ZeroExtend(_) => {
                // No-op for unsigned types
            }
            Transform::ShiftLeft(n) => {
                expr = format!("(({}) << {})", expr, n);
            }
        }
    }

    // Handle sub-decoder fields: call the sub-decoder dispatch function
    if let Some(ref sd_name) = resolved.sub_decoder {
        let decode_fn = format!("decode_{}", to_snake_case(sd_name));
        // The field value is extracted and cast to the sub-decoder's word type, then dispatched
        return format!(
            "{}(({}) as {}).unwrap()",
            decode_fn, expr, resolved.base_type
        );
    }

    // Check if there's a type map for the wrapper type
    let wrapper_type = resolved.alias_name.as_ref().and_then(|a| type_maps.get(a));

    if let Some(wrapper) = wrapper_type {
        let short_name = wrapper.rsplit("::").next().unwrap_or(wrapper);
        expr = format!("{}::from(({}) as {})", short_name, expr, resolved.base_type);
    } else if resolved.base_type == "bool" {
        expr = format!("({}) != 0", expr);
    } else {
        expr = format!("({}) as {}", expr, resolved.base_type);
    }

    expr
}

/// Generate display format helper structs (SignedHex, Hex) if needed.
fn generate_display_helpers(out: &mut String, def: &ValidatedDef) {
    let mut need_signed_hex = false;
    let mut need_hex = false;

    for instr in &def.instructions {
        for field in &instr.resolved_fields {
            match field.resolved_type.display_format {
                Some(DisplayFormat::SignedHex) => need_signed_hex = true,
                Some(DisplayFormat::Hex) => need_hex = true,
                None => {}
            }
        }
    }

    if need_signed_hex {
        writeln!(out, "struct SignedHex<T>(T);").unwrap();
        writeln!(out).unwrap();

        for ty in &["i8", "i16", "i32"] {
            writeln!(out, "impl fmt::Display for SignedHex<{}> {{", ty).unwrap();
            writeln!(
                out,
                "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{"
            )
            .unwrap();
            writeln!(out, "        if self.0 == 0 {{ write!(f, \"0\") }}").unwrap();
            writeln!(
                out,
                "        else if self.0 > 0 {{ write!(f, \"0x{{:X}}\", self.0) }}"
            )
            .unwrap();
            writeln!(
                out,
                "        else {{ write!(f, \"-0x{{:X}}\", (self.0 as i64).wrapping_neg()) }}"
            )
            .unwrap();
            writeln!(out, "    }}").unwrap();
            writeln!(out, "}}").unwrap();
            writeln!(out).unwrap();
        }
    }

    if need_hex {
        if !need_signed_hex {
            writeln!(out, "struct SignedHex<T>(T);").unwrap();
            writeln!(out).unwrap();
        }

        for ty in &["u8", "u16", "u32"] {
            writeln!(out, "impl fmt::Display for SignedHex<{}> {{", ty).unwrap();
            writeln!(
                out,
                "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{"
            )
            .unwrap();
            writeln!(out, "        if self.0 == 0 {{ write!(f, \"0\") }}").unwrap();
            writeln!(out, "        else {{ write!(f, \"0x{{:X}}\", self.0) }}").unwrap();
            writeln!(out, "    }}").unwrap();
            writeln!(out, "}}").unwrap();
            writeln!(out).unwrap();
        }
    }
}

/// Generate map lookup functions.
fn generate_map_functions(out: &mut String, def: &ValidatedDef) {
    if def.maps.is_empty() {
        return;
    }

    // Collect map call sites to infer parameter types
    let map_param_types = infer_map_param_types(def);

    for map in &def.maps {
        let has_interpolation = map.entries.iter().any(|entry| {
            entry
                .output
                .iter()
                .any(|p| matches!(p, FormatPiece::FieldRef { .. }))
        });

        let param_types: Vec<String> = map
            .params
            .iter()
            .enumerate()
            .map(|(i, _)| {
                map_param_types
                    .get(&map.name)
                    .and_then(|types| types.get(i))
                    .cloned()
                    .unwrap_or_else(|| "i64".to_string())
            })
            .collect();

        let return_type = if has_interpolation {
            "String"
        } else {
            "&'static str"
        };

        let params: Vec<String> = map
            .params
            .iter()
            .zip(param_types.iter())
            .map(|(name, ty)| format!("{}: {}", name, ty))
            .collect();

        writeln!(out, "#[allow(dead_code)]").unwrap();
        writeln!(
            out,
            "fn {}({}) -> {} {{",
            map.name,
            params.join(", "),
            return_type
        )
        .unwrap();

        if map.params.len() == 1 {
            writeln!(out, "    match {} {{", map.params[0]).unwrap();
        } else {
            let tuple: Vec<&str> = map.params.iter().map(|s| s.as_str()).collect();
            writeln!(out, "    match ({}) {{", tuple.join(", ")).unwrap();
        }

        // Separate wildcard (default) entries from specific entries
        let mut default_entry = None;

        for entry in &map.entries {
            let is_all_wildcard = entry.keys.iter().all(|k| matches!(k, MapKey::Wildcard));
            if is_all_wildcard {
                default_entry = Some(entry);
                continue;
            }

            let pattern = if map.params.len() == 1 {
                format_map_key(&entry.keys[0])
            } else {
                let keys: Vec<String> = entry.keys.iter().map(|k| format_map_key(k)).collect();
                format!("({})", keys.join(", "))
            };

            let output = format_map_output(&entry.output, has_interpolation);
            writeln!(out, "        {} => {},", pattern, output).unwrap();
        }

        // Default arm
        if let Some(entry) = default_entry {
            let output = format_map_output(&entry.output, has_interpolation);
            writeln!(out, "        _ => {},", output).unwrap();
        } else {
            if has_interpolation {
                writeln!(out, "        _ => String::from(\"???\"),").unwrap();
            } else {
                writeln!(out, "        _ => \"???\",").unwrap();
            }
        }

        writeln!(out, "    }}").unwrap();
        writeln!(out, "}}").unwrap();
        writeln!(out).unwrap();
    }
}

fn format_map_key(key: &MapKey) -> String {
    match key {
        MapKey::Value(v) => format!("{}", v),
        MapKey::Wildcard => "_".to_string(),
    }
}

fn format_map_output(pieces: &[FormatPiece], has_interpolation: bool) -> String {
    if !has_interpolation {
        // All pieces should be literals
        let mut s = String::new();
        for piece in pieces {
            if let FormatPiece::Literal(lit) = piece {
                s.push_str(lit);
            }
        }
        return format!("\"{}\"", s);
    }

    // Build a format!() call
    let mut fmt_str = String::new();
    let mut args = Vec::new();

    for piece in pieces {
        match piece {
            FormatPiece::Literal(lit) => {
                // Escape `{` and `}` for format string
                for ch in lit.chars() {
                    match ch {
                        '{' => fmt_str.push_str("{{"),
                        '}' => fmt_str.push_str("}}"),
                        _ => fmt_str.push(ch),
                    }
                }
            }
            FormatPiece::FieldRef { expr, spec } => {
                if let Some(spec) = spec {
                    fmt_str.push_str(&format!("{{:{}}}", spec));
                } else {
                    fmt_str.push_str("{}");
                }
                args.push(expr_to_rust(expr, &[]));
            }
        }
    }

    if args.is_empty() {
        format!("String::from(\"{}\")", fmt_str)
    } else {
        format!("format!(\"{}\", {})", fmt_str, args.join(", "))
    }
}

/// Infer map parameter types from call sites across all instructions.
fn infer_map_param_types(def: &ValidatedDef) -> HashMap<String, Vec<String>> {
    let mut result: HashMap<String, Vec<String>> = HashMap::new();

    // Build field type lookup per instruction
    for instr in &def.instructions {
        let field_types: HashMap<&str, &ResolvedFieldType> = instr
            .resolved_fields
            .iter()
            .map(|f| (f.name.as_str(), &f.resolved_type))
            .collect();

        for fl in &instr.format_lines {
            for piece in &fl.pieces {
                if let FormatPiece::FieldRef { expr, .. } = piece {
                    collect_map_call_types(expr, &field_types, &mut result);
                }
            }
        }
    }

    result
}

/// Infer the Rust type of a format expression given field type information.
fn infer_expr_type(
    expr: &FormatExpr,
    field_types: &HashMap<&str, &ResolvedFieldType>,
) -> Option<String> {
    match expr {
        FormatExpr::Field(name) => field_types
            .get(name.as_str())
            .map(|ft| ft.base_type.clone()),
        FormatExpr::Arithmetic { left, .. } => infer_expr_type(left, field_types),
        FormatExpr::IntLiteral(_) => Some("i64".to_string()),
        _ => None,
    }
}

fn collect_map_call_types(
    expr: &FormatExpr,
    field_types: &HashMap<&str, &ResolvedFieldType>,
    result: &mut HashMap<String, Vec<String>>,
) {
    match expr {
        FormatExpr::MapCall { map_name, args } => {
            let entry = result
                .entry(map_name.clone())
                .or_insert_with(|| vec!["i64".to_string(); args.len()]);

            for (i, arg) in args.iter().enumerate() {
                if i < entry.len() {
                    if let Some(rust_type) = infer_expr_type(arg, field_types) {
                        entry[i] = rust_type;
                    }
                }
            }

            // Also recurse into map call arguments to find nested map calls
            for arg in args {
                collect_map_call_types(arg, field_types, result);
            }
        }
        FormatExpr::Arithmetic { left, right, .. } => {
            collect_map_call_types(left, field_types, result);
            collect_map_call_types(right, field_types, result);
        }
        FormatExpr::BuiltinCall { args, .. } => {
            for arg in args {
                collect_map_call_types(arg, field_types, result);
            }
        }
        FormatExpr::SubDecoderAccess { .. }
        | FormatExpr::Field(_)
        | FormatExpr::Ternary { .. }
        | FormatExpr::IntLiteral(_) => {}
    }
}

/// Generate the formatting trait with one method per instruction variant.
fn generate_format_trait(
    out: &mut String,
    def: &ValidatedDef,
    trait_name: &str,
    type_maps: &HashMap<String, String>,
) {
    writeln!(out, "pub trait {} {{", trait_name).unwrap();

    for instr in &def.instructions {
        let method_name = format!("fmt_{}", instr.name);
        let params = trait_method_params(&instr.resolved_fields, &type_maps);

        if params.is_empty() {
            writeln!(
                out,
                "    fn {}(f: &mut fmt::Formatter) -> fmt::Result {{",
                method_name
            )
            .unwrap();
        } else {
            writeln!(
                out,
                "    fn {}({}, f: &mut fmt::Formatter) -> fmt::Result {{",
                method_name, params
            )
            .unwrap();
        }

        generate_format_body(out, instr, 2, &type_maps);

        writeln!(out, "    }}").unwrap();
    }

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

/// Generate parameters for a trait method.
fn trait_method_params(fields: &[ResolvedField], type_maps: &HashMap<String, String>) -> String {
    let mut params = Vec::new();
    for field in fields {
        let rust_type = field_rust_type_with_maps(field, type_maps);
        params.push(format!("{}: {}", field.name, rust_type));
    }
    params.join(", ")
}

/// Generate the body of a format trait method from format lines.
fn generate_format_body(
    out: &mut String,
    instr: &ValidatedInstruction,
    indent: usize,
    type_maps: &HashMap<String, String>,
) {
    let pad = "    ".repeat(indent);

    if instr.format_lines.is_empty() {
        // Raw fallback
        if instr.resolved_fields.is_empty() {
            writeln!(out, "{}write!(f, \"{}\")", pad, instr.name).unwrap();
        } else {
            let field_names: Vec<&str> = instr
                .resolved_fields
                .iter()
                .map(|f| f.name.as_str())
                .collect();
            let placeholders: Vec<&str> = field_names.iter().map(|_| "{}").collect();
            let fmt_str = format!("{} {}", instr.name, placeholders.join(", "));
            let args: Vec<String> = instr
                .resolved_fields
                .iter()
                .map(|f| f.name.clone())
                .collect();
            writeln!(
                out,
                "{}write!(f, \"{}\", {})",
                pad,
                fmt_str,
                args.join(", ")
            )
            .unwrap();
        }
        return;
    }

    if instr.format_lines.len() == 1 && instr.format_lines[0].guard.is_none() {
        // Single format line, no guard
        emit_write_call(
            out,
            &instr.format_lines[0].pieces,
            &instr.resolved_fields,
            &pad,
        );
        return;
    }

    // Multiple format lines with guards
    for (i, fl) in instr.format_lines.iter().enumerate() {
        if let Some(guard) = &fl.guard {
            let guard_code = generate_guard_code(guard, &instr.resolved_fields, type_maps);
            if i == 0 {
                writeln!(out, "{}if {} {{", pad, guard_code).unwrap();
            } else {
                writeln!(out, "{}}} else if {} {{", pad, guard_code).unwrap();
            }
            emit_write_call(
                out,
                &fl.pieces,
                &instr.resolved_fields,
                &format!("{}    ", pad),
            );
        } else {
            // Last line without guard = else
            if i > 0 {
                writeln!(out, "{}}} else {{", pad).unwrap();
            }
            emit_write_call(
                out,
                &fl.pieces,
                &instr.resolved_fields,
                &format!("{}    ", pad),
            );
        }
    }

    if instr.format_lines.len() > 1
        || instr
            .format_lines
            .first()
            .map_or(false, |fl| fl.guard.is_some())
    {
        writeln!(out, "{}}}", pad).unwrap();
    }
}

/// Emit a `write!(f, ...)` call for a set of format pieces.
fn emit_write_call(out: &mut String, pieces: &[FormatPiece], fields: &[ResolvedField], pad: &str) {
    let mut fmt_str = String::new();
    let mut args = Vec::new();

    for piece in pieces {
        match piece {
            FormatPiece::Literal(lit) => {
                for ch in lit.chars() {
                    match ch {
                        '{' => fmt_str.push_str("{{"),
                        '}' => fmt_str.push_str("}}"),
                        _ => fmt_str.push(ch),
                    }
                }
            }
            FormatPiece::FieldRef { expr, spec } => {
                if let Some(spec) = spec {
                    fmt_str.push_str(&format!("{{:{}}}", spec));
                    args.push(expr_to_rust(expr, fields));
                } else if let Some(display_fmt) = resolve_display_format(expr, fields) {
                    fmt_str.push_str("{}");
                    let rust_expr = expr_to_rust(expr, fields);
                    match display_fmt {
                        DisplayFormat::SignedHex | DisplayFormat::Hex => {
                            args.push(format!("SignedHex({})", rust_expr));
                        }
                    }
                } else {
                    fmt_str.push_str("{}");
                    args.push(expr_to_rust(expr, fields));
                }
            }
        }
    }

    if args.is_empty() {
        writeln!(out, "{}write!(f, \"{}\")", pad, fmt_str).unwrap();
    } else {
        writeln!(
            out,
            "{}write!(f, \"{}\", {})",
            pad,
            fmt_str,
            args.join(", ")
        )
        .unwrap();
    }
}

/// Resolve display format for a format expression, if the expression is a simple field reference.
fn resolve_display_format(expr: &FormatExpr, fields: &[ResolvedField]) -> Option<DisplayFormat> {
    if let FormatExpr::Field(name) = expr {
        fields
            .iter()
            .find(|f| f.name == *name)
            .and_then(|f| f.resolved_type.display_format)
    } else {
        None
    }
}

/// Convert a format expression to Rust code.
fn expr_to_rust(expr: &FormatExpr, fields: &[ResolvedField]) -> String {
    match expr {
        FormatExpr::Field(name) => {
            let field = fields.iter().find(|f| f.name == *name);
            if let Some(f) = field {
                if f.resolved_type.base_type == "bool" {
                    format!("{} as u8", name)
                } else {
                    name.clone()
                }
            } else {
                name.clone()
            }
        }
        FormatExpr::Ternary {
            field,
            if_nonzero,
            if_zero,
        } => {
            let f = fields.iter().find(|f| f.name == *field);
            let cond = if let Some(f) = f {
                if f.resolved_type.base_type == "bool" {
                    field.clone()
                } else {
                    format!("{} != 0", field)
                }
            } else {
                format!("{} != 0", field)
            };

            let else_val = if_zero
                .as_deref()
                .map(|s| format!("\"{}\"", s))
                .unwrap_or_else(|| "\"\"".to_string());

            format!(
                "if {} {{ \"{}\" }} else {{ {} }}",
                cond, if_nonzero, else_val
            )
        }
        FormatExpr::Arithmetic { left, op, right } => {
            let l = expr_to_rust(left, fields);
            let r = expr_to_rust(right, fields);
            let op_str = match op {
                ArithOp::Add => "+",
                ArithOp::Sub => "-",
                ArithOp::Mul => "*",
                ArithOp::Div => "/",
                ArithOp::Mod => "%",
            };
            format!("{} {} {}", l, op_str, r)
        }
        FormatExpr::IntLiteral(val) => format!("{}", val),
        FormatExpr::MapCall { map_name, args } => {
            let arg_strs: Vec<String> = args.iter().map(|a| expr_to_rust(a, fields)).collect();
            format!("{}({})", map_name, arg_strs.join(", "))
        }
        FormatExpr::BuiltinCall { func, args } => {
            let arg_strs: Vec<String> = args.iter().map(|a| expr_to_rust(a, fields)).collect();
            match func {
                BuiltinFunc::RotateRight => {
                    format!(
                        "({} as u32).rotate_right({} as u32)",
                        arg_strs.get(0).map(|s| s.as_str()).unwrap_or("0"),
                        arg_strs.get(1).map(|s| s.as_str()).unwrap_or("0")
                    )
                }
                BuiltinFunc::RotateLeft => {
                    format!(
                        "({} as u32).rotate_left({} as u32)",
                        arg_strs.get(0).map(|s| s.as_str()).unwrap_or("0"),
                        arg_strs.get(1).map(|s| s.as_str()).unwrap_or("0")
                    )
                }
            }
        }
        FormatExpr::SubDecoderAccess { field, fragment } => {
            format!("{}.{}", field, fragment)
        }
    }
}

/// Generate Rust code for a guard condition.
fn generate_guard_code(
    guard: &Guard,
    fields: &[ResolvedField],
    type_maps: &HashMap<String, String>,
) -> String {
    let conditions: Vec<String> = guard
        .conditions
        .iter()
        .map(|cond| {
            let left = guard_operand_to_rust(&cond.left, fields);
            let right = guard_operand_to_rust(&cond.right, fields);
            let op = match cond.op {
                CompareOp::Eq => "==",
                CompareOp::Ne => "!=",
                CompareOp::Lt => "<",
                CompareOp::Le => "<=",
                CompareOp::Gt => ">",
                CompareOp::Ge => ">=",
            };

            let left_field = match &cond.left {
                GuardOperand::Field(name) => fields.iter().find(|f| f.name == *name),
                _ => None,
            };

            if let Some(f) = left_field {
                // Check for type-mapped wrapper types
                let wrapper = f
                    .resolved_type
                    .alias_name
                    .as_ref()
                    .and_then(|a| type_maps.get(a));
                if let Some(wrapper_path) = wrapper {
                    let short_name = wrapper_path.rsplit("::").next().unwrap_or(wrapper_path);
                    if let GuardOperand::Literal(val) = &cond.right {
                        return format!(
                            "{} {} {}::from({}{})",
                            left, op, short_name, val, f.resolved_type.base_type
                        );
                    }
                } else if f.resolved_type.base_type == "bool" {
                    // For bool fields, generate simpler comparisons
                    if let GuardOperand::Literal(val) = &cond.right {
                        match (cond.op, *val) {
                            (CompareOp::Eq, 0) => return format!("!{}", left),
                            (CompareOp::Eq, _) => return left.clone(),
                            (CompareOp::Ne, 0) => return left.clone(),
                            (CompareOp::Ne, _) => return format!("!{}", left),
                            _ => {}
                        }
                    }
                }
            }

            format!("{} {} {}", left, op, right)
        })
        .collect();

    conditions.join(" && ")
}

fn guard_operand_to_rust(operand: &GuardOperand, fields: &[ResolvedField]) -> String {
    match operand {
        GuardOperand::Field(name) => name.clone(),
        GuardOperand::Literal(val) => format!("{}", val),
        GuardOperand::Expr { left, op, right } => {
            let l = guard_operand_to_rust(left, fields);
            let r = guard_operand_to_rust(right, fields);
            let op_str = match op {
                ArithOp::Add => "+",
                ArithOp::Sub => "-",
                ArithOp::Mul => "*",
                ArithOp::Div => "/",
                ArithOp::Mod => "%",
            };
            format!("({} {} {})", l, op_str, r)
        }
    }
}

/// Generate the `write_asm` method.
fn generate_write_asm(out: &mut String, def: &ValidatedDef, enum_name: &str, trait_name: &str) {
    writeln!(
        out,
        "    pub fn write_asm<F: {}>(&self, f: &mut fmt::Formatter) -> fmt::Result {{",
        trait_name
    )
    .unwrap();
    writeln!(out, "        match self {{").unwrap();

    for instr in &def.instructions {
        let variant_name = to_pascal_case(&instr.name);
        let method_name = format!("fmt_{}", instr.name);

        if instr.resolved_fields.is_empty() {
            writeln!(
                out,
                "            {}::{} => F::{}(f),",
                enum_name, variant_name, method_name
            )
            .unwrap();
        } else {
            let field_names: Vec<String> = instr
                .resolved_fields
                .iter()
                .map(|f| f.name.clone())
                .collect();
            let destructure = format!(
                "{}::{} {{ {} }}",
                enum_name,
                variant_name,
                field_names.join(", ")
            );

            let args: Vec<String> = instr
                .resolved_fields
                .iter()
                .map(|f| format!("*{}", f.name))
                .collect();

            writeln!(
                out,
                "            {} => F::{}({}, f),",
                destructure,
                method_name,
                args.join(", ")
            )
            .unwrap();
        }
    }

    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
}

fn word_type_for_width(width: u32) -> &'static str {
    match width {
        8 => "u8",
        16 => "u16",
        32 => "u32",
        _ => "u32",
    }
}

pub(crate) fn signed_type_for(base: &str) -> &'static str {
    match base {
        "u8" | "i8" => "i8",
        "u16" | "i16" => "i16",
        "u32" | "i32" => "i32",
        _ => "i32",
    }
}

pub(crate) fn type_bits(base: &str) -> u32 {
    match base {
        "u8" | "i8" => 8,
        "u16" | "i16" => 16,
        "u32" | "i32" => 32,
        _ => 32,
    }
}

fn field_rust_type_with_maps(field: &ResolvedField, type_maps: &HashMap<String, String>) -> String {
    if field.resolved_type.sub_decoder.is_some() {
        let sd_name = field.resolved_type.sub_decoder.as_ref().unwrap();
        format!("{}Insn", sd_name)
    } else if let Some(ref alias) = field.resolved_type.alias_name {
        if let Some(rust_type) = type_maps.get(alias) {
            // Use the last segment of the Rust path as the type name
            rust_type
                .rsplit("::")
                .next()
                .unwrap_or(rust_type)
                .to_string()
        } else {
            field.resolved_type.base_type.clone()
        }
    } else {
        field.resolved_type.base_type.clone()
    }
}

/// Convert a DSL instruction name to PascalCase.
/// - `addi` -> `Addi`
/// - `ld_b_c` -> `LdBC`
/// - `ADD` -> `Add`
pub fn to_pascal_case(name: &str) -> String {
    let mut result = String::new();
    let mut capitalize_next = true;

    for ch in name.chars() {
        if ch == '_' {
            capitalize_next = true;
        } else if capitalize_next {
            result.push(ch.to_ascii_uppercase());
            capitalize_next = false;
        } else {
            result.push(ch.to_ascii_lowercase());
        }
    }

    result
}

/// Generate sub-decoder types, pre-baked string arrays, and inline dispatch function.
fn generate_subdecoder(out: &mut String, sd: &ValidatedSubDecoder, dispatch: crate::Dispatch) {
    let opcode_enum = format!("{}Opcode", sd.name);
    let insn_struct = format!("{}Insn", sd.name);
    let decode_fn = format!("decode_{}", to_snake_case(&sd.name));
    let word_type = word_type_for_width(sd.width);
    let lut_size = 1usize << sd.width;
    let use_inline = dispatch == crate::Dispatch::JumpTable;

    // Generate map functions for sub-decoder maps
    if !sd.maps.is_empty() {
        for map in &sd.maps {
            let has_interpolation = map.entries.iter().any(|entry| {
                entry
                    .output
                    .iter()
                    .any(|p| matches!(p, FormatPiece::FieldRef { .. }))
            });
            let return_type = if has_interpolation {
                "String"
            } else {
                "&'static str"
            };
            let params: Vec<String> = map
                .params
                .iter()
                .map(|name| format!("{}: {}", name, word_type))
                .collect();
            writeln!(out, "#[allow(dead_code)]").unwrap();
            writeln!(
                out,
                "fn {}({}) -> {} {{",
                map.name,
                params.join(", "),
                return_type
            )
            .unwrap();
            if map.params.len() == 1 {
                writeln!(out, "    match {} {{", map.params[0]).unwrap();
            } else {
                let tuple: Vec<&str> = map.params.iter().map(|s| s.as_str()).collect();
                writeln!(out, "    match ({}) {{", tuple.join(", ")).unwrap();
            }
            let mut default_entry = None;
            for entry in &map.entries {
                let is_all_wildcard = entry.keys.iter().all(|k| matches!(k, MapKey::Wildcard));
                if is_all_wildcard {
                    default_entry = Some(entry);
                    continue;
                }
                let pattern = if map.params.len() == 1 {
                    format_map_key(&entry.keys[0])
                } else {
                    let keys: Vec<String> = entry.keys.iter().map(|k| format_map_key(k)).collect();
                    format!("({})", keys.join(", "))
                };
                let output = format_map_output(&entry.output, has_interpolation);
                writeln!(out, "        {} => {},", pattern, output).unwrap();
            }
            if let Some(entry) = default_entry {
                let output = format_map_output(&entry.output, has_interpolation);
                writeln!(out, "        _ => {},", output).unwrap();
            } else if has_interpolation {
                writeln!(out, "        _ => String::from(\"???\"),").unwrap();
            } else {
                writeln!(out, "        _ => \"???\",").unwrap();
            }
            writeln!(out, "    }}").unwrap();
            writeln!(out, "}}").unwrap();
            writeln!(out).unwrap();
        }
    }

    // Opcode enum
    writeln!(out, "#[derive(Debug, Clone, Copy, PartialEq, Eq)]").unwrap();
    writeln!(out, "pub enum {} {{", opcode_enum).unwrap();
    for instr in &sd.instructions {
        writeln!(out, "    {},", to_pascal_case(&instr.name)).unwrap();
    }
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // Instruction struct
    writeln!(out, "#[derive(Debug, Clone, Copy, PartialEq, Eq)]").unwrap();
    writeln!(out, "pub struct {} {{", insn_struct).unwrap();
    writeln!(out, "    pub opcode: {},", opcode_enum).unwrap();
    for frag_name in &sd.fragment_names {
        writeln!(out, "    pub {}: &'static str,", frag_name).unwrap();
    }
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // Pre-baked static string arrays and handler functions for each instruction
    for instr in &sd.instructions {
        let handler_name = format!("_sd_{}", instr.name);

        // For each fragment, check if it has interpolation
        // If it does, we need to pre-compute all possible strings
        let has_fields = !instr.resolved_fields.is_empty();

        if has_fields {
            // Generate static arrays for fragments that reference fields
            for frag in &instr.fragments {
                let has_field_refs = frag
                    .pieces
                    .iter()
                    .any(|p| matches!(p, FormatPiece::FieldRef { .. }));
                if has_field_refs {
                    // Compute all possible string values
                    generate_prebaked_fragment_array(out, sd, instr, frag);
                }
            }
        }

        // Generate handler function
        if use_inline {
            writeln!(out, "#[inline(always)]").unwrap();
        }
        let param_name = if has_fields { "val" } else { "_val" };
        writeln!(
            out,
            "fn {}({}: {}) -> {} {{",
            handler_name, param_name, word_type, insn_struct
        )
        .unwrap();

        // Extract fields
        for field in &instr.resolved_fields {
            let extract = extract_expression("val", &field.ranges, word_type, sd.width / 8, "be");
            let expr = apply_transforms(&extract, &field.resolved_type);
            writeln!(out, "    let {} = {};", field.name, expr).unwrap();
        }

        // Build the struct
        writeln!(out, "    {} {{", insn_struct).unwrap();
        writeln!(
            out,
            "        opcode: {}::{},",
            opcode_enum,
            to_pascal_case(&instr.name)
        )
        .unwrap();

        for frag in &instr.fragments {
            let has_field_refs = frag
                .pieces
                .iter()
                .any(|p| matches!(p, FormatPiece::FieldRef { .. }));
            if has_field_refs {
                // Use pre-baked array lookup
                let array_name = prebaked_array_name(&instr.name, &frag.name);
                let index_expr = prebaked_index_expr(instr);
                writeln!(
                    out,
                    "        {}: {}[{} as usize],",
                    frag.name, array_name, index_expr
                )
                .unwrap();
            } else {
                // Static literal string
                let literal = pieces_to_static_str(&frag.pieces);
                writeln!(out, "        {}: \"{}\",", frag.name, literal).unwrap();
            }
        }

        writeln!(out, "    }}").unwrap();
        writeln!(out, "}}").unwrap();
        writeln!(out).unwrap();
    }

    // Build the dispatch match table
    // First, build a mapping from each value 0..2^width to the instruction that matches it
    let mut dispatch_table: Vec<Option<usize>> = vec![None; lut_size];

    for (instr_idx, instr) in sd.instructions.iter().enumerate() {
        // For each possible value, check if this instruction's fixed bits match
        for val in 0..lut_size {
            let matches = instr.segments.iter().all(|seg| {
                if let Segment::Fixed {
                    ranges, pattern, ..
                } = seg
                {
                    let mut bit_idx = 0;
                    for range in ranges {
                        for i in 0..range.width() as usize {
                            if bit_idx < pattern.len() {
                                let hw_bit = range.start - i as u32;
                                let bit_val = (val >> hw_bit) & 1;
                                match pattern[bit_idx] {
                                    Bit::Zero if bit_val != 0 => return false,
                                    Bit::One if bit_val != 1 => return false,
                                    _ => {}
                                }
                                bit_idx += 1;
                            }
                        }
                    }
                    true
                } else {
                    true
                }
            });
            if matches && dispatch_table[val].is_none() {
                dispatch_table[val] = Some(instr_idx);
            }
        }
    }

    match dispatch {
        crate::Dispatch::JumpTable => {
            // Generate a match statement with #[inline(always)]
            writeln!(out, "#[inline(always)]").unwrap();
            writeln!(
                out,
                "pub fn {}(val: {}) -> Option<{}> {{",
                decode_fn, word_type, insn_struct
            )
            .unwrap();
            writeln!(out, "    match val {{").unwrap();

            let mut i = 0;
            while i < lut_size {
                let current = dispatch_table[i];
                let start = i;
                while i < lut_size && dispatch_table[i] == current {
                    i += 1;
                }
                let end = i - 1;

                let handler_name = current.map(|idx| format!("_sd_{}", sd.instructions[idx].name));

                let pattern = if start == end {
                    format!("{:#x}", start)
                } else {
                    format!("{:#x}..={:#x}", start, end)
                };

                match handler_name {
                    Some(h) => writeln!(out, "        {} => Some({}(val)),", pattern, h).unwrap(),
                    None => writeln!(out, "        {} => None,", pattern).unwrap(),
                }
            }

            writeln!(out, "    }}").unwrap();
            writeln!(out, "}}").unwrap();
            writeln!(out).unwrap();
        }
        crate::Dispatch::FnPtrLut => {
            // Generate a static fn ptr array and a dispatch function that indexes into it
            let handler_type = format!("fn({}) -> {}", word_type, insn_struct);
            let table_name = format!("_SD_{}_LUT", sd.name.to_uppercase());

            writeln!(
                out,
                "static {}: [Option<{}>; {}] = [",
                table_name, handler_type, lut_size
            )
            .unwrap();
            for val in 0..lut_size {
                match dispatch_table[val] {
                    Some(idx) => {
                        writeln!(out, "    Some(_sd_{}),", sd.instructions[idx].name).unwrap()
                    }
                    None => writeln!(out, "    None,").unwrap(),
                }
            }
            writeln!(out, "];").unwrap();
            writeln!(out).unwrap();

            writeln!(
                out,
                "pub fn {}(val: {}) -> Option<{}> {{",
                decode_fn, word_type, insn_struct
            )
            .unwrap();
            writeln!(out, "    {}[val as usize].map(|f| f(val))", table_name).unwrap();
            writeln!(out, "}}").unwrap();
            writeln!(out).unwrap();
        }
    }
}

/// Generate a pre-baked static string array for a fragment that contains field interpolation.
fn generate_prebaked_fragment_array(
    out: &mut String,
    sd: &ValidatedSubDecoder,
    instr: &ValidatedSubInstruction,
    frag: &FragmentLine,
) {
    let array_name = prebaked_array_name(&instr.name, &frag.name);
    // Total combinations = product of all field value counts
    let total = prebaked_total_entries(instr);

    writeln!(out, "static {}: [&str; {}] = [", array_name, total).unwrap();

    for combo_idx in 0..total {
        // Compute field values for this combination
        let field_values = prebaked_field_values(instr, combo_idx);
        let s = evaluate_fragment_pieces(&frag.pieces, &field_values, &sd.maps);
        writeln!(out, "    \"{}\",", s).unwrap();
    }

    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();
}

/// Evaluate fragment pieces with concrete field values to produce a string.
fn evaluate_fragment_pieces(
    pieces: &[FormatPiece],
    field_values: &HashMap<String, i64>,
    maps: &[MapDef],
) -> String {
    evaluate_fragment_pieces_str(pieces, field_values, maps)
}

/// Evaluate a format expression with concrete field values.
fn evaluate_format_expr(
    expr: &FormatExpr,
    field_values: &HashMap<String, i64>,
    maps: &[MapDef],
) -> i64 {
    match expr {
        FormatExpr::Field(name) => *field_values.get(name).unwrap_or(&0),
        FormatExpr::IntLiteral(v) => *v,
        FormatExpr::Arithmetic { left, op, right } => {
            let l = evaluate_format_expr(left, field_values, maps);
            let r = evaluate_format_expr(right, field_values, maps);
            match op {
                ArithOp::Add => l + r,
                ArithOp::Sub => l - r,
                ArithOp::Mul => l * r,
                ArithOp::Div => {
                    if r != 0 {
                        l / r
                    } else {
                        0
                    }
                }
                ArithOp::Mod => {
                    if r != 0 {
                        l % r
                    } else {
                        0
                    }
                }
            }
        }
        _ => 0,
    }
}

/// Evaluate a format expression that produces a string (for map calls in fragments).
fn evaluate_fragment_pieces_str(
    pieces: &[FormatPiece],
    field_values: &HashMap<String, i64>,
    maps: &[MapDef],
) -> String {
    let mut result = String::new();
    for piece in pieces {
        match piece {
            FormatPiece::Literal(s) => result.push_str(s),
            FormatPiece::FieldRef { expr, spec } => {
                match expr {
                    FormatExpr::MapCall { map_name, args } => {
                        // Evaluate map call
                        if let Some(map) = maps.iter().find(|m| m.name == *map_name) {
                            let arg_vals: Vec<i64> = args
                                .iter()
                                .map(|a| evaluate_format_expr(a, field_values, maps))
                                .collect();
                            let map_result = evaluate_map(map, &arg_vals);
                            result.push_str(&map_result);
                        }
                    }
                    FormatExpr::Field(name) => {
                        let val = *field_values.get(name).unwrap_or(&0);
                        match spec {
                            Some(s) if s.contains('x') || s.contains('X') => {
                                result.push_str(&format!("{:x}", val));
                            }
                            _ => result.push_str(&val.to_string()),
                        }
                    }
                    _ => {
                        let val = evaluate_format_expr(expr, field_values, maps);
                        result.push_str(&val.to_string());
                    }
                }
            }
        }
    }
    result
}

/// Evaluate a map lookup with concrete key values.
fn evaluate_map(map: &MapDef, keys: &[i64]) -> String {
    for entry in &map.entries {
        let matches = entry.keys.iter().zip(keys.iter()).all(|(k, v)| match k {
            MapKey::Value(expected) => *expected == *v,
            MapKey::Wildcard => true,
        });
        if matches {
            // Entry output is Vec<FormatPiece>, for static maps it should be all literals
            let mut s = String::new();
            for piece in &entry.output {
                if let FormatPiece::Literal(lit) = piece {
                    s.push_str(lit);
                }
            }
            return s;
        }
    }
    "???".to_string()
}

/// Get total number of pre-baked entries (product of all field ranges).
fn prebaked_total_entries(instr: &ValidatedSubInstruction) -> usize {
    let mut total = 1usize;
    for field in &instr.resolved_fields {
        let field_bits: u32 = field.ranges.iter().map(|r| r.width()).sum();
        total *= 1 << field_bits;
    }
    total
}

/// Compute field values for a given combination index.
fn prebaked_field_values(
    instr: &ValidatedSubInstruction,
    combo_idx: usize,
) -> HashMap<String, i64> {
    let mut values = HashMap::new();
    let mut remaining = combo_idx;

    for field in instr.resolved_fields.iter().rev() {
        let field_bits: u32 = field.ranges.iter().map(|r| r.width()).sum();
        let field_range = 1 << field_bits;
        let val = remaining % field_range;
        remaining /= field_range;
        values.insert(field.name.clone(), val as i64);
    }

    values
}

/// Generate the index expression for looking up in the pre-baked array.
fn prebaked_index_expr(instr: &ValidatedSubInstruction) -> String {
    if instr.resolved_fields.len() == 1 {
        return instr.resolved_fields[0].name.clone();
    }

    // Multi-field: combine into a single index
    let mut parts = Vec::new();
    let mut accumulated_bits = 0u32;

    for field in instr.resolved_fields.iter().rev() {
        let field_bits: u32 = field.ranges.iter().map(|r| r.width()).sum();
        if accumulated_bits > 0 {
            parts.push(format!(
                "(({} as usize) << {})",
                field.name, accumulated_bits
            ));
        } else {
            parts.push(format!("({} as usize)", field.name));
        }
        accumulated_bits += field_bits;
    }

    parts.reverse();
    parts.join(" | ")
}

/// Get the name of a pre-baked array for a specific instruction and fragment.
fn prebaked_array_name(instr_name: &str, frag_name: &str) -> String {
    format!(
        "_SD_{}_{}",
        instr_name.to_uppercase(),
        frag_name.to_uppercase()
    )
}

/// Convert pieces that are all literals into a single string.
fn pieces_to_static_str(pieces: &[FormatPiece]) -> String {
    let mut s = String::new();
    for piece in pieces {
        if let FormatPiece::Literal(lit) = piece {
            s.push_str(lit);
        }
    }
    s
}

/// Convert a name to snake_case (for function names).
fn to_snake_case(name: &str) -> String {
    let mut result = String::new();
    for (i, ch) in name.chars().enumerate() {
        if ch.is_ascii_uppercase() && i > 0 {
            result.push('_');
        }
        result.push(ch.to_ascii_lowercase());
    }
    result
}

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

    #[test]
    fn test_to_pascal_case() {
        assert_eq!(to_pascal_case("addi"), "Addi");
        assert_eq!(to_pascal_case("ld_b_c"), "LdBC");
        assert_eq!(to_pascal_case("ADD"), "Add");
        assert_eq!(to_pascal_case("nop"), "Nop");
    }

    #[test]
    fn test_extract_expression() {
        // Extract bits [5:0] (6 bits from position 5 down to 0)
        let range = BitRange::new(5, 0);
        assert_eq!(
            extract_expression("opcode", &[range], "u32", 4, "be"),
            "opcode & 0x3f"
        );

        // Extract bits [31:26] (6 bits from position 31 down to 26)
        let range = BitRange::new(31, 26);
        assert_eq!(
            extract_expression("opcode", &[range], "u32", 4, "be"),
            "(opcode >> 26) & 0x3f"
        );

        // Extract bits from unit 1 in variable-length mode (width=16, unit_bytes=2)
        let range = BitRange::new_in_unit(1, 15, 0);
        assert_eq!(
            extract_expression("opcode", &[range], "u16", 2, "be"),
            "u16::from_be_bytes(data[2..4].try_into().unwrap()) & 0xffff"
        );

        // Extract cross-unit field (width=16, bits [8:23] -> 2 units)
        // Unit 0: bits 8-15 (8 bits), Unit 1: bits 0-7 (8 bits)
        let range0 = BitRange::new_in_unit(0, 7, 0); // 8 bits from unit 0
        let range1 = BitRange::new_in_unit(1, 15, 8); // 8 bits from unit 1
        let result = extract_expression("opcode", &[range0, range1], "u16", 2, "be");
        assert!(result.contains("opcode & 0xff"));
        assert!(result.contains("u16::from_be_bytes(data[2..4].try_into().unwrap())"));
    }
}