elf-magic 0.5.1

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

use minijinja::{context, Environment};

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

/// Template for the generated lib.rs file
const LIB_RS_TEMPLATE: &str = r#"
// Your generated code here
// This file is auto-generated by elf-magic
// DO NOT EDIT - Assume changes will be lost.
//
// Build Status:
{% for build_status in build_statuses -%}
// {{ build_status.icon }} {{ build_status.program_name }} - {{ build_status.message }}
{% endfor -%}
// ------------------------------------------------------------
// Constants
{% 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: &ProgramBuildResult) -> Result<String, Error> {
    // Collect all programs with render data in one pass
    let mut program_specs: Vec<(String, Option<serde_json::Value>, serde_json::Value)> = Vec::new();

    // Process successful programs
    for (program, _path) in &build_result.successful {
        let constant = Some(serde_json::json!({
            "constant_name": program.constant_name,
            "env_var": program.env_var_name(),
            "program_name": program.target_name.clone()
        }));

        let build_status = serde_json::json!({
            "icon": "✓",
            "program_name": program.target_name,
            "message": "SUCCESS"
        });

        program_specs.push((program.target_name.clone(), constant, build_status));
    }

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

        program_specs.push((program.target_name.clone(), None, build_status));
    }

    // Sort once by target_name
    program_specs.sort_by(|a, b| a.0.cmp(&b.0));

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

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

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

    let constants: Vec<_> = program_specs
        .iter()
        .filter_map(|(_, constant, _)| constant.clone())
        .collect();

    let build_statuses: Vec<_> = program_specs
        .iter()
        .map(|(_, _, build_status)| build_status)
        .collect();

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

    Ok(rendered_content)
}

/// Write generated code to OUT_DIR
pub fn save(manifest_dir: &Path, code: &str) -> Result<(), Error> {
    // Generate to OUT_DIR instead of src/ (standard Rust pattern)
    let out_dir = std::env::var("OUT_DIR").map_err(|_| {
        Error::CodeGeneration(
            "OUT_DIR not set - this should only be called from build scripts".to_string(),
        )
    })?;
    let output_path = std::path::Path::new(&out_dir).join("generated.rs");

    fs::write(&output_path, code).map_err(|e| {
        let message = format!("Failed to write generated.rs: {}", e);
        Error::CodeGeneration(message)
    })?;

    // Tell Cargo where the generated file is
    println!(
        "cargo:rustc-env=ELF_MAGIC_GENERATED_PATH={}",
        output_path.display()
    );

    // Format the generated code to prevent unexpected changes when users run cargo fmt
    format_generated_code(manifest_dir, &output_path);

    Ok(())
}

