metamorph 1.0.0

use rand::prelude::SliceRandom;
use rand::{Rng, thread_rng};
use regex::Regex;
use std::collections::HashMap;
use std::env;
use std::fs;
use std::io::{self, Write};
use std::path::Path;
use std::process::Command;
use std::thread;
use std::time::Duration;

fn _generate_random_ident(prefix: &str) -> String {
    let mut rng = thread_rng();
    let chars: Vec<char> = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
        .chars()
        .collect();
    let len = rng.gen_range(6..15);
    let random_part: String = (0..len).map(|_| *chars.choose(&mut rng).unwrap()).collect();
    if prefix.is_empty() {
        random_part
    } else {
        format!("{}_{}", prefix, random_part)
    }
}

fn _generate_random_comment() -> String {
    let mut rng = thread_rng();
    let chars: Vec<char> = "abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*"
        .chars()
        .collect();
    let len = rng.gen_range(10..50);
    let random_str: String = (0..len).map(|_| *chars.choose(&mut rng).unwrap()).collect();
    format!("// {}", random_str)
}

fn _generate_junk_function() -> String {
    let mut rng = thread_rng();
    let func_name = _generate_random_ident("func");
    let var_count = rng.gen_range(1..4);
    let mut body = Vec::new();

    for _ in 0..var_count {
        let var_name = _generate_random_ident("var");
        let value_type = rng.gen_range(0..3);
        let value = match value_type {
            0 => format!("{}", rng.gen_range(i32::MIN..i32::MAX)),
            1 => format!("{:.2}", rng.gen_range(0.0..1.0)),
            2 => format!("\"{}\"", _generate_random_ident("str").replace("\"", "")),
            _ => unreachable!(),
        };

        let declaration = format!("    let {} = {};", var_name, value);
        body.push(declaration);
    }

    if rng.gen_bool(0.5) {
        body.insert(0, format!("    {}", _generate_random_comment()));
    }

    format!("fn {}() {{\n{}\n}}", func_name, body.join("\n"))
}

fn _transform_source(
    source_code: &str,
    original_content: &str,
    junk_function_count: usize,
) -> io::Result<String> {
    let marker_start = "// METAMORPHIC_MARKER_START";
    let marker_end = "// METAMORPHIC_MARKER_END";

    let start_idx = source_code.find(marker_start);
    let end_idx = source_code.find(marker_end);

    let (before, after) = if let (Some(start), Some(end)) = (start_idx, end_idx) {
        (
            &source_code[..start],
            &source_code[end + marker_end.len()..],
        )
    } else {
        return Err(io::Error::new(io::ErrorKind::Other, "Marker not found"));
    };

    // Apply variable renaming to original_content
    let re = Regex::new(r"\br_[a-zA-Z0-9_]*\b").unwrap();
    let mut replacements = HashMap::new();
    for cap in re.captures_iter(original_content) {
        let var_name = cap[0].to_string();
        if !replacements.contains_key(&var_name) {
            replacements.insert(
                var_name.clone(),
                format!("r_{}", _generate_random_ident("")),
            );
        }
    }

    let mut modified_content = original_content.to_string();
    for (old_name, new_name) in &replacements {
        modified_content = modified_content.replace(old_name, new_name);
    }

    // Remove existing func_ functions from modified_content
    let func_re = Regex::new(r"(?s)fn func_[a-zA-Z0-9_]+\s*\(\)\s*\{.*?\}").unwrap();
    let cleaned_content = func_re
        .replace_all(&modified_content, "")
        .trim()
        .to_string();

    // Generate new junk functions
    let mut junk_functions = Vec::new();
    for _ in 0..junk_function_count {
        junk_functions.push(_generate_junk_function());
    }

    // Combine cleaned content and new junk functions
    let mut combined_blocks = vec![cleaned_content];
    combined_blocks.extend(junk_functions);

    // Reconstruct the full source code
    let modified_code = format!(
        "{}{}{}\n\n{}\n\n{}{}{}",
        before.trim_end(),
        "\n",
        marker_start,
        combined_blocks.join("\n\n"),
        marker_end,
        "\n",
        after.trim_start()
    );

    Ok(modified_code)
}

