run_benchmarks.py

#!/usr/bin/env python3
"""
Root-level benchmark runner for MCode.
Runs benchmarks from the dataset/ directory and saves results to results/results.jsonl.
"""

import yaml
import subprocess
import time
import sys
import shutil
import argparse
import json
import re
import os
import tomllib
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional, Dict, List

from generate_prompts import generate_prompt, load_benchmark_config as load_benchmark_config_for_prompt


@dataclass
class ImplConfig:
    """Configuration for a source or destination implementation."""
    language: str
    description: str
    install_cmd: Optional[str] = None
    build_cmd: Optional[str] = None
    run_cmd: Optional[str] = None
    port: Optional[int] = None
    port_env_var: Optional[str] = None
    startup_wait: Optional[int] = None
    env: Dict[str, str] = field(default_factory=dict)


@dataclass
class BenchmarkConfig:
    """Top-level benchmark configuration."""
    id: str
    name: str
    benchmark_type: str
    source: ImplConfig
    destination: ImplConfig
    test_dir: str
    benchmark_dir: Path
    pytest_args: Optional[str] = None
    source_repo: Optional[str] = None
    source_commit: Optional[str] = None


@dataclass
class TestResult:
    """Result of testing an implementation."""
    language: str
    build_success: bool
    build_time_seconds: float
    tests_passed: int
    tests_failed: int
    tests_error: int
    tests_skipped: int
    tests_total: int
    pass_rate: float
    test_time_seconds: float
    failed_tests: List[str]


def load_benchmark_config(benchmark_dir: Path) -> Optional[BenchmarkConfig]:
    """Load benchmark.yml and metadata.yml from a benchmark directory."""
    benchmark_file = benchmark_dir / "benchmark.yml"
    metadata_file = benchmark_dir / "metadata.yml"

    if not benchmark_file.exists():
        return None

    with open(benchmark_file) as f:
        config = yaml.safe_load(f)

    # Skip benchmarks missing required fields
    if 'type' not in config.get('benchmark', {}):
        return None

    # Load ID and provenance from metadata.yml
    benchmark_id = "000"
    source_repo = None
    source_commit = None
    if metadata_file.exists():
        with open(metadata_file) as f:
            metadata = yaml.safe_load(f)
            benchmark_id = metadata.get("id", "000")
            provenance = metadata.get("provenance", {})
            source_repo = provenance.get("source_repo")
            source_commit = provenance.get("source_commit")

    def extract_impl_config(impl_dict) -> ImplConfig:
        return ImplConfig(
            language=impl_dict['language'],
            description=impl_dict.get('description', ''),
            install_cmd=impl_dict.get('install_cmd'),
            build_cmd=impl_dict.get('build_cmd'),
            run_cmd=impl_dict.get('run_cmd'),
            port=impl_dict.get('port'),
            port_env_var=impl_dict.get('port_env_var'),
            startup_wait=impl_dict.get('startup_wait'),
            env=impl_dict.get('env') or {},
        )

    return BenchmarkConfig(
        id=benchmark_id,
        name=config['benchmark']['name'],
        benchmark_type=config['benchmark']['type'],
        source=extract_impl_config(config['source']),
        destination=extract_impl_config(config['destination']),
        test_dir=config['testing']['test_dir'],
        benchmark_dir=benchmark_dir,
        pytest_args=config['testing'].get('pytest_args'),
        source_repo=source_repo,
        source_commit=source_commit,
    )


def discover_benchmarks(dataset_dir: Path) -> Dict[str, BenchmarkConfig]:
    """Discover all benchmarks in the dataset directory."""
    benchmarks = {}
    for item in dataset_dir.iterdir():
        if item.is_dir() and (item / "benchmark.yml").exists():
            config = load_benchmark_config(item)
            if config:
                benchmarks[config.id] = config
    return benchmarks


def run_cmd(cmd: list, **kwargs) -> subprocess.CompletedProcess:
    """Run a command and return the result."""
    return subprocess.run(cmd, **kwargs)


