import argparse
import asyncio
import contextlib
import csv
import importlib
import logging
import os
import sys
import threading
import time
from datetime import datetime
from pathlib import Path

from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

from clients.harness.problem_id import HARNESS_ARTIFACT_ID_ENV, HARNESS_PROBLEM_ID_ENV
from logger import console, init_logger
from sregym.agent_launcher import AgentLauncher
from sregym.agent_registry import get_agent, list_agents
from sregym.conductor.conductor import Conductor, ConductorConfig
from sregym.conductor.conductor_api import request_shutdown, run_api
from sregym.conductor.constants import StartProblemResult
from sregym.run_artifacts import ArtifactFinalizationError, RunArtifacts
from sregym.service.container_runner import ContainerRunner, ExecInput
from sregym.traces import postprocess as trace_postprocess
from sregym.traces import store as trace_store

LAUNCHER = AgentLauncher()
logger = logging.getLogger(__name__)
_driver_results: list[dict] = []
_driver_base_dir: Path | None = None


def run_preflight_check(
    agent_name: str,
    container_runner: ContainerRunner | None = None,
    install_script: str | None = None,
) -> None:
    """Run the agent's pre-flight check inside the container."""

    # Agents that need pre-flight check
    agent_driver_modules: dict[str, str] = {
        "stratus": "clients.stratus.stratus_agent.driver.driver",
        "claudecode": "clients.claudecode.driver",
        "codex": "clients.codex.driver",
        "copilot": "clients.copilot.driver",
        "opencode": "clients.opencode.driver",
        "gemini": "clients.geminicli.driver",
    }

    module_path = agent_driver_modules.get(agent_name)
    if not module_path:
        return

    driver_mod = importlib.import_module(module_path)
    if not hasattr(driver_mod, "run_preflight"):
        return

    if container_runner is None:
        logger.warning(f"⚠️  No container runner — skipping pre-flight check for '{agent_name}'")
        return

    check_cmd = f"python3 -c 'from {module_path} import run_preflight; run_preflight()'"
    if install_script:
        check_cmd = f"/opt/sregym/install-scripts/{install_script} > /dev/null 2>&1 && {check_cmd}"

    logger.info(f"🔍 Running pre-flight check for '{agent_name}'...")
    result = container_runner.run_sync(ExecInput(command=check_cmd, label="preflight", timeout=180))
    if result.returncode != 0:
        if result.stdout:
            print(result.stdout.strip())
        if result.stderr:
            print(result.stderr.strip())
        logger.error(f"❌ Pre-flight check failed for '{agent_name}'")
        sys.exit(1)

    logger.info(f"✅ Pre-flight check passed for '{agent_name}'")


def run_judge_preflight_check() -> None:
    """Validate the judge model and credentials through its LiteLLM backend."""
    from llm_backend.init_backend import get_llm_backend_for_judge

    logger.info("🔍 Running pre-flight check for judge model...")
    try:
        get_llm_backend_for_judge().inference("Say ok.", system_prompt="Reply with exactly 'ok'.")
    except Exception as e:
        logger.error(f"❌ Judge pre-flight check failed: {e}")
        logger.error(
            "The judge uses LiteLLM. Ensure the model is LiteLLM-compatible or pass a compatible --judge-model."
        )
        sys.exit(1)

    logger.info("✅ Judge pre-flight check passed")


def get_current_datetime_formatted():
    now = datetime.now()
    formatted_datetime = now.strftime("%m%d_%H%M")
    return formatted_datetime


def _restore_env_var(name: str, previous_value: str | None) -> None:
    if previous_value is None:
        os.environ.pop(name, None)
    else:
        os.environ[name] = previous_value


def _env_status(name: str) -> str:
    return "set" if os.environ.get(name) else "unset"


def _is_opencode_local_model(agent: str | None, model: str) -> bool:
    return agent == "opencode" and model.startswith("local/")


def _normalize_opencode_local_model_for_litellm(model: str) -> str:
    return f"openai/{model.split('/', 1)[1]}"