fn _save_new_executable<P: AsRef<Path>>(
    modified_source: &str,
    cargo_toml_content: &str,
    temp_dir_path: P,
    project_name: &str,
) -> io::Result<()> {
    let temp_dir = temp_dir_path.as_ref();
    let src_dir = temp_dir.join("src");

    println!("[LOG] Creating directory: {:?}", src_dir);
    fs::create_dir_all(&src_dir)?;

    let source_path = src_dir.join("main.rs");
    println!("[LOG] Writing source to: {:?}", source_path);
    fs::write(&source_path, modified_source)?;

    let cargo_path = temp_dir.join("Cargo.toml");
    println!("[LOG] Writing Cargo.toml to: {:?}", cargo_path);
    fs::write(&cargo_path, cargo_toml_content)?;

    println!("[LOG] Checking for Cargo...");
    let cargo_check = Command::new("cargo").arg("--version").output();
    if cargo_check.is_err() {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            "Cargo not found. Ensure Rust is installed and added to PATH.",
        ));
    }

    println!("[LOG] Running cargo build in: {:?}", temp_dir);
    let build_output = Command::new("cargo")
        .arg("build")
        .current_dir(&temp_dir)
        .output()?;

    if !build_output.status.success() {
        eprintln!("[ERROR] Build failed!");
        io::stderr().write_all(&build_output.stderr)?;
        return Err(io::Error::new(io::ErrorKind::Other, "Build failed"));
    }

    let exe_name = if cfg!(target_os = "windows") {
        format!("{}.exe", project_name)
    } else {
        project_name.to_string()
    };
    let target_dir = temp_dir.join("target/debug");
    let new_exe = target_dir.join(&exe_name);

    println!("[LOG] Checking for new executable: {:?}", new_exe);
    if !new_exe.exists() {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            format!("New executable not found at {:?}", new_exe),
        ));
    }

    let exe_path = env::current_exe()?;
    println!("[LOG] Current executable path: {:?}", exe_path);

    // Attempt to replace the current executable
    let temp_exe_path = exe_path.with_extension("tmp");
    println!(
        "[LOG] Attempting to rename current executable to: {:?}",
        temp_exe_path
    );

    let mut attempts = 0;
    const MAX_ATTEMPTS: u32 = 10;
    const RETRY_DELAY_MS: u64 = 500;

    // Step 1: Rename current executable
    let rename_success = loop {
        match fs::rename(&exe_path, &temp_exe_path) {
            Ok(_) => {
                println!(
                    "[LOG] Successfully renamed current executable to {:?}",
                    temp_exe_path
                );
                break true;
            }
            Err(e) if attempts < MAX_ATTEMPTS => {
                attempts += 1;
                println!(
                    "[LOG] Rename attempt {}/{} failed: {}. Retrying after {}ms...",
                    attempts, MAX_ATTEMPTS, e, RETRY_DELAY_MS
                );
                thread::sleep(Duration::from_millis(RETRY_DELAY_MS));
                continue;
            }
            Err(e) => {
                println!(
                    "[ERROR] Failed to rename current executable after {} attempts: {}. This may require running as administrator.",
                    MAX_ATTEMPTS, e
                );
                break false;
            }
        }
    };

    // Step 2: Copy new executable
    attempts = 0;
    println!("[LOG] Attempting to copy new executable to: {:?}", exe_path);
    let copy_success = loop {
        match fs::copy(&new_exe, &exe_path) {
            Ok(_) => {
                println!("[LOG] Successfully copied new executable to {:?}", exe_path);
                break true;
            }
            Err(e) if attempts < MAX_ATTEMPTS => {
                attempts += 1;
                println!(
                    "[LOG] Copy attempt {}/{} failed: {}. Retrying after {}ms...",
                    attempts, MAX_ATTEMPTS, e, RETRY_DELAY_MS
                );
                thread::sleep(Duration::from_millis(RETRY_DELAY_MS));
                continue;
            }
            Err(e) => {
                println!(
                    "[ERROR] Failed to copy executable after {} attempts: {}.",
                    MAX_ATTEMPTS, e
                );
                break false;
            }
        }
    };

    // Step 3: Cleanup
    if rename_success && copy_success {
        let _ = fs::remove_file(&temp_exe_path);
        println!("[LOG] Cleaned up temporary executable: {:?}", temp_exe_path);
    } else if rename_success {
        // If copy failed, restore the original executable
        let _ = fs::rename(&temp_exe_path, &exe_path);
        return Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            "Failed to replace executable. Try running as administrator.",
        ));
    } else {
        return Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            "Failed to rename current executable. Try running as administrator.",
        ));
    }

    Ok(())
}

pub fn morph(
    source_code: &str,
    cargo_toml_content: &str,
    junk_function_count: usize,
    project_name: &str,
) -> io::Result<()> {
    println!("Program running!");
    thread::sleep(Duration::from_secs(2));

    println!(
        "[LOG] Using provided source code (length: {})",
        source_code.len()
    );
    let source_content = source_code.to_string();

    // Extract original_content from source_content between markers
    let marker_start = "// METAMORPHIC_MARKER_START";
    let marker_end = "// METAMORPHIC_MARKER_END";
    let start_idx = source_content.find(marker_start);
    let end_idx = source_content.find(marker_end);

    let original_content = if let (Some(start), Some(end)) = (start_idx, end_idx) {
        let content = &source_content[start + marker_start.len()..end];
        // Remove the batnet::metamorphic call to prevent recursion
        let re = Regex::new(r"batnet::metamorphic\([^)]*\);").unwrap();
        re.replace_all(content, "").trim().to_string()
    } else {
        return Err(io::Error::new(
            io::ErrorKind::Other,
            "Markers not found in source content",
        ));
    };

    println!("[LOG] Extracted original content:\n{}", original_content);

    println!("[LOG] Transforming source code...");
    let modified_source =
        _transform_source(&source_content, &original_content, junk_function_count)?;
    println!(
        "[LOG] Modified source code length: {}",
        modified_source.len()
    );
    println!("[LOG] Modified source code:\n{}", modified_source);

    println!("[LOG] Saving new executable...");
    let temp_dir = env::temp_dir().join(format!("m{}", _generate_random_ident("")));
    println!("[LOG] Temporary directory: {:?}", temp_dir);

    // Ensure the temporary directory is removed after use
    let result = _save_new_executable(
        &modified_source,
        cargo_toml_content,
        &temp_dir,
        project_name,
    );

    // Cleanup temporary directory
    let _ = fs::remove_dir_all(&temp_dir);
    println!("[LOG] Cleaned up temporary directory: {:?}", temp_dir);

    result?;
    println!("Program finished!");
    Ok(())
}