#!/usr/bin/env python3
import os
import sys
import subprocess
import tempfile
from pathlib import Path

def run_test(test_file: Path, compiler_bin: Path):
    with open(test_file, "r", encoding="utf-8") as f:
        lines = f.readlines()

    expected_code = 0
    code_lines = []
    harness_lines = None
    current_section = None

    for line in lines:
        trimmed = line.strip()
        if trimmed.startswith("[") and trimmed.endswith("]"):
            current_section = trimmed[1:-1]
            if current_section == "harness":
                harness_lines = []
        elif current_section:
            if current_section == "expected_return_code":
                if trimmed:
                    expected_code = int(trimmed)
            elif current_section == "code":
                code_lines.append(line)
            elif current_section == "harness":
                harness_lines.append(line)

    code = "".join(code_lines)
    harness = "".join(harness_lines) if harness_lines is not None else None

    with tempfile.TemporaryDirectory(prefix=f"compiler_test_{test_file.stem}_") as temp_dir_name:
        temp_dir = Path(temp_dir_name)
        
        code_path = temp_dir / "code.src"
        obj_path = temp_dir / "code.o"
        exec_path = temp_dir / "exec"

        code_path.write_text(code, encoding="utf-8")

        # 1. Invoke the custom compiler
        comp_res = subprocess.run(
            [str(compiler_bin), str(code_path), "-o", str(obj_path)],
            capture_output=True,
            text=True
        )
        
        if comp_res.returncode != 0:
            out = comp_res.stdout.strip()
            err = comp_res.stderr.strip()
            error_msg = f"{out}\n{err}".strip() or "Compiler failed with no output"
            raise Exception(error_msg)

        # 2. Invoke GCC to link the object file (and compile harness if provided)
        gcc_cmd = ["gcc", str(obj_path), "-o", str(exec_path)]
        if harness is not None:
            harness_path = temp_dir / "harness.c"
            harness_path.write_text(harness, encoding="utf-8")
            gcc_cmd.append(str(harness_path))

        gcc_res = subprocess.run(gcc_cmd, capture_output=True, text=True)
        
        if gcc_res.returncode != 0:
            out = gcc_res.stdout.strip()
            err = gcc_res.stderr.strip()
            error_msg = f"{out}\n{err}".strip() or "GCC failed with no output"
            raise Exception(error_msg)

        # 3. Run the resulting executable
        exec_res = subprocess.run([str(exec_path)], capture_output=True, text=True)
        
        if exec_res.returncode != expected_code:
            raise Exception(f"Expected return code {expected_code}, but got {exec_res.returncode}")

def main():
    GREEN = '\033[92m'
    RED = '\033[91m'
    RESET = '\033[0m'

    project_root = Path(__file__).resolve().parent
    tests_dir = project_root / "tests"
    compiler_bin = project_root / "target" / "debug" / "compiler"

    # Automatically run cargo build to ensure the compiler is compiled and up to date
    build_status = subprocess.run(["cargo", "build"], cwd=project_root)
    if build_status.returncode != 0:
        print("Failed to build the compiler binary.", file=sys.stderr)
        sys.exit(1)

    test_files = sorted(tests_dir.glob("*.test"))
    
    all_passed = True
    failed_tests = []

    for test_file in test_files:
        test_name = test_file.name
        
        # Format the output to align results, padding with dots
        msg = f"Running {test_name} "
        print(f"{msg:.<50} ", end="", file=sys.stderr, flush=True)
        
        try:
            run_test(test_file, compiler_bin)
            print(f"{GREEN}ok{RESET}", file=sys.stderr)
        except Exception as e:
            print(f"{RED}FAIL{RESET}", file=sys.stderr)
            print(f"  Reason: {e}", file=sys.stderr)
            all_passed = False
            failed_tests.append(test_name)
            
    if not all_passed:
        print(f"\nOne or more E2E tests failed: {failed_tests}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()