def _configure_model_environment(args) -> tuple[str, str]:
    agent_model = args.model
    raw_judge_model = args.judge_model or args.model
    normalizes_opencode_local_judge = _is_opencode_local_model(args.agent, raw_judge_model)
    judge_model = (
        _normalize_opencode_local_model_for_litellm(raw_judge_model)
        if normalizes_opencode_local_judge
        else raw_judge_model
    )

    os.environ["AGENT_MODEL_ID"] = agent_model
    os.environ["JUDGE_MODEL_ID"] = judge_model

    if not getattr(args, "judge_model", None) or normalizes_opencode_local_judge:
        if not os.environ.get("JUDGE_API_BASE") and os.environ.get("AGENT_API_BASE"):
            os.environ["JUDGE_API_BASE"] = os.environ["AGENT_API_BASE"]
        if not os.environ.get("JUDGE_API_KEY") and os.environ.get("AGENT_API_KEY"):
            os.environ["JUDGE_API_KEY"] = os.environ["AGENT_API_KEY"]

    if normalizes_opencode_local_judge:
        if not os.environ.get("JUDGE_API_BASE"):
            raise ValueError("AGENT_API_BASE or JUDGE_API_BASE is required to use an OpenCode local model as the judge")
        if not os.environ.get("JUDGE_API_KEY"):
            os.environ["JUDGE_API_KEY"] = "dummy"

    return agent_model, judge_model


@contextlib.contextmanager
def _artifact_environment(run: RunArtifacts):
    previous = {
        "AGENT_LOGS_DIR": os.environ.get("AGENT_LOGS_DIR"),
        HARNESS_ARTIFACT_ID_ENV: os.environ.get(HARNESS_ARTIFACT_ID_ENV),
        HARNESS_PROBLEM_ID_ENV: os.environ.get(HARNESS_PROBLEM_ID_ENV),
    }
    os.environ["AGENT_LOGS_DIR"] = str(run.active_dir.resolve())
    os.environ[HARNESS_ARTIFACT_ID_ENV] = run.artifact_id
    os.environ.pop(HARNESS_PROBLEM_ID_ENV, None)
    try:
        yield
    finally:
        for name, value in previous.items():
            _restore_env_var(name, value)


