rustnetconf-yang 0.1.4

//! Build script: generates Rust structs from YANG models using libyang2.

use std::collections::HashSet;
use std::env;
use std::fmt::Write as FmtWrite;
use std::fs;
use std::path::PathBuf;

use yang2::context::{Context, ContextFlags};
use yang2::schema::{DataValueType, SchemaNode, SchemaNodeKind};

fn main() {
    let yang_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()).join("yang-models");
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
    let output_file = out_dir.join("yang_generated.rs");

    println!("cargo:rerun-if-changed={}", yang_dir.display());

    if !yang_dir.exists() {
        fs::write(&output_file, "// No yang-models/ directory found.\n").unwrap();
        return;
    }

    // Collect .yang filenames (just the module names, not full paths)
    let yang_files: Vec<(String, Option<String>)> = fs::read_dir(&yang_dir)
        .unwrap()
        .filter_map(|entry| {
            let entry = entry.ok()?;
            let path = entry.path();
            if path.extension().map(|e| e == "yang").unwrap_or(false) {
                println!("cargo:rerun-if-changed={}", path.display());
                let stem = path.file_stem()?.to_str()?.to_string();
                // Parse "module-name@revision" format
                if let Some((name, rev)) = stem.split_once('@') {
                    Some((name.to_string(), Some(rev.to_string())))
                } else {
                    Some((stem, None))
                }
            } else {
                None
            }
        })
        .collect();

    if yang_files.is_empty() {
        fs::write(&output_file, "// No YANG models found in yang-models/\n").unwrap();
        return;
    }

    // Create libyang2 context with search path
    let mut ctx =
        Context::new(ContextFlags::NO_YANGLIBRARY).expect("failed to create libyang2 context");
    ctx.set_searchdir(&yang_dir)
        .expect("failed to set yang search directory");

    // Load each module; collect failures and hard-error at the end so users
    // never silently end up with missing generated types.
    let mut module_names = Vec::new();
    let mut load_errors: Vec<String> = Vec::new();
    for (name, revision) in &yang_files {
        match ctx.load_module(name, revision.as_deref(), &[]) {
            Ok(_module) => {
                println!("cargo:warning=Loaded YANG module: {name}");
                module_names.push(name.clone());
            }
            Err(e) => {
                let msg = format!("Failed to load YANG module '{name}': {e}");
                println!("cargo:warning={msg}");
                load_errors.push(msg);
            }
        }
    }
    if !load_errors.is_empty() {
        panic!(
            "rustnetconf-yang: {} YANG module(s) failed to load — types were NOT generated.\n{}",
            load_errors.len(),
            load_errors.join("\n")
        );
    }

    // Generate Rust code
    let mut output = String::new();
    writeln!(output, "// Auto-generated from YANG models. Do not edit.").unwrap();
    writeln!(output, "// Generated by rustnetconf-yang build.rs").unwrap();
    writeln!(output).unwrap();
    writeln!(output, "use serde::{{Serialize, Deserialize}};").unwrap();
    writeln!(output).unwrap();

    // Iterate loaded modules and generate code
    for module in ctx.modules(true) {
        let name = module.name().to_string();
        // Only generate for explicitly loaded models (skip internal/imported-only)
        if !module_names.contains(&name) {
            continue;
        }
        let namespace = module.namespace().to_string();
        let mod_name = name.replace('-', "_");

        writeln!(output, "/// Generated from YANG module `{name}`").unwrap();
        writeln!(output, "/// Namespace: `{namespace}`").unwrap();
        writeln!(output, "pub mod {mod_name} {{").unwrap();
        writeln!(output, "    #[allow(unused_imports)]").unwrap();
        writeln!(output, "    use super::*;").unwrap();
        writeln!(output, "    #[allow(unused_imports)]").unwrap();
        writeln!(output, "    use crate::serialize::*;").unwrap();
        writeln!(output).unwrap();
        writeln!(output, "    /// Namespace URI for this YANG module.").unwrap();
        writeln!(output, "    pub const NAMESPACE: &str = \"{namespace}\";").unwrap();
        writeln!(output).unwrap();

        // Generate structs for top-level data nodes
        let mut emitted = HashSet::new();
        for node in module.data() {
            generate_node(&mut output, &node, 1, true, &mut emitted);
        }

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

    fs::write(&output_file, &output).unwrap();
    eprintln!(
        "cargo:warning=Generated YANG types to {}",
        output_file.display()
    );
}

