elf-magic 0.2.7

use std::{fs, path::Path};

use minijinja::{context, Environment};

use crate::{error::Error, programs::BuildResult};

/// Template for the generated lib.rs file
const LIB_TEMPLATE: &str = r#"
// This file is auto-generated by elf-magic
// DO NOT EDIT MANUALLY - Changes will be overwritten
//
// Build Status:
{% for status in build_status -%}
// {{ status.icon }} {{ status.program_name }} - {{ status.message }}
{% endfor -%}

{% for constant in constants -%}
/// ELF binary for the {{ constant.program_name }} Solana program
pub const {{ constant.constant_name }}: &[u8] = include_bytes!(env!("{{ constant.env_var }}"));

{% endfor -%}

/// Get all available Solana program ELF binaries
/// Returns a vector of (program_name, elf_bytes) tuples
pub fn elves() -> Vec<(&'static str, &'static [u8])> {
    vec![
{%- for constant in constants %}
        ("{{ constant.program_name }}", {{ constant.constant_name }}),
{%- endfor %}
    ]
}
"#;

/// Generate code for Solana programs from build results
pub fn generate(build_result: &BuildResult) -> Result<String, Error> {
    // Convert successful programs to template data
    let template_data = build_result
        .successful
        .iter()
        .map(|(program, _path)| {
            serde_json::json!({
                "constant_name": program.constant_name(),
                "env_var": program.env_var_name(),
                "program_name": program.target_name.clone()
            })
        })
        .collect::<Vec<_>>();

    // Create build status entries
    let mut build_status = Vec::new();

    // Add successful builds
    for (program, _path) in &build_result.successful {
        build_status.push(serde_json::json!({
            "icon": "✓",
            "program_name": program.target_name,
            "message": "SUCCESS"
        }));
    }

    // Add failed builds
    for (program, error) in &build_result.failed {
        build_status.push(serde_json::json!({
            "icon": "✗",
            "program_name": program.target_name,
            "message": format!("FAILED: {}", error)
        }));
    }

    // Create minijinja environment and render template
    let mut env = Environment::new();

    env.add_template("lib", LIB_TEMPLATE).map_err(|e| {
        let msg = format!("Failed to add template: {}", e);
        Error::CodeGeneration(msg)
    })?;

    let template = env.get_template("lib").map_err(|e| {
        let msg = format!("Failed to get template: {}", e);
        Error::CodeGeneration(msg)
    })?;

    let rendered_content = template
        .render(context! {
            constants => template_data,
            build_status => build_status,
        })
        .map_err(|e| {
            let msg = format!("Failed to render template: {}", e);
            Error::CodeGeneration(msg)
        })?;

    Ok(rendered_content)
}