def driver_loop(
    conductor: Conductor,
    problem_filter: str | None = None,
    agent_to_run: str | None = None,
    use_external_harness: bool = False,
    n_attempts: int = 1,
    agent_timeout: int = 1800,
    resume_csv: str | None = None,
):
    """
    Deploy each problem and wait for HTTP grading via POST /submit.
    Returns a list of flattened dicts with results per problem.

    Args:
        conductor: The Conductor instance
        problem_filter: Optional problem ID to run. If specified, only this problem will be run.
        agent_to_run: Agent name to run (required unless use_external_harness is True).
        use_external_harness: If True, inject fault and exit without running evaluation logic.
        n_attempts: Number of end-to-end attempts to run each problem.
        resume_csv: Path to a previous results CSV to resume from (skip completed problems).
    """

    async def driver():
        base_dir = Path("results") / get_current_datetime_formatted()
        base_dir.mkdir(parents=True, exist_ok=True)
        global _driver_base_dir
        _driver_base_dir = base_dir
        # give the API a moment to bind
        await asyncio.sleep(1)

        # Verify agent exists in registry (skip if using external harness)
        if not use_external_harness:
            available_agents = list_agents(path=Path(os.path.dirname(os.path.abspath(__file__))) / "agents.yaml").keys()
            if agent_to_run not in available_agents:
                console.log(f"⚠️ Agent '{agent_to_run}' not found in registry. Available agents: {available_agents}")
                sys.exit(1)

            console.log(f"Starting agent now: {agent_to_run}")
            conductor.register_agent(agent_to_run)

            # Start K8s API proxy to hide chaos engineering namespaces from the agent
            console.log("🔒 Starting Kubernetes API proxy to hide chaos namespaces...")
            conductor.start_k8s_proxy()
            LAUNCHER.set_agent_kubeconfig(conductor.get_agent_kubeconfig_path())

        all_results_for_agent = []

        # Get all problem IDs and filter if needed
        problem_ids = conductor.problems.get_problem_ids()
        all_problem_ids = conductor.problems.get_problem_ids(all=True)
        if problem_filter:
            if problem_filter not in all_problem_ids:
                console.log(f"⚠️  Problem '{problem_filter}' not found in registry. Available problems: {problem_ids}")
                sys.exit(1)
            problem_ids = [problem_filter]
            console.log(f"🎯 Running single problem: {problem_filter}")

        # sanity check: are there any specified problem ids that do not exist in the registry?
        unknown_problem_ids = set(problem_ids) - set(all_problem_ids)
        if unknown_problem_ids:
            console.log(
                f"⚠️  These problem ids do not exist in the registry and they will be skipped: {unknown_problem_ids}"
            )
        for unknown_problem_id in unknown_problem_ids:
            problem_ids.remove(unknown_problem_id)

        # Resume support: load completed problems from previous CSV and pre-seed results
        from collections import Counter

        completed_problems: set[str] = set()
        attempt_counts: Counter[str] = Counter()
        if resume_csv:
            try:
                with open(resume_csv, newline="") as f:
                    reader = csv.DictReader(f)
                    resume_rows = list(reader)
                # Group by problem_id and count attempts
                attempt_counts = Counter(r["problem_id"] for r in resume_rows)
                completed_problems = {pid for pid, count in attempt_counts.items() if count >= n_attempts}

                for row in resume_rows:
                    all_results_for_agent.append(row)
                console.log(
                    f"📋 Resuming from {resume_csv}: {len(completed_problems)} problems already done, skipping them"
                )
            except Exception as e:
                console.log(f"⚠️  Failed to load resume CSV: {e}")

        # Bar tracks attempts (problems × n_attempts), not problems, so the
        # 0/N total reflects total work even when n_attempts > 1.
        already_done = sum(min(attempt_counts.get(p, 0), n_attempts) for p in problem_ids)
        progress = Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
            console=console,
        )
        task_id = progress.add_task(
            f"[cyan]Benchmarking {agent_to_run or 'agent'}",
            total=len(problem_ids) * n_attempts,
            completed=already_done,
        )
        progress.start()

        for pid in problem_ids:
            if pid in completed_problems:
                console.log(f"⏭️  Skipping already-completed problem: {pid}")
                progress.advance(task_id, n_attempts)
                continue

            conductor.problem_id = pid

            # Keep a record of results for this problem in a temp file in case an attempt fails
            tmp_path = f"_running_{pid}_{agent_to_run}_results.csv"

            for attempt in range(1, n_attempts + 1):
                progress.update(
                    task_id,
                    description=f"[cyan]Benchmarking {agent_to_run or 'agent'} — {pid} (attempt {attempt}/{n_attempts})",
                )
                console.log(f"\n🔍 Starting problem: {pid} (Attempt {attempt} of {n_attempts})")

                # Retry start_problem up to 3 times to handle transient deploy failures
                max_deploy_retries = 3
                result = None
                for deploy_attempt in range(1, max_deploy_retries + 1):
                    try:
                        result = await conductor.start_problem()
                        break  # Success — exit retry loop
                    except Exception as e:
                        console.log(
                            f"❌ start_problem failed for '{pid}' "
                            f"(deploy attempt {deploy_attempt}/{max_deploy_retries}): {e}"
                        )
                        if deploy_attempt < max_deploy_retries:
                            console.log("🧹 Cleaning up before retry...")
                            try:
                                conductor._finish_problem()
                            except Exception as cleanup_err:
                                console.log(f"⚠️  Cleanup error (non-fatal): {cleanup_err}")
                            console.log(f"🔄 Retrying start_problem for '{pid}'...")
                        else:
                            console.log(
                                f"⛔ All {max_deploy_retries} deploy attempts failed for '{pid}', skipping this attempt"
                            )
                            result = None

                if result is None:
                    # The inner retry loop already logged the failure. Don't crash the
                    # entire driver — record the failure, clean up cluster state, and
                    # move on to the next problem so the benchmark can keep making progress.
                    try:
                        conductor._finish_problem()
                    except Exception as cleanup_err:
                        console.log(f"⚠️  Cleanup after exhausted deploy retries failed (non-fatal): {cleanup_err}")
                    snapshot = {
                        "problem_id": pid,
                        "attempt": attempt,
                        "deploy_failed": True,
                    }
                    all_results_for_agent.append(snapshot)
                    fieldnames = sorted({key for row in all_results_for_agent for key in row})
                    with open(tmp_path, "w", newline="") as csvfile:
                        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                        writer.writeheader()
                        writer.writerows(all_results_for_agent)
                    console.log(f"⏭️  Skipping remaining attempts for '{pid}' and moving to next problem")
                    # Account for this attempt + remaining skipped attempts on the bar.
                    progress.advance(task_id, n_attempts - attempt + 1)
                    break

                if result == StartProblemResult.SKIPPED_KHAOS_REQUIRED:
                    console.log(f"⏭️  Skipping problem '{pid}': requires Khaos but running on emulated cluster")
                    progress.advance(task_id, n_attempts - attempt + 1)
                    break  # Skip to next problem

                # If using external harness, fault is injected - exit now
                if use_external_harness:
                    console.log(f"✅ Fault injected for problem '{pid}'. Exiting for external harness.")
                    progress.stop()
                    return []

                assert agent_to_run is not None

                run = RunArtifacts.create(
                    staging_root=Path(".runtime"),
                    results_root=base_dir,
                    problem_id=pid,
                    agent=agent_to_run,
                    attempt=attempt,
                )
                agent_proc = None

                with _artifact_environment(run):
                    reg = get_agent(
                        agent_to_run,
                        path=Path(os.path.dirname(os.path.abspath(__file__))) / "agents.yaml",
                    )
                    if reg:
                        agent_proc = await LAUNCHER.ensure_started(reg)

                timed_out = False
                agent_start_time = time.time()
                while conductor.submission_stage != "done":
                    if time.time() - agent_start_time > agent_timeout:
                        timed_out = True
                        console.log(f"⏰ Agent timeout ({agent_timeout}s) exceeded, killing agent")
                        LAUNCHER.cleanup_agent(agent_to_run)
                        conductor.results["timed_out"] = True
                        conductor.results["agent_timeout_seconds"] = agent_timeout
                        console.log("🧹 Running conductor cleanup after agent timeout...")
                        conductor._finish_problem()
                        break

                    tracked_proc = LAUNCHER._procs.get(agent_to_run) or agent_proc
                    if tracked_proc:
                        tracked_proc.proc.poll()
                        if tracked_proc.proc.returncode is not None:
                            agent_proc = tracked_proc
                            console.log(f"⚠️  Agent process exited with return code {tracked_proc.proc.returncode}")
                            if conductor._submit_future is not None and not conductor._submit_future.done():
                                console.log("⏳ Waiting for conductor evaluation to complete...")
                                try:
                                    await asyncio.wait_for(
                                        asyncio.wrap_future(conductor._submit_future),
                                        timeout=300,
                                    )
                                except TimeoutError:
                                    console.log("⚠️  Conductor evaluation did not finish within 300s")
                                except Exception as e:
                                    console.log(f"⚠️  Conductor evaluation raised: {e}")
                            console.log("🧹 Running conductor cleanup after agent exit...")
                            conductor._finish_problem()
                            break
                    await asyncio.sleep(1)

                console.log(f"✅ Completed {pid}: results={conductor.results}", markup=False)
                # Wait for agent process to complete naturally before cleanup
                # This allows the agent to finish saving trajectories and other cleanup tasks
                if not use_external_harness:
                    agent_proc = LAUNCHER._procs.get(agent_to_run) or agent_proc
                    if agent_proc and not timed_out:
                        console.log("⏳ Waiting for agent process to complete...")
                        timeout = 60  # seconds
                        elapsed = 0
                        while elapsed < timeout:
                            agent_proc.proc.poll()
                            if agent_proc.proc.returncode is not None:
                                console.log(f"✅ Agent process completed with return code {agent_proc.proc.returncode}")
                                break
                            await asyncio.sleep(1)
                            elapsed += 1
                        else:
                            console.log(f"⚠️  Agent process did not complete within {timeout}s, will force cleanup")

                    # Publication must happen only after cleanup_agent has reaped
                    # the tracked docker client or shell process tree.
                    LAUNCHER.cleanup_agent(agent_to_run)
                    console.log(f"🧹 Cleaned up agent process for {agent_to_run}")

                snapshot = {
                    "problem_id": pid,
                    "attempt": attempt,
                }
                for stage, outcome in conductor.results.items():
                    if isinstance(outcome, dict):
                        for k, v in outcome.items():
                            snapshot[f"{stage}.{k}"] = v
                    else:
                        snapshot[stage] = outcome

                fieldnames = sorted({key for row in [*all_results_for_agent, snapshot] for key in row})
                ownership_image = (
                    LAUNCHER._container_runner.config.image
                    if LAUNCHER._container_runner is not None
                    else "sregym-agent-base:latest"
                )
                try:
                    published_run_dir = run.finalize_and_publish(
                        snapshot=snapshot,
                        fieldnames=fieldnames,
                        ownership_image=ownership_image,
                    )
                except ArtifactFinalizationError as exc:
                    snapshot["artifact_finalization_failed"] = True
                    snapshot["artifact_staging_path"] = str(run.active_dir)
                    logger.error(
                        "Artifact finalization failed for %s attempt %s; staging retained at %s: %s",
                        pid,
                        attempt,
                        run.active_dir,
                        exc,
                    )
                    published_run_dir = None

                all_results_for_agent.append(snapshot)
                fieldnames = sorted({key for row in all_results_for_agent for key in row})
                with open(tmp_path, "w", newline="") as csvfile:
                    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                    writer.writeheader()
                    writer.writerows(all_results_for_agent)

                if published_run_dir is not None:
                    logger.info(
                        f"⏳ Attempt {attempt} of {n_attempts} for problem {pid} complete - "
                        f"Intermediate results written to {tmp_path}; artifacts published to {published_run_dir}"
                    )
                    # Normalize the agent's raw logs into an ATIF trajectory.json
                    # right after publication. Non-fatal: a conversion failure or
                    # an unsupported tool must never break the run.
                    trajectory_path = trace_postprocess.write_trajectory(published_run_dir)
                    if trajectory_path is not None:
                        logger.info(f"📝 ATIF trajectory written to {trajectory_path}")
                        try:
                            trace_store.ingest_trajectory_file(
                                trajectory_path,
                                base_dir.parent / trace_store.DEFAULT_DB_PATH,  # results/traces.db
                            )
                        except Exception as exc:  # defensive: never abort a run
                            logger.warning(f"⚠️ ATIF trajectory DB ingest failed: {exc}")
                    else:
                        logger.warning(f"⚠️ ATIF trajectory conversion skipped for {published_run_dir}")

                if attempt == n_attempts:
                    final_csv_path = base_dir / agent_to_run / pid / f"{pid}_{agent_to_run}_results.csv"
                    final_csv_path.parent.mkdir(parents=True, exist_ok=True)
                    os.replace(tmp_path, final_csv_path)
                    logger.info(
                        f"✅ Problem {pid} for agent {agent_to_run} complete! Results written to {final_csv_path}"
                    )

                progress.advance(task_id)

        progress.stop()

        # Stop K8s API proxy when all problems are done
        if not use_external_harness:
            console.log("🔓 Stopping Kubernetes API proxy...")
            conductor.stop_k8s_proxy()

        return [{agent_to_run: all_results_for_agent}]

    return asyncio.run(driver())


