rem-verification 0.1.0

use std::{
    path::PathBuf,
    fs::File,
    io::Write as ioWrite,
    process::Command,
    io::BufRead,
    io::BufReader,
    collections::HashSet,
};

use regex::Regex;

use log::{
    error, 
    info
};

/// Returns the paths to:
/// - The _CoQProject file
/// - The EquivCheck.v file
/// - The Primitives.v file
/// These are used to ensure the files are deleted at the end of the
/// verification process.
pub fn coq_verification(
    original_coq: &PathBuf,
    refactored_coq: &PathBuf,
    top_level_function: &String,
) -> Result<(PathBuf, PathBuf, PathBuf, bool), Box<dyn std::error::Error>> {

    let dir: PathBuf = original_coq.parent().unwrap().to_path_buf(); // The directory where the CoQ files are saved

    // First we need to create the _CoqProject file
    let coq_path: PathBuf = create_coq_project_file(
        &original_coq,
        &refactored_coq,
        &dir,
    )?;

    info!("Created _CoqProject file");

    // Create the EquivCheck.v file
    let equiv_check: PathBuf = create_equiv_check_file(
        &original_coq,
        &refactored_coq,
        &top_level_function,
        &dir
    )?;

    info!("Created EquivCheck.v file");

    // Now we need to call CoQ to verify the files
    let (result, primitives) = verify_coq_files(
        &original_coq,
        &refactored_coq,
        &equiv_check,
        &dir,
    )?;

    if result {
        info!("CoQ verification completed successfully");
    } else {
        error!("CoQ verification failed");
    }

    // Finally cleanup the directory (remove all files other than the CoQ files
    // and the _CoQ project file)
    cleanup_directory(&dir)?;

    Ok((
        coq_path,
        equiv_check,
        primitives,
        result
    ))
}

/// Creates the _CoqProject file for the CoQ project. This file is used by coqc
/// to understand the project structure and dependencies, and helps with
/// compilation.
fn create_coq_project_file(
    original_coq: &PathBuf,
    refactored_coq: &PathBuf,
    dir: &PathBuf,
) -> Result<PathBuf, Box<dyn std::error::Error>> {
    // The _CoqProject is a text file that contains:
    // -R . AutoEquiv
    // Primitives.v
    // <Original CoQ file>
    // <Refactored CoQ file>
    // Any other CoQ files that are needed for the verification
    // EquivCheck.v

    // Create the file
    let mut coq_project = File::create(dir.join("_CoqProject"))?;
    writeln!(coq_project, "-R . AutoEquiv")?;
    writeln!(coq_project, "")?;
    writeln!(coq_project, "Primitives.v")?;
    writeln!(coq_project, "{}", original_coq.file_name().unwrap().to_str().unwrap())?;
    writeln!(coq_project, "{}", refactored_coq.file_name().unwrap().to_str().unwrap())?;
    writeln!(coq_project, "EquivCheck.v")?;

    // Return the path to the _CoqProject file
    let coq_path = dir.join("_CoqProject");
    Ok(coq_path)
}

/// Creates the EquivCheck.v file for the CoQ project. This is the most involved
/// part of the verification process, as it requires writing a CoQ script that
/// can check the equivalence of two functions.
/// The script will be saved in the same directory as the original CoQ file.
fn create_equiv_check_file(
    original_coq: &PathBuf,
    refactored_coq: &PathBuf,
    top_level_function: &String,
    dir: &PathBuf,
) -> Result<PathBuf, Box<dyn std::error::Error>> {
    // The EquivCheck.v file will contain the following:
    // Require Import Primitives.
    // Import Primitives.
    // Require Import <Original CoQ file>.
    // Require Import <Refactored CoQ file>.
    // Require Import Coq.ZArith.ZArith.
    // Local Open Scope Primitives_scope.

    // Theorem equiv_check: forall (args: list nat), <Top level function> args = <Top level function> args.
    // Proof.
    //   reflexivity.
    // Qed.

    // Create the file
    info!("Directory: {:?}", dir);

    let equiv_check_path = dir.join("EquivCheck.v");
    let mut equiv_check = File::create(&equiv_check_path)?;

    // Write the header comments
    writeln!(equiv_check, "(** EquivCheck.v Autogenerated by REM-verification **)")?;

    // Always include the Primitives.v file
    // These files are removed from the list of imports we collect to ensure
    // they aren't duplicated.
    writeln!(equiv_check, "Require Import Primitives.")?;
    writeln!(equiv_check, "Import Primitives.")?;

    // Also include the original and refactored CoQ files so that we have access
    // to the functions we need to check.
    writeln!(
        equiv_check,
        "Require Import {}.",
        original_coq.file_stem().unwrap().to_str().unwrap()
    )?;
    writeln!(
        equiv_check,
        "Require Import {}.",
        refactored_coq.file_stem().unwrap().to_str().unwrap()
    )?;

    // Extract additional import lines from both files.
    let original_imports: Vec<String> = extract_imports(original_coq)?;
    let refactored_imports: Vec<String> = extract_imports(refactored_coq)?;
    let combined_imports: String = combine_imports(original_imports, refactored_imports);

    if !combined_imports.is_empty() {
        writeln!(
            equiv_check,
            "{}",
            combined_imports
        )?;
    }

    writeln!(equiv_check, "Local Open Scope Primitives_scope.")?;
    writeln!(equiv_check, "\n")?;

    // From here we need to evaluate the Lemma / Theorem needed to check
    // equivalence. This will depend on the top level function provided.
    // TODO add in more than just a forall check
    let proof: String = generate_equiv_check_proof(&top_level_function, &original_coq, &refactored_coq)?;
    writeln!(equiv_check, "{}", proof)?;
    info!("Proof generated successfully");

    Ok(dir.join("EquivCheck.v"))
}