class BenchmarkRunner:
    """Runs a single benchmark."""

    def __init__(self, config: BenchmarkConfig, log_file: Path, workspace_dir: Path = None):
        self.config = config
        self.log_file = log_file
        self.temp_workspace = None
        self.temp_compose_file = None
        # Use absolute path to avoid path issues with cwd in subprocess calls
        self.benchmark_dir = config.benchmark_dir.absolute()
        # If workspace_dir provided, use it directly (already isolated)
        self.provided_workspace = workspace_dir
        self.owns_workspace = workspace_dir is None  # Only cleanup workspace if we created it

    def initialize_workspace(self) -> bool:
        """Initialize empty workspace - will clone inside container."""
        print(f"  Initializing workspace...")

        self.temp_workspace.mkdir(parents=True, exist_ok=True)

        # Create empty src and dst directories
        (self.temp_workspace / "src").mkdir(exist_ok=True)
        (self.temp_workspace / "dst").mkdir(exist_ok=True)

        # Copy existing workspace/dst if it exists (for evaluating pre-generated dst)
        workspace_dst = self.benchmark_dir / "workspace" / "dst"
        if workspace_dst.exists() and any(workspace_dst.iterdir()):
            print(f"  Copying existing workspace/dst...")
            shutil.copytree(workspace_dst, self.temp_workspace / "dst", dirs_exist_ok=True)

        # Generate prompt.md from benchmark.yml (do this before container starts)
        benchmark_config = load_benchmark_config_for_prompt(self.benchmark_dir)
        if benchmark_config:
            prompt_content = generate_prompt(benchmark_config, "v1")
            with open(self.temp_workspace / "prompt.md", 'w') as f:
                f.write(prompt_content)

        return True

    def clone_and_setup_in_container(self) -> bool:
        """Clone repository and setup git inside the container."""
        print(f"  Setting up source repository in container...")

        if self.config.source_repo and self.config.source_commit:
            print(f"    Cloning {self.config.source_repo}...")

            # Clone inside container
            clone_cmd = f"cd /workspace && rm -rf src && git clone --recurse-submodules {self.config.source_repo} src"
            exit_code, stdout, stderr = self.docker_exec(clone_cmd, capture_output=True)

            if exit_code != 0:
                print(f"    Clone failed: {stderr}")
                return False

            # Checkout specific commit
            checkout_cmd = f"cd /workspace/src && git checkout {self.config.source_commit}"
            exit_code, stdout, stderr = self.docker_exec(checkout_cmd, capture_output=True)

            if exit_code != 0:
                print(f"    Checkout failed: {stderr}")
                return False

            # Update submodules
            submodule_cmd = "cd /workspace/src && git submodule update --init --recursive"
            self.docker_exec(submodule_cmd, capture_output=True)

            print(f"    Cloned at {self.config.source_commit[:8]}")
        else:
            # Copy from workspace/src if it exists (for benchmarks without provenance)
            workspace_src = self.benchmark_dir / "workspace" / "src"
            if workspace_src.exists():
                print(f"    Copying workspace/src (no provenance)...")
                # Copy files into mounted workspace
                shutil.copytree(workspace_src, self.temp_workspace / "src", dirs_exist_ok=True)
            else:
                print(f"    No source_repo in metadata and no workspace/src")
                return False

        # Initialize dst/ with git repo inside container
        print(f"    Initializing destination directory...")
        init_dst_cmd = """
            cd /workspace/dst && \
            git init && \
            git config user.email "agent@benchmark.local" && \
            git config user.name "Benchmark Agent"
        """
        exit_code, stdout, stderr = self.docker_exec(init_dst_cmd, capture_output=True)

        if exit_code != 0:
            print(f"    Warning: Failed to initialize dst/: {stderr}")

        return True

    def docker_exec(self, cmd: str, env: Dict[str, str] = None, detached: bool = False, capture_output: bool = False) -> tuple:
        """Execute a command inside the Docker container."""
        docker_cmd = ["docker", "compose", "-f", str(self.temp_compose_file), "exec"]
        if detached:
            docker_cmd.append("-d")

        if env:
            for key, value in env.items():
                docker_cmd.extend(["-e", f"{key}={value}"])

        docker_cmd.extend(["dev", "bash", "-c", cmd])

        if capture_output:
            result = run_cmd(docker_cmd, cwd=self.benchmark_dir, capture_output=True, text=True)
            return result.returncode, result.stdout, result.stderr
        else:
            result = run_cmd(docker_cmd, cwd=self.benchmark_dir)
            return result.returncode, "", ""

    def setup(self):
        """Set up the test environment."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        self.temp_compose_file = self.benchmark_dir / f"docker-compose-{timestamp}.yml"
        self.container_name = f"{self.benchmark_dir.name}-dev-{timestamp}"

        # Use provided workspace or initialize from scratch
        if self.provided_workspace:
            print(f"  Using provided workspace...")
            self.temp_workspace = self.provided_workspace
        else:
            self.temp_workspace = self.benchmark_dir / f"workspace-{timestamp}"
            if not self.initialize_workspace():
                return False

        # Copy tests into workspace
        shutil.copytree(self.benchmark_dir / self.config.test_dir, self.temp_workspace / "tests", dirs_exist_ok=True)

        with open(self.benchmark_dir / "docker-compose.dev.yml", "r") as f:
            compose_content = f.read()

        # Replace workspace mount path
        if self.provided_workspace:
            # Use absolute path for provided workspace (may be in different directory)
            compose_content = re.sub(
                r'\./workspace:/workspace',
                f'{self.temp_workspace.absolute()}:/workspace',
                compose_content
            )
        else:
            # Use relative path for temp workspace (in same directory)
            compose_content = compose_content.replace("./workspace:", f"./{self.temp_workspace.name}:")

        # Remove port mapping to avoid conflicts when running in parallel
        # Tests run inside the container, so host port mapping is not needed
        compose_content = re.sub(
            r'    ports:\n      - "\d+:\d+"\n',
            "",
            compose_content
        )

        # Make container name unique to avoid conflicts
        compose_content = re.sub(
            r'container_name:\s*\S+',
            f'container_name: {self.container_name}',
            compose_content
        )

        # Replace /bin/bash with a command that keeps container running
        # This is needed because bash exits immediately when run detached
        compose_content = compose_content.replace(
            "command: /bin/bash",
            "command: tail -f /dev/null"
        )

        with open(self.temp_compose_file, "w") as f:
            f.write(compose_content)

        # Clean up any existing containers with the same name first
        run_cmd(
            ["docker", "compose", "-f", str(self.temp_compose_file), "down", "--remove-orphans"],
            cwd=self.benchmark_dir,
            capture_output=True,
            text=True
        )

        print(f"  Building and starting Docker container...")
        result = run_cmd(["docker", "compose", "-f", str(self.temp_compose_file), "up", "--build", "-d"],
                cwd=self.benchmark_dir, capture_output=True, text=True)

        if result.returncode != 0:
            error_msg = result.stderr
            if "address already in use" in error_msg.lower() or "bind" in error_msg.lower():
                print(f"  ❌ Failed to start container: Port 3000 is already in use")
                print(f"     Please free up port 3000 or stop any services using it")
                print(f"     Error details: {error_msg.split('Error')[0] if 'Error' in error_msg else error_msg[-200:]}")
            else:
                print(f"  ❌ Failed to start container: {error_msg}")
            return False

        time.sleep(2)

        # Verify container is running
        check_result = run_cmd(
            ["docker", "compose", "-f", str(self.temp_compose_file), "ps", "-q"],
            cwd=self.benchmark_dir,
            capture_output=True,
            text=True
        )

        if not check_result.stdout.strip():
            print(f"  ❌ Container is not running")
            return False

        # Clone repository and setup inside container (only if not using provided workspace)
        if not self.provided_workspace:
            if not self.clone_and_setup_in_container():
                return False

        return True

    def cleanup(self):
        """Clean up temporary files."""
        if self.temp_compose_file and self.temp_compose_file.exists():
            # Fix permissions inside container before stopping it
            # This ensures files created by root can be deleted by regular user
            try:
                # Check if container is still running
                check_result = run_cmd(
                    ["docker", "compose", "-f", str(self.temp_compose_file), "ps", "-q"],
                    cwd=self.benchmark_dir,
                    capture_output=True,
                    text=True
                )

                if check_result.stdout.strip():
                    # Container is running - fix permissions before stopping
                    run_cmd([
                        "docker", "compose", "-f", str(self.temp_compose_file), "exec", "-T",
                        "dev", "chmod", "-R", "777", "/workspace"
                    ], cwd=self.benchmark_dir, capture_output=True)
            except Exception:
                # If we can't fix permissions, continue anyway
                pass

            # Stop and remove container
            run_cmd(["docker", "compose", "-f", str(self.temp_compose_file), "down"],
                    cwd=self.benchmark_dir, capture_output=True)
            self.temp_compose_file.unlink()

        # Only delete workspace if we created it (not if it was provided)
        if self.owns_workspace and self.temp_workspace and self.temp_workspace.exists():
            # Delete the workspace - permissions should be fixed now
            try:
                shutil.rmtree(self.temp_workspace)
            except PermissionError:
                # If still can't delete, warn but don't fail
                print(f"  ⚠️  Warning: Could not delete temp workspace (files may be owned by root)")
                print(f"     Clean up manually: sudo rm -rf {self.temp_workspace}")
            except Exception as e:
                # Other errors - still warn but don't fail
                print(f"  ⚠️  Warning: Could not delete temp workspace: {e}")
                print(f"     Clean up manually: rm -rf {self.temp_workspace}")

    def parse_pytest_output(self, output: str) -> tuple:
        """Parse pytest output to extract test results."""
        passed = 0
        failed = 0
        errors = 0
        skipped = 0
        total = 0
        failed_tests = []

        # Get total from "collected X items"
        collected_match = re.search(r'collected (\d+) item', output)
        if collected_match:
            total = int(collected_match.group(1))

        summary_match = re.search(r'(\d+) passed', output)
        if summary_match:
            passed = int(summary_match.group(1))

        failed_match = re.search(r'(\d+) failed', output)
        if failed_match:
            failed = int(failed_match.group(1))

        errors_match = re.search(r'(\d+) error', output)
        if errors_match:
            errors = int(errors_match.group(1))

        skipped_match = re.search(r'(\d+) skipped', output)
        if skipped_match:
            skipped = int(skipped_match.group(1))

        # Fallback if collected not found
        if total == 0:
            total = passed + failed + errors + skipped

        failed_test_matches = re.findall(r'^FAILED\s+(\S+)', output, re.MULTILINE)
        failed_tests = failed_test_matches

        return passed, failed, errors, skipped, total, failed_tests

    def test_server(self, impl: ImplConfig, dir_name: str) -> TestResult:
        """Test a server implementation."""
        print(f"\n  Testing {dir_name} ({impl.language})...")

        build_time = 0.0
        test_time = 0.0

        if impl.install_cmd:
            print(f"    Installing dependencies...")
            start = time.time()
            returncode, _, _ = self.docker_exec(f"cd /workspace/{dir_name} && {impl.install_cmd}", env=impl.env)
            if returncode != 0:
                print(f"    ❌ Install failed")
                return TestResult(impl.language, False, 0, 0, 0, 0, 0, 0, 0.0, 0, [])
            build_time += time.time() - start
            print(f"    ✓ Install successful")

        if impl.build_cmd:
            print(f"    Building...")
            start = time.time()
            returncode, _, _ = self.docker_exec(f"cd /workspace/{dir_name} && {impl.build_cmd}", env=impl.env)
            build_time += time.time() - start
            if returncode != 0:
                print(f"    ❌ Build failed ({build_time:.1f}s)")
                return TestResult(impl.language, False, build_time, 0, 0, 0, 0, 0, 0.0, 0, [])
            print(f"    ✓ Build successful ({build_time:.1f}s)")

        print(f"    Starting server on port {impl.port}...")
        start_cmd = f"cd /workspace/{dir_name} && {impl.port_env_var}={impl.port} {impl.run_cmd}"
        self.docker_exec(start_cmd, env=impl.env, detached=True)
        time.sleep(impl.startup_wait)

        extra_args = self.config.pytest_args if self.config.pytest_args else "-v --tb=short"
        test_cmd = f"cd /workspace/tests && python3 -m pip install -q -r requirements.txt && python3 -m pytest --api-port={impl.port} {extra_args}"

        print(f"    Running tests...")
        start = time.time()
        returncode, stdout, stderr = self.docker_exec(test_cmd, capture_output=True)
        test_time = time.time() - start

        with open(self.log_file, 'a') as f:
            f.write(f"\n{'='*70}\n")
            f.write(f"Testing {dir_name} ({impl.language})\n")
            f.write(f"{'='*70}\n")
            f.write(stdout)
            f.write(stderr)

        output = stdout + stderr
        passed, failed, errors, skipped, total, failed_tests = self.parse_pytest_output(output)

        self.docker_exec(f"pkill -f '{impl.run_cmd}'")
        time.sleep(1)

        runnable = total - skipped
        pass_rate = passed / runnable if runnable > 0 else 0.0
        test_success = failed == 0 and errors == 0 and passed > 0

        if test_success:
            print(f"    ✅ Tests passed: {passed}/{runnable} ({test_time:.1f}s)")
        else:
            print(f"    ❌ Tests failed: {passed}/{runnable} ({test_time:.1f}s)")

        return TestResult(impl.language, True, build_time, passed, failed, errors, skipped, runnable, pass_rate, test_time, failed_tests)

    def test_cli(self, impl: ImplConfig, dir_name: str) -> TestResult:
        """Test a CLI implementation."""
        print(f"\n  Testing {dir_name} ({impl.language})...")

        build_time = 0.0
        test_time = 0.0

        if impl.install_cmd:
            print(f"    Installing dependencies...")
            start = time.time()
            returncode, _, _ = self.docker_exec(f"cd /workspace/{dir_name} && {impl.install_cmd}", env=impl.env)
            if returncode != 0:
                print(f"    ❌ Install failed")
                return TestResult(impl.language, False, 0, 0, 0, 0, 0, 0, 0.0, 0, [])
            build_time += time.time() - start
            print(f"    ✓ Install successful")

        if impl.build_cmd:
            print(f"    Building...")
            start = time.time()
            returncode, _, _ = self.docker_exec(f"cd /workspace/{dir_name} && {impl.build_cmd}", env=impl.env)
            build_time += time.time() - start
            if returncode != 0:
                print(f"    ❌ Build failed ({build_time:.1f}s)")
                return TestResult(impl.language, False, build_time, 0, 0, 0, 0, 0, 0.0, 0, [])
            print(f"    ✓ Build successful ({build_time:.1f}s)")

        work_dir = f"/workspace/{dir_name}"
        run_cmd = impl.run_cmd
        extra_args = self.config.pytest_args if self.config.pytest_args else "-v --tb=short"
        test_cmd = f"cd /workspace/tests && python3 -m pip install -q -r requirements.txt && python3 -m pytest --work-dir='{work_dir}' --run-cmd='{run_cmd}' {extra_args}"

        print(f"    Running tests...")
        start = time.time()
        returncode, stdout, stderr = self.docker_exec(test_cmd, capture_output=True)
        test_time = time.time() - start

        with open(self.log_file, 'a') as f:
            f.write(f"\n{'='*70}\n")
            f.write(f"Testing {dir_name} ({impl.language})\n")
            f.write(f"{'='*70}\n")
            f.write(stdout)
            f.write(stderr)

        output = stdout + stderr
        passed, failed, errors, skipped, total, failed_tests = self.parse_pytest_output(output)

        runnable = total - skipped
        pass_rate = passed / runnable if runnable > 0 else 0.0
        test_success = failed == 0 and errors == 0 and passed > 0

        if test_success:
            print(f"    ✅ Tests passed: {passed}/{runnable} ({test_time:.1f}s)")
        else:
            print(f"    ❌ Tests failed: {passed}/{runnable} ({test_time:.1f}s)")

        return TestResult(impl.language, True, build_time, passed, failed, errors, skipped, runnable, pass_rate, test_time, failed_tests)

    def test_impl(self, impl: ImplConfig, dir_name: str) -> TestResult:
        """Test an implementation based on benchmark type."""
        if self.config.benchmark_type == 'server':
            return self.test_server(impl, dir_name)
        elif self.config.benchmark_type == 'cli':
            return self.test_cli(impl, dir_name)
        else:
            print(f"    ❌ Unknown benchmark type: {self.config.benchmark_type}")
            return TestResult(impl.language, False, 0, 0, 0, 0, 0, 0, 0.0, 0, [])

    def run(self, test_target: str) -> dict:
        """Run the benchmark and return results."""
        src_result = None
        dst_result = None

        try:
            if not self.setup():
                # Setup failed, return empty results
                return self._build_result(None, None)

            if test_target in ['src', 'both']:
                print(f"\n  === TESTING SOURCE ===")
                src_result = self.test_impl(self.config.source, "src")

            if test_target in ['dst', 'both']:
                print(f"\n  === TESTING DESTINATION ===")
                dst_result = self.test_impl(self.config.destination, "dst")

            return self._build_result(src_result, dst_result)
        finally:
            self.cleanup()

    def _build_result(self, src_result: Optional[TestResult], dst_result: Optional[TestResult]) -> dict:
        """Build the result dictionary."""
        result = {
            "id": self.config.id,
            "benchmark": self.config.name,
            "type": self.config.benchmark_type,
            "timestamp": datetime.now().isoformat(),
        }

        if src_result:
            result["source"] = {
                "language": src_result.language,
                "build_success": src_result.build_success,
                "build_time_seconds": round(src_result.build_time_seconds, 2),
                "tests_passed": src_result.tests_passed,
                "tests_failed": src_result.tests_failed,
                "tests_error": src_result.tests_error,
                "tests_skipped": src_result.tests_skipped,
                "tests_total": src_result.tests_total,
                "pass_rate": round(src_result.pass_rate, 4),
                "test_time_seconds": round(src_result.test_time_seconds, 2),
                "failed_tests": src_result.failed_tests,
            }

        if dst_result:
            result["destination"] = {
                "language": dst_result.language,
                "build_success": dst_result.build_success,
                "build_time_seconds": round(dst_result.build_time_seconds, 2),
                "tests_passed": dst_result.tests_passed,
                "tests_failed": dst_result.tests_failed,
                "tests_error": dst_result.tests_error,
                "tests_skipped": dst_result.tests_skipped,
                "tests_total": dst_result.tests_total,
                "pass_rate": round(dst_result.pass_rate, 4),
                "test_time_seconds": round(dst_result.test_time_seconds, 2),
                "failed_tests": dst_result.failed_tests,
            }

        return result


def print_summary(results: List[dict]):
    """Print a summary of all benchmark results."""
    print(f"\n{'='*70}")
    print("SUMMARY")
    print("="*70)

    for r in results:
        status_parts = []
        if "source" in r:
            src = r["source"]
            src_status = "✅" if src["build_success"] and src["tests_failed"] == 0 and src["tests_error"] == 0 else "❌"
            status_parts.append(f"src:{src_status} {src['tests_passed']}/{src['tests_total']}")
        if "destination" in r:
            dst = r["destination"]
            dst_status = "✅" if dst["build_success"] and dst["tests_failed"] == 0 and dst["tests_error"] == 0 else "❌"
            status_parts.append(f"dst:{dst_status} {dst['tests_passed']}/{dst['tests_total']}")

        print(f"[{r['id']}] {r['benchmark']}: {' | '.join(status_parts)}")

    print("="*70)


def run_tests_for_benchmark(
    benchmark_name: str,
    benchmark_dir: Path,
    output_dir: Path,
    test_target: str = "dst",
    workspace_dir: Path = None
) -> dict:
    """
    Run tests for a single benchmark. Designed to be called from run_experiment.py.

    Args:
        benchmark_name: Name of the benchmark
        benchmark_dir: Path to benchmark directory (e.g., dataset/jq-gojq)
        output_dir: Directory to save test_result.json
        test_target: Which implementation to test ('src', 'dst', 'both')
        workspace_dir: Optional custom workspace path for isolation

    Returns:
        dict with test results including build_success, tests_passed, tests_total, pass_rate
    """
    # Load benchmark config
    config = load_benchmark_config(benchmark_dir)
    if not config:
        print(f"  ❌ Failed to load benchmark config from {benchmark_dir}")
        return {
            "build_success": False,
            "tests_passed": 0,
            "tests_total": 0,
            "pass_rate": 0.0,
            "error": "Failed to load benchmark config"
        }

    # Create log file
    log_file = output_dir / "test_log.txt"

    # Run benchmark
    runner = BenchmarkRunner(config, log_file, workspace_dir)
    result = runner.run(test_target)

    # Save test_result.json
    result_file = output_dir / "test_result.json"
    with open(result_file, 'w') as f:
        json.dump(result, f, indent=2)

    # Extract key metrics for return
    if test_target == "dst" and "destination" in result:
        dst = result["destination"]
        return {
            "build_success": dst["build_success"],
            "tests_passed": dst["tests_passed"],
            "tests_failed": dst["tests_failed"],
            "tests_total": dst["tests_total"],
            "pass_rate": dst["pass_rate"],
            "test_time_seconds": dst["test_time_seconds"],
            "failed_tests": dst["failed_tests"]
        }
    elif test_target == "src" and "source" in result:
        src = result["source"]
        return {
            "build_success": src["build_success"],
            "tests_passed": src["tests_passed"],
            "tests_failed": src["tests_failed"],
            "tests_total": src["tests_total"],
            "pass_rate": src["pass_rate"],
            "test_time_seconds": src["test_time_seconds"],
            "failed_tests": src["failed_tests"]
        }
    else:
        # Return full result for 'both' or unexpected cases
        return result


def load_config(config_path: Path) -> dict:
    """Load configuration from TOML file."""
    if not config_path.exists():
        return {}
    with open(config_path, 'rb') as f:
        return tomllib.load(f)


def run_single_benchmark(benchmark_id: str, benchmark_dir: Path, log_session_dir: Path, test_target: str) -> dict:
    """Run a single benchmark and return results."""
    config = load_benchmark_config(benchmark_dir)
    if not config:
        return {
            "id": benchmark_id,
            "benchmark": benchmark_dir.name,
            "type": "unknown",
            "timestamp": datetime.now().isoformat(),
            "error": "Failed to load benchmark config",
        }

    log_file = log_session_dir / f"{config.name}.log"
    with open(log_file, 'w') as f:
        f.write("MCode Benchmark Log\n")
        f.write(f"Benchmark: {config.name}\n")
        f.write(f"ID: {benchmark_id}\n")
        f.write(f"Timestamp: {datetime.now().isoformat()}\n")

    runner = BenchmarkRunner(config, log_file)
    return runner.run(test_target)


def main():
    parser = argparse.ArgumentParser(description='Run MCode benchmarks')
    parser.add_argument('--config', type=str, default='config.toml',
                        help='Path to config file (default: config.toml)')
    parser.add_argument('--test', choices=['src', 'dst', 'both'], default='src',
                        help='Which implementation to test (default: src)')
    parser.add_argument('--results-dir', type=str, default='results',
                        help='Results output directory (default: results)')
    parser.add_argument('--logs-dir', type=str, default='logs',
                        help='Logs output directory (default: logs)')
    parser.add_argument('--parallel', type=int, default=1,
                        help='Number of parallel workers (default: 1)')
    args = parser.parse_args()

    # Setup paths
    root_dir = Path(__file__).parent
    dataset_dir = root_dir / "dataset"
    results_dir = root_dir / args.results_dir
    logs_dir = root_dir / args.logs_dir
    config_path = root_dir / args.config

    # Load config
    config = load_config(config_path)
    benchmark_ids = config.get('benchmarks', {}).get('ids', [])

    # Create output directories
    results_dir.mkdir(exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_session_dir = logs_dir / timestamp
    log_session_dir.mkdir(parents=True, exist_ok=True)

    # Discover benchmarks
    benchmarks = discover_benchmarks(dataset_dir)

    if not benchmarks:
        print("No benchmarks found in dataset/")
        sys.exit(1)

    # Filter by IDs from config
    if benchmark_ids:
        benchmarks = {k: v for k, v in benchmarks.items() if k in benchmark_ids}
        if not benchmarks:
            print(f"No benchmarks found matching IDs: {benchmark_ids}")
            print(f"Available IDs: {list(discover_benchmarks(dataset_dir).keys())}")
            sys.exit(1)

    # Sort by ID
    sorted_benchmarks = sorted(benchmarks.items(), key=lambda x: x[0])

    print("="*70)
    print("MCode Benchmark Runner")
    print(f"Running {len(sorted_benchmarks)} benchmark(s)")
    print(f"Test target: {args.test}")
    print(f"Parallel: {args.parallel}")
    print("="*70)

    results = []

    if args.parallel <= 1:
        for benchmark_id, config in sorted_benchmarks:
            print(f"\n{'='*70}")
            print(f"[{benchmark_id}] {config.name}")
            print(f"Type: {config.benchmark_type} | {config.source.language} → {config.destination.language}")
            print("="*70)

            log_file = log_session_dir / f"{config.name}.log"
            with open(log_file, 'w') as f:
                f.write("MCode Benchmark Log\n")
                f.write(f"Benchmark: {config.name}\n")
                f.write(f"ID: {benchmark_id}\n")
                f.write(f"Timestamp: {datetime.now().isoformat()}\n")

            runner = BenchmarkRunner(config, log_file)
            result = runner.run(args.test)
            results.append(result)
    else:
        print(f"\nRunning {len(sorted_benchmarks)} benchmark(s) with {args.parallel} workers...")
        with ProcessPoolExecutor(max_workers=args.parallel) as executor:
            future_to_benchmark = {
                executor.submit(
                    run_single_benchmark,
                    benchmark_id,
                    config.benchmark_dir,
                    log_session_dir,
                    args.test
                ): (benchmark_id, config.name)
                for benchmark_id, config in sorted_benchmarks
            }

            for future in as_completed(future_to_benchmark):
                benchmark_id, benchmark_name = future_to_benchmark[future]
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    results.append({
                        "id": benchmark_id,
                        "benchmark": benchmark_name,
                        "type": "unknown",
                        "timestamp": datetime.now().isoformat(),
                        "error": str(e),
                    })

    # Save results
    results = sorted(results, key=lambda r: r.get("id", ""))
    results_file = results_dir / "results.jsonl"
    with open(results_file, 'w') as f:
        for r in results:
            f.write(json.dumps(r) + "\n")

    # Print summary
    print_summary(results)

    print(f"\nResults saved to: {results_file}")
    print(f"Logs saved to: {log_session_dir}/")

    # Exit code: 0 if all passed, 1 if any failed
    all_passed = all(
        (r.get("source", {}).get("tests_failed", 0) == 0 if "source" in r else True) and
        (r.get("destination", {}).get("tests_failed", 0) == 0 if "destination" in r else True)
        for r in results
    )
    sys.exit(0 if all_passed else 1)


if __name__ == "__main__":
    main()