def _run_driver_and_shutdown(
    conductor: Conductor,
    problem_filter: str | None = None,
    agent_to_run: str | None = None,
    use_external_harness: bool = False,
    n_attempts: int = 1,
    agent_timeout: int = 1800,
    resume_csv: str | None = None,
):
    """Run the benchmark driver, stash results, then tell the API to exit."""
    try:
        results = driver_loop(
            conductor,
            problem_filter=problem_filter,
            agent_to_run=agent_to_run,
            use_external_harness=use_external_harness,
            n_attempts=n_attempts,
            agent_timeout=agent_timeout,
            resume_csv=resume_csv,
        )
        global _driver_results
        _driver_results = results
    except Exception:
        logger.exception("Driver thread crashed")
    finally:
        LAUNCHER.cleanup_all()
        request_shutdown()


def main(args):
    # set up the logger
    init_logger()

    agent_model, judge_model = _configure_model_environment(args)

    if args.noise:
        logger.info("Noise injection enabled.")
    os.environ["API_HOSTNAME"] = "0.0.0.0"
    os.environ["API_PORT"] = "8000"
    os.environ["MCP_SERVER_PORT"] = "9954"
    os.environ["MCP_SERVER_URL"] = "http://127.0.0.1:9954"

    logger.info(
        f"🔧 Config — agent: {args.agent}, agent_model: {agent_model}, judge_model: {judge_model}, "
        f"agent_api_base: {_env_status('AGENT_API_BASE')}, judge_api_base: {_env_status('JUDGE_API_BASE')}"
    )

    if not args.use_external_harness:
        run_judge_preflight_check()

    # Only build/check agent container image if the agent requires it
    agent_reg = (
        get_agent(args.agent, path=Path(os.path.dirname(os.path.abspath(__file__))) / "agents.yaml")
        if args.agent
        else None
    )
    if not agent_reg or agent_reg.container_isolation:
        LAUNCHER.enable_container_isolation(force_build=args.force_build)

    # Pre-flight check — makes a real (minimal) API call inside the agent
    # container to validate model and credentials in one shot.
    run_preflight_check(
        args.agent,
        container_runner=LAUNCHER._container_runner,
        install_script=agent_reg.install_script if agent_reg else None,
    )

    conductor_config = ConductorConfig(deploy_loki=not args.use_external_harness, enable_noise=args.noise)
    conductor = Conductor(config=conductor_config)

    # Start the driver in the background; it will call request_shutdown() when finished
    driver_thread = threading.Thread(
        target=_run_driver_and_shutdown,
        args=(
            conductor,
            args.problem,
            args.agent,
            args.use_external_harness,
            args.n_attempts,
            args.agent_timeout,
            args.resume,
        ),
        name="driver",
        daemon=True,
    )
    driver_thread.start()

    # Start the Conductor HTTP API in the MAIN thread (blocking)
    try:
        run_api(conductor)
    except KeyboardInterrupt:
        # If interrupted, still try to shut down cleanly
        LAUNCHER.cleanup_all()
        request_shutdown()
    finally:
        # Stop any remaining agent containers/processes
        LAUNCHER.cleanup_all()

        # Stop noise manager if it was enabled
        if args.noise:
            try:
                from sregym.generators.noise.manager import get_noise_manager

                logger.info("Stopping noise manager...")
                get_noise_manager().stop()
            except Exception as e:
                logger.error(f"⚠️ Error stopping noise manager: {e}")

        # Give driver a moment to finish setting results
        driver_thread.join(timeout=5)

    # When API shuts down, collect results from driver
    results = _driver_results

    if results:
        aggregated = {}
        for entry in results:
            for agent_name, agent_rows in entry.items():
                aggregated.setdefault(agent_name, []).extend(agent_rows)

        for agent_name, agent_results in aggregated.items():
            fieldnames = sorted({key for row in agent_results for key in row})
            out_dir = _driver_base_dir if _driver_base_dir else Path("results")
            csv_path = out_dir / f"{agent_name}_ALL_results.csv"
            with open(csv_path, "w", newline="") as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(agent_results)
            logger.info(f"✅ Benchmark complete! Results for {agent_name} written to {csv_path}")
    else:
        logger.warning("⚠️ No results to write.")

    if __name__ == "__main__":
        # separate run, use exit
        sys.exit(0)
    else:
        # function call run, return results
        return results