/// Extracts import lines from the beginning of a Coq file.
/// It collects lines starting with "Require", "Import", or "Require Export"
/// until it reaches a line that starts with "Local Open Scope" or a non-import line.
fn extract_imports(coq_path: &PathBuf) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let file: File = File::open(coq_path)?;
    let reader = BufReader::new(file);
    let mut imports: Vec<String> = Vec::new();

    for line in reader.lines() {
        let line = line?;
        let trimmed = line.trim();
        // Stop if we encounter the "Local Open Scope" directive or a non-import statement.
        if trimmed.starts_with("Local Open Scope") {
            break;
        }
        // Stop if we encounter the Module declaration.
        if trimmed.starts_with("Module") {
            break;
        }
        // If the line starts with "Require", "Import", or "Require Export", consider it an import.
        if trimmed.starts_with("Require") || trimmed.starts_with("Import") {
            imports.push(trimmed.to_string());
        } else if !trimmed.is_empty() {
            // If we hit a non-empty, non-import line, we assume the import section is done.
            break;
        }
    }
    Ok(imports)
}

/// Combines two lists of import lines into one, deduplicating identical lines.
/// Also removes the primitives imports from the output list.
fn combine_imports(imports1: Vec<String>, imports2: Vec<String>) -> String {
    let mut set = HashSet::new();
    // Insert each line into the set (this automatically deduplicates).
    for line in imports1.into_iter().chain(imports2.into_iter()) {
        set.insert(line);
    }
    // Convert the set into a vector.
    let mut combined: Vec<_> = set.into_iter().collect::<Vec<String>>();
    // Remove the Primitives imports.
    combined.retain(|line| !line.contains("Primitives"));
    // Combine the set into a single string.
    combined.join("\n")
}


struct CoqArgument {
    name: String,
    ty: String,
}

/// Generates the proof for the equivalence check.
fn generate_equiv_check_proof(
    top_level_function: &str,
    original_coq: &PathBuf,
    refactored_coq: &PathBuf,
) -> Result<String, Box<dyn std::error::Error>> {
    // Read the original Coq file as a string.
    let content = std::fs::read_to_string(original_coq)?;


    // Use a regex to locate the definition of the function.
    // We assume the definition looks like:
    //   Definition top_level_function (arg1 : type1) (arg2 : type2) ... : ...
    let def_pattern = format!(r"Definition\s+{}\s+((?:\([^)]*\)\s*)+):", top_level_function);
    let def_re = Regex::new(&def_pattern)?;

    // Find the arguments portion.
    let caps = def_re.captures(&content)
        .ok_or("Could not find the function definition in the original Coq file")?;
    let args_str = caps.get(1).unwrap().as_str();

    // Now extract each argument of the form: (name : type)
    let arg_re = Regex::new(r"\(\s*([^:\s]+)\s*:\s*([^)]+)\s*\)")?;
    let mut args: Vec<CoqArgument> = Vec::new();
    for cap in arg_re.captures_iter(args_str) {
        let name = cap.get(1).unwrap().as_str().trim().to_string();
        let ty = cap.get(2).unwrap().as_str().trim().to_string();
        args.push(CoqArgument { name, ty });
    }

    // Build the "forall" part of the lemma and the argument list for function application.
    let forall_part = args.iter()
        .map(|arg| format!("({} : {})", arg.name, arg.ty))
        .collect::<Vec<_>>()
        .join(" ");
    let args_names = args.iter()
        .map(|arg| arg.name.clone())
        .collect::<Vec<_>>()
        .join(" ");

    // Name the Lemma
    let lemma_name = format!("{}_equiv_check", top_level_function);

    // Get the module names for the original and refactored CoQ files.
    let original_module = original_coq
        .file_stem()
        .unwrap()
        .to_str()
        .unwrap();
    let refactored_module = refactored_coq
        .file_stem()
        .unwrap()
        .to_str()
        .unwrap();

    // Generate the proof script as a string.
    let proof = format!(
r"Lemma {} : forall {},
  {}.{} {} = {}.{} {}.
Proof.
  intros {}.
  unfold {}.{}.
  unfold {}.{}.
  reflexivity.
Qed.",
        lemma_name, forall_part, // Line 1
        original_module, top_level_function, args_names, // Line 2
        refactored_module, top_level_function, args_names, // Line 2
        args_names, // Line 4
        original_module, top_level_function, // Line 5
        refactored_module, top_level_function, // Line 6
    );

    // Log the generated proof script
    info!("Proof script:\n{}", proof);

    Ok(proof)
}