/// Write generated code to lib.rs
pub fn save(manifest_dir: &Path, code: &str) -> Result<(), Error> {
    let output_path = manifest_dir.join("src").join("lib.rs");
    fs::write(&output_path, code).map_err(|e| {
        let message = format!("Failed to write lib.rs: {}", e);
        Error::CodeGeneration(message)
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::programs::SolanaProgram;
    use std::path::PathBuf;

    fn sample_programs() -> Vec<SolanaProgram> {
        vec![
            SolanaProgram {
                package_name: "package1".to_string(),
                target_name: "target1".to_string(),
                manifest_path: PathBuf::from("/path/to/Cargo.toml"),
            },
            SolanaProgram {
                package_name: "package2".to_string(),
                target_name: "target2".to_string(),
                manifest_path: PathBuf::from("/path/to/other/Cargo.toml"),
            },
        ]
    }

    #[test]
    fn test_generate_empty_programs() {
        let result = generate(&BuildResult {
            successful: Vec::new(),
            failed: Vec::new(),
        })
        .unwrap();

        // Should generate valid Rust code with empty elves function
        assert!(result.contains("pub fn elves() -> Vec<(&'static str, &'static [u8])> {"));
        assert!(result.contains("vec![\n    ]"));
        assert!(result.contains("// This file is auto-generated by elf-magic"));
    }

    #[test]
    fn test_generate_single_program() {
        let programs = [SolanaProgram {
            package_name: "my_package".to_string(),
            target_name: "my_target".to_string(),
            manifest_path: PathBuf::from("/path/to/Cargo.toml"),
        }];

        let result = generate(&BuildResult {
            successful: programs
                .iter()
                .map(|p| (p.clone(), PathBuf::from("/path/to/Cargo.toml")))
                .collect(),
            failed: Vec::new(),
        })
        .unwrap();

        // Check constant definition
        assert!(result.contains(
            "pub const MY_TARGET_ELF: &[u8] = include_bytes!(env!(\"MY_TARGET_ELF_PATH\"));"
        ));

        // Check elves function includes the program
        assert!(result.contains("(\"my_target\", MY_TARGET_ELF),"));

        // Check doc comment
        assert!(result.contains("/// ELF binary for the my_target Solana program"));
    }

    #[test]
    fn test_generate_multiple_programs() {
        let programs = sample_programs();
        let result = generate(&BuildResult {
            successful: programs
                .iter()
                .map(|p| (p.clone(), PathBuf::from("/path/to/Cargo.toml")))
                .collect(),
            failed: Vec::new(),
        })
        .unwrap();

        // Should have both constants
        assert!(result.contains("pub const TARGET1_ELF"));
        assert!(result.contains("pub const TARGET2_ELF"));

        // Should have both environment variables
        assert!(result.contains("env!(\"TARGET1_ELF_PATH\")"));
        assert!(result.contains("env!(\"TARGET2_ELF_PATH\")"));

        // Should include both in elves function
        assert!(result.contains("(\"target1\", TARGET1_ELF),"));
        assert!(result.contains("(\"target2\", TARGET2_ELF),"));

        // Should have proper doc comments
        assert!(result.contains("/// ELF binary for the target1 Solana program"));
        assert!(result.contains("/// ELF binary for the target2 Solana program"));
    }

    #[test]
    fn test_generate_with_special_characters() {
        let programs = [SolanaProgram {
            package_name: "my-special-package".to_string(),
            target_name: "my_target_name".to_string(),
            manifest_path: PathBuf::from("/path/to/Cargo.toml"),
        }];

        let result = generate(&BuildResult {
            successful: programs
                .iter()
                .map(|p| (p.clone(), PathBuf::from("/path/to/Cargo.toml")))
                .collect(),
            failed: Vec::new(),
        })
        .unwrap();

        // Target name should be preserved as-is in program name
        assert!(result.contains("(\"my_target_name\", MY_TARGET_NAME_ELF),"));

        // But constant and env var should follow their respective conventions
        assert!(result.contains("pub const MY_TARGET_NAME_ELF"));
        assert!(result.contains("env!(\"MY_TARGET_NAME_ELF_PATH\")"));
    }

    #[test]
    fn test_generated_code_is_valid_rust() {
        let programs = sample_programs();
        let result = generate(&BuildResult {
            successful: programs
                .iter()
                .map(|p| (p.clone(), PathBuf::from("/path/to/Cargo.toml")))
                .collect(),
            failed: Vec::new(),
        })
        .unwrap();

        // Basic syntax checks
        assert!(result.contains("pub const"));
        assert!(result.contains("pub fn elves()"));
        assert!(result.contains("vec!["));
        assert!(!result.contains("{{")); // No unresolved template variables
        assert!(!result.contains("}}"));

        // Should be properly formatted
        let lines: Vec<&str> = result.lines().collect();
        assert!(lines.len() > 5); // Should have multiple lines

        // Comments should be present
        assert!(result.contains("// This file is auto-generated by elf-magic"));
        assert!(result.contains("// DO NOT EDIT MANUALLY"));
    }

    #[test]
    fn test_template_render_error_handling() {
        // This is harder to test without breaking the template
        // But we can at least verify the function signature works
        let _programs: Vec<SolanaProgram> = vec![];
        let result = generate(&BuildResult {
            successful: Vec::new(),
            failed: Vec::new(),
        });
        assert!(result.is_ok());
    }

    #[test]
    fn test_generate_with_build_status() {
        let successful_program = SolanaProgram {
            package_name: "good_package".to_string(),
            target_name: "good_program".to_string(),
            manifest_path: PathBuf::from("/path/to/good/Cargo.toml"),
        };

        let failed_program = SolanaProgram {
            package_name: "bad_package".to_string(),
            target_name: "bad_program".to_string(),
            manifest_path: PathBuf::from("/path/to/bad/Cargo.toml"),
        };

        let build_result = BuildResult {
            successful: vec![(successful_program, PathBuf::from("/tmp/good_program.so"))],
            failed: vec![(
                failed_program,
                Error::ProgramBuild {
                    program: "bad_program".to_string(),
                    error: "wasi crate contains multiple cdylib targets".to_string(),
                },
            )],
        };

        let result = generate(&build_result).unwrap();

        // Should contain build status comments
        assert!(result.contains("// Build Status:"));
        assert!(result.contains("// ✓ good_program - SUCCESS"));
        assert!(result.contains("// ✗ bad_program - FAILED: Failed to build program bad_program: wasi crate contains multiple cdylib targets"));

        // Should only include successful program in constants
        assert!(result.contains("pub const GOOD_PROGRAM_ELF"));
        assert!(!result.contains("pub const BAD_PROGRAM_ELF"));

        // Should only include successful program in elves function
        assert!(result.contains("(\"good_program\", GOOD_PROGRAM_ELF),"));
        assert!(!result.contains("(\"bad_program\""));
    }
}