/// Generate a Rust struct for a YANG container or list node.
fn generate_node(
    output: &mut String,
    node: &SchemaNode,
    indent: usize,
    is_top_level: bool,
    emitted: &mut HashSet<String>,
) {
    let ind = "    ".repeat(indent);
    let node_name = node.name().to_string();
    let rust_name = to_rust_type_name(&node_name);

    // Skip if already emitted (augmented nodes can appear in multiple modules)
    if emitted.contains(&rust_name) {
        return;
    }

    match node.kind() {
        SchemaNodeKind::Container => {
            emitted.insert(rust_name.clone());

            writeln!(output, "{ind}/// YANG container: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[derive(Debug, Clone, Default, Serialize, Deserialize)]"
            )
            .unwrap();
            writeln!(output, "{ind}pub struct {rust_name} {{").unwrap();

            for child in node.children() {
                generate_field(output, &child, indent + 1);
            }

            writeln!(output, "{ind}}}").unwrap();
            writeln!(output).unwrap();

            // Generate WriteXmlFields for all containers (top-level and nested)
            generate_write_xml_fields_impl(output, node, &rust_name, indent);

            if is_top_level {
                generate_to_xml_impl(output, &rust_name, &node_name, indent);
            }

            for child in node.children() {
                if matches!(
                    child.kind(),
                    SchemaNodeKind::Container | SchemaNodeKind::List
                ) {
                    generate_node(output, &child, indent, false, emitted);
                }
            }
        }
        SchemaNodeKind::List => {
            emitted.insert(rust_name.clone());

            writeln!(output, "{ind}/// YANG list entry: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[derive(Debug, Clone, Default, Serialize, Deserialize)]"
            )
            .unwrap();
            writeln!(output, "{ind}pub struct {rust_name} {{").unwrap();

            for child in node.children() {
                generate_field(output, &child, indent + 1);
            }

            writeln!(output, "{ind}}}").unwrap();
            writeln!(output).unwrap();

            // Generate WriteXmlFields for list entries so they can be embedded
            generate_write_xml_fields_impl(output, node, &rust_name, indent);

            for child in node.children() {
                if matches!(
                    child.kind(),
                    SchemaNodeKind::Container | SchemaNodeKind::List
                ) {
                    generate_node(output, &child, indent, false, emitted);
                }
            }
        }
        _ => {}
    }
}

/// Generate a struct field for a YANG child node.
fn generate_field(output: &mut String, node: &SchemaNode, indent: usize) {
    let ind = "    ".repeat(indent);
    let node_name = node.name().to_string();
    let field_name = to_rust_field_name(&node_name);

    match node.kind() {
        SchemaNodeKind::Leaf => {
            let rust_type = yang_type_to_rust(node);
            writeln!(output, "{ind}/// YANG leaf: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[serde(skip_serializing_if = \"Option::is_none\")]"
            )
            .unwrap();
            writeln!(output, "{ind}pub {field_name}: Option<{rust_type}>,").unwrap();
        }
        SchemaNodeKind::LeafList => {
            let rust_type = yang_type_to_rust(node);
            writeln!(output, "{ind}/// YANG leaf-list: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[serde(default, skip_serializing_if = \"Vec::is_empty\")]"
            )
            .unwrap();
            writeln!(output, "{ind}pub {field_name}: Vec<{rust_type}>,").unwrap();
        }
        SchemaNodeKind::Container => {
            let type_name = to_rust_type_name(&node_name);
            writeln!(output, "{ind}/// YANG container: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[serde(skip_serializing_if = \"Option::is_none\")]"
            )
            .unwrap();
            writeln!(output, "{ind}pub {field_name}: Option<{type_name}>,").unwrap();
        }
        SchemaNodeKind::List => {
            let type_name = to_rust_type_name(&node_name);
            writeln!(output, "{ind}/// YANG list: `{node_name}`").unwrap();
            writeln!(
                output,
                "{ind}#[serde(default, skip_serializing_if = \"Vec::is_empty\")]"
            )
            .unwrap();
            writeln!(output, "{ind}pub {field_name}: Vec<{type_name}>,").unwrap();
        }
        _ => {} // Skip choice, anyxml, etc. for now
    }
}

/// Generate WriteXmlFields implementation for a container or list entry struct.
///
/// This impl writes all child fields (leaves, leaf-lists, containers, lists)
/// into a caller-supplied writer. It does not write the surrounding element tags.
fn generate_write_xml_fields_impl(
    output: &mut String,
    node: &SchemaNode,
    rust_name: &str,
    indent: usize,
) {
    let ind = "    ".repeat(indent);

    writeln!(output, "{ind}impl WriteXmlFields for {rust_name} {{").unwrap();
    writeln!(output, "{ind}    fn write_xml_fields(&self, writer: &mut Writer<Cursor<Vec<u8>>>) -> Result<(), XmlError> {{").unwrap();

    for child in node.children() {
        let child_name = child.name().to_string();
        let field = to_rust_field_name(&child_name);

        match child.kind() {
            SchemaNodeKind::Leaf => {
                writeln!(
                    output,
                    "{ind}        if let Some(ref val) = self.{field} {{"
                )
                .unwrap();
                writeln!(output, "{ind}            write_text_element(writer, \"{child_name}\", &val.to_string())?;").unwrap();
                writeln!(output, "{ind}        }}").unwrap();
            }
            SchemaNodeKind::LeafList => {
                writeln!(output, "{ind}        for val in &self.{field} {{").unwrap();
                writeln!(output, "{ind}            write_text_element(writer, \"{child_name}\", &val.to_string())?;").unwrap();
                writeln!(output, "{ind}        }}").unwrap();
            }
            SchemaNodeKind::Container => {
                writeln!(
                    output,
                    "{ind}        if let Some(ref child) = self.{field} {{"
                )
                .unwrap();
                writeln!(
                    output,
                    "{ind}            write_element_with_fields(writer, \"{child_name}\", child)?;"
                )
                .unwrap();
                writeln!(output, "{ind}        }}").unwrap();
            }
            SchemaNodeKind::List => {
                writeln!(output, "{ind}        for item in &self.{field} {{").unwrap();
                writeln!(
                    output,
                    "{ind}            write_element_with_fields(writer, \"{child_name}\", item)?;"
                )
                .unwrap();
                writeln!(output, "{ind}        }}").unwrap();
            }
            _ => {
                // Skip choice, anyxml, etc. for now
            }
        }
    }

    writeln!(output, "{ind}        Ok(())").unwrap();
    writeln!(output, "{ind}    }}").unwrap();
    writeln!(output, "{ind}}}").unwrap();
    writeln!(output).unwrap();
}