/// Verifies the CoQ files using CoQ's coqc command.
/// This function will call coqc on the original and refactored CoQ files, as well
/// as the EquivCheck.v file. If any of the files fail to compile, the function
/// will return false.
/// The function will return true if all files compile successfully, and it has
/// checked the return value of the EquivCheck.v file.
/// Verifies the CoQ files using CoQ's coqc command.
/// This function will call coqc on the original and refactored CoQ files,
/// as well as the EquivCheck.v file. If any of the files fail to compile,
/// the function will return false. The function will return true if all
/// files compile successfully.
/// Returns the result of the equivalence check and the path to the Primitives.v file.
fn verify_coq_files(
    original_coq: &PathBuf,
    refactored_coq: &PathBuf,
    equiv_check: &PathBuf,
    dir: &PathBuf, // This should be the directory that contains Primitives.v.
) -> Result<(bool, PathBuf), Box<dyn std::error::Error>> {
    // The mapping: -R <dir> Primitives
    let dir_str = dir.to_str().ok_or("Invalid directory path")?;

    info!("Directory: {}", dir_str);

    // Start with the Primitives.v file
    let primitives: PathBuf = dir.join("Primitives.v");
    let primitives_coq_output = Command::new("coqc")
        .arg("-R")
        .arg(dir_str)
        .arg("Primitives")
        .arg(&primitives)
        .output()?;
    if !primitives_coq_output.status.success() {
        let stdout = String::from_utf8_lossy(&primitives_coq_output.stdout);
        let stderr = String::from_utf8_lossy(&primitives_coq_output.stderr);
        error!("Primitives.v failed to compile:\nstdout: {}\nstderr: {}", stdout, stderr);
        return Ok((false, primitives));
    }
    info!("Primitives.v compiled successfully");

    // Compile the original Coq file with the same mapping
    let original_coq_output = Command::new("coqc")
        .arg("-R")
        .arg(dir_str)
        .arg("Primitives")
        .arg(&original_coq)
        .output()?;
    if !original_coq_output.status.success() {
        let stdout = String::from_utf8_lossy(&original_coq_output.stdout);
        let stderr = String::from_utf8_lossy(&original_coq_output.stderr);
        error!("Original Coq file failed to compile:\nstdout: {}\nstderr: {}", stdout, stderr);
        return Ok((false, primitives));
    }
    info!("Original Coq file compiled successfully");

    // Compile the refactored Coq file
    let refactored_coq_output = Command::new("coqc")
        .arg("-R")
        .arg(dir_str)
        .arg("Primitives")
        .arg(&refactored_coq)
        .output()?;
    if !refactored_coq_output.status.success() {
        let stdout = String::from_utf8_lossy(&refactored_coq_output.stdout);
        let stderr = String::from_utf8_lossy(&refactored_coq_output.stderr);
        error!("Refactored Coq file failed to compile:\nstdout: {}\nstderr: {}", stdout, stderr);
        return Ok((false, primitives));
    }
    info!("Refactored Coq file compiled successfully");

    // Compile the EquivCheck.v file
    let equiv_check_output = Command::new("coqc")
        .arg("-R")
        .arg(dir_str)
        .arg("Primitives")
        .arg(&equiv_check)
        .output()?;
    if !equiv_check_output.status.success() {
        let stdout = String::from_utf8_lossy(&equiv_check_output.stdout);
        let stderr = String::from_utf8_lossy(&equiv_check_output.stderr);
        error!("EquivCheck.v failed to compile:\nstdout: {}\nstderr: {}", stdout, stderr);
        return Ok((false, primitives));
    }
    info!("EquivCheck.v compiled successfully");

    Ok((true, primitives))
}

/// Cleans up the directory by removing all files other than the CoQ files and
/// the _CoqProject file.
/// eqivalent to rm -f *.glob *.vo *.vok *.vos *.aux .*.aux .*.cache, but
/// platform independent.
fn cleanup_directory(dir: &PathBuf) -> Result<(), Box<dyn std::error::Error>> {
    let files = std::fs::read_dir(dir)?;
    for file in files {
        let file = file?;
        let path = file.path();
        let extension = path.extension();
        if let Some(extension) = extension {
            if extension == "glob"
            || extension == "vo"
            || extension == "vok"
            || extension == "vos"
            || extension == "aux"
            || extension == "cache"
            {
                std::fs::remove_file(path)?;
            }
        }
    }

    Ok(())
}