if __name__ == "__main__":
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Run SREGym benchmark suite")
    parser.add_argument(
        "--problem",
        type=str,
        default=None,
        help="Run only a specific problem by its ID (e.g., 'target_port')",
    )
    parser.add_argument(
        "--agent",
        type=str,
        default="stratus",
        help="Agent to run (default: stratus)",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="gpt-5",
        help="LiteLLM model string (e.g. anthropic/claude-sonnet-4-6-20250627, gpt-5, gemini/gemini-2.5-pro)",
    )
    parser.add_argument(
        "--judge-model",
        type=str,
        default=None,
        help="Model for the LLM-as-a-judge evaluator (defaults to --model if not set)",
    )
    parser.add_argument(
        "--use-external-harness", action="store_true", help="For use in external harnesses, deploy the fault and exit."
    )
    parser.add_argument(
        "--noise",
        action="store_true",
        help="Enable transient noise injection via Chaos Mesh during problem runs",
    )
    parser.add_argument(
        "--n-attempts",
        type=int,
        default=1,
        help="Number of attempts to run each problem (default: 1)",
    )
    parser.add_argument(
        "--force-build",
        action="store_true",
        help="Force rebuild the agent Docker image even if it already exists (use after updating dependencies or build scripts)",
    )
    parser.add_argument(
        "--agent-timeout",
        type=int,
        default=1800,
        help="Agent timeout in seconds after deployment (default: 1800)",
    )
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Resume from a previous results CSV file. Problems already in the CSV will be skipped.",
    )
    args = parser.parse_args()

    # Validate that n_attempts is positive
    if args.n_attempts is not None and args.n_attempts < 1:
        parser.error("--n-attempts must be a positive integer")

    main(args)
