#!/usr/bin/env python3
import os
import subprocess
import sys

def run_all_tests(lean4export_path: str, tests_root: str, lean_proj_dir: str, lean_mod_filename: str):
    if not os.path.isdir(tests_root):
        raise FileNotFoundError(f"Tests root directory not found: {tests_root}")
    if not os.path.isdir(lean_proj_dir):
        raise FileNotFoundError(f"Lean project directory not found: {lean_proj_dir}")
    if not os.path.isfile(lean4export_path) or not os.access(lean4export_path, os.X_OK):
        raise FileNotFoundError(f"lean4export not found or not executable: {lean4export_path}")

    os.chdir(lean_proj_dir)

    lean_mod_filepath = os.path.join(lean_proj_dir, lean_mod_filename)

    original_contents = ""
    if os.path.isfile(lean_mod_filepath):
        with open(lean_mod_filepath, "r") as f:
            original_contents = f.read()

    # Iterate over each entry in the tests_root directory
    for entry in os.listdir(tests_root):
        test_dir = os.path.join(tests_root, entry)
        if not os.path.isdir(test_dir):
            continue  # Skip non-directories

        source_path = os.path.join(test_dir, "source")
        export_path = os.path.join(test_dir, "export")

        # Skip this test directory if source file is missing
        if not os.path.isfile(source_path):
            print(f"[WARNING] 'source' not found in {test_dir}, skipping.")
            continue

        # Read the contents of the test's "source" file
        with open(source_path, "r") as f_src:
            source_contents = f_src.read()

        # Overwrite the selected .lean file with the tests's `source`
        with open(lean_mod_filepath, "w") as f_dot_lean:
            f_dot_lean.write(source_contents)

        # Extract module name from lean_mod_filename (remove .lean extension)
        module_name = os.path.splitext(lean_mod_filename)[0]
        exec_args = ["lake", "env", lean4export_path, module_name]

        if "unsafe" in source_contents:
            exec_args.append("--export-unsafe")
            continue
        try:
            subprocess.run(
                ["lake", "build", module_name],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True
            )
        except subprocess.CalledProcessError as e:
            print(f"[ERROR] Build failed for test '{entry}':")
            print(f"        Exit code: {e.returncode}")
            print(f"        stderr: {e.stderr.decode().strip()}")
            # Skip running the executable and continue to the next test
            continue

        with open(export_path, "w") as f_out:
            try:
                subprocess.run(
                    exec_args,

                    stdout=f_out,
                    stderr=subprocess.PIPE,
                    check=True
                )
            except subprocess.CalledProcessError as e:
                print(f"[ERROR] Test '{entry}' failed when running executable:")
                print(f"        Exit code: {e.returncode}")
                print(f"        stderr: {e.stderr.decode().strip()}")
                # Continue to next test even if this one fails
                continue

        print(f"[OK] Ran test '{entry}': output written to {export_path}")

    # Restore the original contents of the lean module file
    with open(lean_mod_filepath, "w") as f:
        f.write(original_contents)


if __name__ == "__main__":
    if len(sys.argv) != 5:
        print("Error: Missing required arguments", file=sys.stderr)
        print(f"Usage: {sys.argv[0]} <lean4export_path> <tests_root> <lean_proj_dir> <lean_mod_filename>", file=sys.stderr)
        print("", file=sys.stderr)
        print("Arguments:", file=sys.stderr)
        print("  lean4export_path  - Path to the lean4export executable", file=sys.stderr)
        print("  tests_root        - Path to the directory containing test subdirectories (each should contain a `source` file)", file=sys.stderr)
        print("  lean_proj_dir     - Path to the scrap Lean project directory", file=sys.stderr)
        print("  lean_mod_filename - Name of the file inside lean_proj_dir to write test cases to", file=sys.stderr)
        sys.exit(1)
    
    lean4export_path = sys.argv[1]
    tests_root = sys.argv[2]
    lean_proj_dir = sys.argv[3]
    lean_mod_filename = sys.argv[4]
    
    try:
        run_all_tests(lean4export_path, tests_root, lean_proj_dir, lean_mod_filename)
    except FileNotFoundError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)