/// Format generated code with cargo fmt, ignoring errors
fn format_generated_code(manifest_dir: &Path, file_path: &Path) {
    let result = Command::new("cargo")
        .args(["fmt", "--manifest-path"])
        .arg(manifest_dir.join("Cargo.toml"))
        .arg("--")
        .arg(file_path)
        .output();

    if let Err(e) = result {
        eprintln!("Warning: Could not run cargo fmt: {}", e);
    } else if let Ok(output) = result {
        if !output.status.success() {
            eprintln!(
                "Warning: Failed to format generated code: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }
    }
}

#[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"),
                constant_name: "TARGET1_ELF".to_string(),
            },
            SolanaProgram {
                package_name: "package2".to_string(),
                target_name: "target2".to_string(),
                manifest_path: PathBuf::from("/path/to/other/Cargo.toml"),
                constant_name: "TARGET2_ELF".to_string(),
            },
        ]
    }

    #[test]
    fn test_generate_empty_programs() {
        let result = generate(&ProgramBuildResult {
            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"),
            constant_name: "MY_TARGET_ELF".to_string(),
        }];

        let result = generate(&ProgramBuildResult {
            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(&ProgramBuildResult {
            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"),
            constant_name: "MY_TARGET_NAME_ELF".to_string(),
        }];

        let result = generate(&ProgramBuildResult {
            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(&ProgramBuildResult {
            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 - Assume changes will be lost."));
    }

    #[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(&ProgramBuildResult {
            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"),
            constant_name: "GOOD_PROGRAM_ELF".to_string(),
        };

        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"),
            constant_name: "BAD_PROGRAM_ELF".to_string(),
        };

        let build_result = ProgramBuildResult {
            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\""));
    }

    #[test]
    fn test_alphabetical_sorting_mixed_success_and_failure() {
        // Create programs with names that should be sorted alphabetically
        let zebra_program = SolanaProgram {
            package_name: "zebra".to_string(),
            target_name: "zebra".to_string(),
            manifest_path: PathBuf::from("/path/to/zebra/Cargo.toml"),
            constant_name: "ZEBRA_ELF".to_string(),
        };

        let alpha_program = SolanaProgram {
            package_name: "alpha".to_string(),
            target_name: "alpha".to_string(),
            manifest_path: PathBuf::from("/path/to/alpha/Cargo.toml"),
            constant_name: "ALPHA_ELF".to_string(),
        };

        let beta_program = SolanaProgram {
            package_name: "beta".to_string(),
            target_name: "beta".to_string(),
            manifest_path: PathBuf::from("/path/to/beta/Cargo.toml"),
            constant_name: "BETA_ELF".to_string(),
        };

        // Mix success and failure to test unified sorting
        let build_result = ProgramBuildResult {
            successful: vec![
                (zebra_program, PathBuf::from("/tmp/zebra.so")),
                (beta_program, PathBuf::from("/tmp/beta.so")),
            ],
            failed: vec![(
                alpha_program,
                Error::ProgramBuild {
                    program: "alpha".to_string(),
                    error: "dependency not found".to_string(),
                },
            )],
        };

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

        // Find the build status section
        let build_status_start = result.find("// Build Status:").unwrap();
        let build_status_section = &result[build_status_start..];

        // Build status should be alphabetical: alpha (failed), beta (success), zebra (success)
        let alpha_pos = build_status_section.find("// ✗ alpha").unwrap();
        let beta_pos = build_status_section.find("// ✓ beta").unwrap();
        let zebra_pos = build_status_section.find("// ✓ zebra").unwrap();

        assert!(
            alpha_pos < beta_pos,
            "alpha should come before beta in build status"
        );
        assert!(
            beta_pos < zebra_pos,
            "beta should come before zebra in build status"
        );

        // Constants should be alphabetical too: beta, zebra (alpha failed so no constant)
        let beta_const_pos = result.find("pub const BETA_ELF").unwrap();
        let zebra_const_pos = result.find("pub const ZEBRA_ELF").unwrap();

        assert!(
            beta_const_pos < zebra_const_pos,
            "BETA_ELF should come before ZEBRA_ELF"
        );

        // elves() function should be alphabetical: beta, zebra
        let elves_section = result.find("pub fn elves()").unwrap();
        let elves_content = &result[elves_section..];

        let beta_elves_pos = elves_content.find("(\"beta\", BETA_ELF)").unwrap();
        let zebra_elves_pos = elves_content.find("(\"zebra\", ZEBRA_ELF)").unwrap();

        assert!(
            beta_elves_pos < zebra_elves_pos,
            "beta should come before zebra in elves()"
        );
    }

    #[test]
    fn test_alphabetical_sorting_case_sensitive() {
        let lower_program = SolanaProgram {
            package_name: "lowercase".to_string(),
            target_name: "lowercase".to_string(),
            manifest_path: PathBuf::from("/path/to/lowercase/Cargo.toml"),
            constant_name: "LOWERCASE_ELF".to_string(),
        };

        let upper_program = SolanaProgram {
            package_name: "UPPERCASE".to_string(),
            target_name: "UPPERCASE".to_string(),
            manifest_path: PathBuf::from("/path/to/UPPERCASE/Cargo.toml"),
            constant_name: "UPPERCASE_ELF".to_string(),
        };

        let build_result = ProgramBuildResult {
            successful: vec![
                (upper_program, PathBuf::from("/tmp/UPPERCASE.so")),
                (lower_program, PathBuf::from("/tmp/lowercase.so")),
            ],
            failed: Vec::new(),
        };

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

        // Build status should be sorted by target_name: UPPERCASE, lowercase (uppercase comes first in ASCII)
        let build_status_start = result.find("// Build Status:").unwrap();
        let build_status_section = &result[build_status_start..];

        let uppercase_pos = build_status_section.find("// ✓ UPPERCASE").unwrap();
        let lowercase_pos = build_status_section.find("// ✓ lowercase").unwrap();

        assert!(
            uppercase_pos < lowercase_pos,
            "UPPERCASE should come before lowercase (uppercase letters come first in ASCII sort)"
        );

        // Constants should follow same order
        let uppercase_const_pos = result.find("pub const UPPERCASE_ELF").unwrap();
        let lowercase_const_pos = result.find("pub const LOWERCASE_ELF").unwrap();

        assert!(
            uppercase_const_pos < lowercase_const_pos,
            "UPPERCASE_ELF should come before LOWERCASE_ELF"
        );
    }
}