/// Generate ToNetconfXml implementation for a top-level container.
///
/// Delegates to WriteXmlFields for field serialization so containers and
/// lists at any nesting depth are serialized correctly.
fn generate_to_xml_impl(output: &mut String, rust_name: &str, node_name: &str, indent: usize) {
    let ind = "    ".repeat(indent);

    writeln!(output, "{ind}impl ToNetconfXml for {rust_name} {{").unwrap();
    writeln!(
        output,
        "{ind}    fn namespace(&self) -> &str {{ NAMESPACE }}"
    )
    .unwrap();
    writeln!(
        output,
        "{ind}    fn root_element(&self) -> &str {{ \"{node_name}\" }}"
    )
    .unwrap();
    writeln!(
        output,
        "{ind}    fn to_xml(&self) -> Result<String, XmlError> {{"
    )
    .unwrap();
    writeln!(output, "{ind}        let mut writer = new_writer();").unwrap();
    writeln!(
        output,
        "{ind}        write_start_with_ns(&mut writer, \"{node_name}\", NAMESPACE)?;"
    )
    .unwrap();
    writeln!(output, "{ind}        self.write_xml_fields(&mut writer)?;").unwrap();
    writeln!(
        output,
        "{ind}        write_end(&mut writer, \"{node_name}\")?;"
    )
    .unwrap();
    writeln!(output, "{ind}        finish_writer(writer)").unwrap();
    writeln!(output, "{ind}    }}").unwrap();
    writeln!(output, "{ind}}}").unwrap();
    writeln!(output).unwrap();
}

/// Convert YANG name to Rust PascalCase type name.
fn to_rust_type_name(yang_name: &str) -> String {
    yang_name
        .split('-')
        .map(|part| {
            let mut chars = part.chars();
            match chars.next() {
                None => String::new(),
                Some(c) => c.to_uppercase().to_string() + chars.as_str(),
            }
        })
        .collect()
}

/// Convert YANG name to Rust snake_case field name.
fn to_rust_field_name(yang_name: &str) -> String {
    let name = yang_name.replace('-', "_");
    // Full set of Rust keywords (stable + reserved) that must be escaped.
    match name.as_str() {
        "as" | "async" | "await" | "break" | "const" | "continue" | "crate" | "dyn" | "else"
        | "enum" | "extern" | "false" | "fn" | "for" | "if" | "impl" | "in" | "let" | "loop"
        | "match" | "mod" | "move" | "mut" | "pub" | "ref" | "return" | "self" | "Self"
        | "static" | "struct" | "super" | "trait" | "true" | "type" | "unsafe" | "use"
        | "where" | "while" | "abstract" | "become" | "box" | "do" | "final" | "macro"
        | "override" | "priv" | "try" | "typeof" | "unsized" | "virtual" | "yield" => {
            format!("{name}_")
        }
        _ => name,
    }
}

/// Map YANG leaf type to Rust type.
fn yang_type_to_rust(node: &SchemaNode) -> String {
    if let Some(leaf_type) = node.leaf_type() {
        match leaf_type.base_type() {
            DataValueType::String => "String",
            DataValueType::Bool => "bool",
            DataValueType::Uint8 => "u8",
            DataValueType::Uint16 => "u16",
            DataValueType::Uint32 => "u32",
            DataValueType::Uint64 => "u64",
            DataValueType::Int8 => "i8",
            DataValueType::Int16 => "i16",
            DataValueType::Int32 => "i32",
            DataValueType::Int64 => "i64",
            DataValueType::Empty => "bool",
            DataValueType::Enum => "String",
            DataValueType::Union => "String",
            DataValueType::Binary => "String",
            DataValueType::IdentityRef => "String",
            DataValueType::LeafRef => "String",
            DataValueType::Dec64 => "f64",
            DataValueType::Bits => "String",
            DataValueType::InstanceId => "String",
            _ => "String",
        }
        .to_string()
    } else {
        "String".to_string()
    }
}