# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import argparse
import asyncio
import contextlib
import json
import logging
import os
import shutil
import sys
import tempfile
import textwrap
import time
from importlib import resources
from pathlib import Path

import yaml
from pydantic import ValidationError

from arksim import __version__
from arksim.evaluator import (
    Evaluation,
    EvaluationInput,
    check_numeric_thresholds,
    check_qualitative_failure_labels,
    run_evaluation,
)
from arksim.simulation_engine import SimulationInput, run_simulation
from arksim.utils.logger import get_logger

# Exit codes
EXIT_OK = 0
EXIT_EVAL_FAILED = 1
EXIT_CONFIG_ERROR = 2
EXIT_INTERNAL_ERROR = 3


logger = get_logger("arksim")

# Suppress various logging and warnings
logging.getLogger("azure").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("openai").setLevel(logging.WARNING)
logging.getLogger("a2a.client.card_resolver").setLevel(logging.WARNING)

_version_tag = f"v{__version__}" if "+" not in __version__ else "main"
_EXAMPLES_REPO_URL = (
    f"https://github.com/arklexai/arksim/archive/refs/tags/{_version_tag}.tar.gz"
)
_EXAMPLES_PREFIX = "examples/"


def _enforce_thresholds(
    evaluator_output: Evaluation, evaluation_input: EvaluationInput
) -> None:
    """Run all threshold gates and exit with EXIT_EVAL_FAILED if any fail."""
    metric_ok = check_numeric_thresholds(
        evaluator_output, evaluation_input.numeric_thresholds
    )
    qual_ok = check_qualitative_failure_labels(
        evaluator_output, evaluation_input.qualitative_failure_labels
    )
    if not metric_ok or not qual_ok:
        sys.exit(EXIT_EVAL_FAILED)


def _merge_cli_overrides(yaml_settings: dict, cli_overrides: dict) -> dict:
    """Merge CLI overrides into YAML settings.
    CLI values take priority.

    Args:
        yaml_settings: Settings loaded from YAML file
        cli_overrides: Settings provided via CLI options
            (None values are ignored)

    Returns:
        Merged settings dictionary
    """
    merged = yaml_settings.copy()
    for key, value in cli_overrides.items():
        if value is not None:
            merged[key] = value
    return merged


# ============================================================================
# Argparse CLI - Dynamic argument parsing
# ============================================================================


def parse_extra_args(extra_args: list) -> dict:
    """Parse extra CLI arguments in --key value format."""
    overrides = {}
    i = 0
    while i < len(extra_args):
        arg = extra_args[i]
        if arg.startswith("--"):
            # Convert --key-name to key_name (match YAML format)
            key = arg[2:].replace("-", "_")

            # Check for --key=value format
            if "=" in key:
                key, value = key.split("=", 1)
                overrides[key] = _parse_value(value)
                i += 1
            elif i + 1 < len(extra_args) and not extra_args[i + 1].startswith("--"):
                # --key value format; single-dash values like -0.5 are treated as values
                value = extra_args[i + 1]
                overrides[key] = _parse_value(value)
                i += 2
            else:
                # Boolean flag without value (--flag means True)
                overrides[key] = True
                i += 1
        else:
            i += 1
    return overrides


def _parse_value(value: str) -> bool | int | float | str:
    """Parse a string value to appropriate type."""
    # Boolean
    if value.lower() in ("true", "yes"):
        return True
    if value.lower() in ("false", "no"):
        return False
    # Integer
    try:
        return int(value)
    except ValueError:
        pass
    # Float
    try:
        return float(value)
    except ValueError:
        pass
    # String
    return value


def _coerce_list_overrides(overrides: dict, model_cls: type) -> None:
    """Wrap scalar CLI values in lists for model fields that expect list types."""
    import types
    from typing import Union, get_args, get_origin

    for key in list(overrides):
        if key not in model_cls.model_fields:
            continue
        annotation = model_cls.model_fields[key].annotation
        # Unwrap Optional/Union (e.g. list[str] | None) to find the inner list type.
        origin = get_origin(annotation)
        if origin is Union or isinstance(annotation, types.UnionType):
            for arg in get_args(annotation):
                if get_origin(arg) is list or arg is list:
                    origin = list
                    break
        if origin is not list:
            continue
        val = overrides[key]
        # Comma-separated values are split into a list (e.g. "a.py,b.py" → ["a.py", "b.py"]).
        if isinstance(val, str):
            overrides[key] = [v.strip() for v in val.split(",")]
        elif not isinstance(val, list):
            overrides[key] = [val]


def validate_overrides(overrides: dict, valid_keys: set) -> None:
    invalid_keys = set(overrides.keys()) - valid_keys
    if invalid_keys:
        logger.error(f"Unknown options: {', '.join(sorted(invalid_keys))}")
        logger.info(f"Valid options: {', '.join(sorted(valid_keys))}")
        sys.exit(EXIT_CONFIG_ERROR)


def _log_config_summary(label: str, settings: dict) -> None:
    """Log a compact summary of resolved configuration."""
    logger.info(f"\n{label} configuration:")
    for key, value in sorted(settings.items()):
        logger.info(f"  {key}: {value}")
    logger.info("")


def _run_show_prompts(category: str | None) -> None:
    """Print evaluation prompts, optionally filtered by category."""
    from arksim.evaluator.prompt_registry import (
        get_categories,
        get_prompts_by_category,
    )

    matches = get_prompts_by_category(category)
    if not matches:
        print(
            f"Unknown category: '{category}'. Available: {', '.join(get_categories())}"
        )
        sys.exit(EXIT_CONFIG_ERROR)

    for cat in matches:
        print(f"{'=' * 60}")
        print(f"Category: {cat.category}")
        print(f"Description: {cat.description}")
        print(f"{'=' * 60}")
        for entry in cat.prompts:
            print(f"\n--- {entry.name} ---")
            print(entry.text.strip())
            print()


def _run_examples(
    name: str | None = None,
    list_only: bool = False,
) -> None:
    """Download example projects from the arksim GitHub repo."""
    import io
    import tarfile
    from urllib.request import urlopen

    logger.info("Fetching examples from arksim GitHub repo...")
    with urlopen(_EXAMPLES_REPO_URL) as resp:  # nosec B310
        data = resp.read()

    with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
        top = tar.getmembers()[0].name.split("/")[0]
        example_root = f"{top}/{_EXAMPLES_PREFIX}"

        available = sorted(
            {
                m.name.removeprefix(example_root).split("/")[0]
                for m in tar.getmembers()
                if m.name.startswith(example_root)
                and m.name != example_root.rstrip("/")
            }
        )

        if list_only:
            print("Available examples:")
            for ex in available:
                print(f"  {ex}")
            return

        if name and name not in available:
            logger.error(f"Unknown example '{name}'. Available: {', '.join(available)}")
            sys.exit(EXIT_CONFIG_ERROR)

        if name:
            filter_prefix = f"{top}/{_EXAMPLES_PREFIX}{name}/"
            dest_path = os.path.join("examples", name)
        else:
            filter_prefix = example_root
            dest_path = "examples"

        if os.path.exists(dest_path):
            logger.error(
                f"'{dest_path}' already exists. "
                "Remove it first or choose a "
                "different directory."
            )
            sys.exit(EXIT_CONFIG_ERROR)

        for member in tar.getmembers():
            if member.name.startswith(filter_prefix):
                member.name = member.name.removeprefix(f"{top}/")
                # Block path traversal (CVE-2007-4559)
                if os.path.isabs(member.name) or ".." in member.name.split("/"):
                    continue
                if sys.version_info >= (3, 12):
                    tar.extract(member, ".", filter="data")
                else:
                    tar.extract(member, ".")

    logger.info(f"Downloaded to {os.path.abspath(dest_path)}")


_INIT_AGENT_TYPES = ("custom", "chat_completions", "a2a")

_INIT_CONFIG_MAP = {
    "custom": "config.yaml",
    "chat_completions": "config_chat_completions.yaml",
    "a2a": "config_a2a.yaml",
}

_INIT_NEXT_STEPS = {
    "custom": (
        "\nNext steps:\n"
        "  1. Open my_agent.py and replace the execute() body with your agent logic\n"
        "  2. Run: arksim simulate-evaluate config.yaml"
    ),
    "chat_completions": (
        "\nNext steps:\n"
        "  1. Edit config.yaml with your agent's Chat Completions endpoint\n"
        "  2. Run: arksim simulate-evaluate config.yaml"
    ),
    "a2a": (
        "\nNext steps:\n"
        "  1. Edit config.yaml with your agent's A2A endpoint\n"
        "  2. Run: arksim simulate-evaluate config.yaml"
    ),
}


def _run_init(agent_type: str, force: bool = False) -> None:
    """Scaffold starter files for agent testing in the current directory.

    Generates a config.yaml, scenarios.json, and (for custom agent type)
    a starter my_agent.py. Files are copied from package templates.
    Existing files are not overwritten unless --force is passed.
    """
    cwd = os.getcwd()
    config_dest = os.path.join(cwd, "config.yaml")
    scenarios_dest = os.path.join(cwd, "scenarios.json")
    agent_dest = os.path.join(cwd, "my_agent.py")

    # Only check files that this agent type would write
    targets = [("config.yaml", config_dest), ("scenarios.json", scenarios_dest)]
    if agent_type == "custom":
        targets.append(("my_agent.py", agent_dest))

    if not force:
        existing = [name for name, path in targets if os.path.exists(path)]
        if existing:
            logger.error(
                f"{', '.join(existing)} already exist in the current directory. "
                "Use --force to overwrite, or run from a different directory."
            )
            sys.exit(EXIT_CONFIG_ERROR)

    templates = resources.files("arksim.templates")
    config_template = _INIT_CONFIG_MAP[agent_type]

    with resources.as_file(templates.joinpath(config_template)) as src:
        shutil.copy2(src, config_dest)
    with resources.as_file(templates.joinpath("scenarios.json")) as src:
        shutil.copy2(src, scenarios_dest)

    created = [config_dest, scenarios_dest]

    if agent_type == "custom":
        with resources.as_file(templates.joinpath("my_agent.py")) as src:
            shutil.copy2(src, agent_dest)
        created.append(agent_dest)

    for path in created:
        logger.info(f"Created {path}")

    logger.info(_INIT_NEXT_STEPS[agent_type])


# ============================================================================
# setup-claude - Install/uninstall Claude Code integration
# ============================================================================


# Project-marker files used to confirm a directory is a real project root.
_PROJECT_MARKERS = (
    ".git",
    "pyproject.toml",
    "setup.py",
    "setup.cfg",
    "package.json",
    "Cargo.toml",
    "go.mod",
)

_MANIFEST_FILENAME = ".arksim-manifest.json"
_MANIFEST_VERSION = 1


def _build_mcp_server_config(root: Path) -> dict[str, object]:
    """Build the MCP server config.

    Uses the ``arksim-mcp`` entry point installed by
    ``pip install arksim[claude]`` and pins ``cwd`` to ``root`` so the
    server's ``_project_root()`` does not depend on launch-time cwd
    inheritance from Claude Code.

    Exits with ``EXIT_CONFIG_ERROR`` if ``arksim-mcp`` is not found on PATH.
    """
    arksim_mcp = shutil.which("arksim-mcp")
    if arksim_mcp:
        return {"command": arksim_mcp, "args": [], "cwd": str(root)}

    logger.error("arksim-mcp not found. Run: pip install arksim[claude]")
    sys.exit(EXIT_CONFIG_ERROR)


def _find_integration_dir() -> Path:
    """Locate the integrations/claude_code directory.

    Checks the development layout first (relative to the arksim package),
    then falls back to importlib.resources for installed packages.
    """
    dev_path = Path(__file__).resolve().parent.parent / "integrations" / "claude_code"
    if dev_path.is_dir():
        return dev_path

    try:
        pkg_path = Path(str(resources.files("integrations.claude_code")))
        if pkg_path.is_dir():
            return pkg_path
    except (ModuleNotFoundError, TypeError):
        pass

    logger.error(
        "Could not locate integrations/claude_code directory. "
        "Ensure arksim is installed with the claude-code extra."
    )
    sys.exit(EXIT_CONFIG_ERROR)


def _sanity_check_project_dir(root: Path, force: bool) -> None:
    """Refuse install into $HOME, /, or non-project directories."""
    home = Path.home().resolve()
    resolved = root.resolve()
    if resolved == Path(resolved.anchor or "/"):
        logger.error(f"Refusing to install into filesystem root: {resolved}")
        sys.exit(EXIT_CONFIG_ERROR)
    if resolved == home:
        logger.error(
            f"Refusing to install into the user home directory: {resolved}\n"
            "Use --project-dir=PATH to point at a specific project."
        )
        sys.exit(EXIT_CONFIG_ERROR)
    if not resolved.is_dir():
        logger.error(f"Project directory does not exist: {resolved}")
        sys.exit(EXIT_CONFIG_ERROR)
    if not any((resolved / marker).exists() for marker in _PROJECT_MARKERS):
        if not force:
            markers = " / ".join(_PROJECT_MARKERS[:4])
            logger.error(
                f"Directory {resolved} does not look like a project root\n"
                f"(no {markers} found).\n"
                "Re-run with --force to install anyway."
            )
            sys.exit(EXIT_CONFIG_ERROR)
        logger.warning(
            f"Directory {resolved} has no project marker; "
            "installing because --force was set."
        )


def _atomic_write_text(path: Path, content: str) -> None:
    """Write text to ``path`` atomically via a temp file and os.replace."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_path = tempfile.mkstemp(
        dir=str(path.parent), prefix=path.name + ".", suffix=".tmp"
    )
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            fh.write(content)
        os.replace(tmp_path, path)
    except Exception:
        with contextlib.suppress(OSError):
            os.unlink(tmp_path)
        raise


def _atomic_write_json(path: Path, data: object) -> None:
    """Write JSON to ``path`` atomically."""
    _atomic_write_text(path, json.dumps(data, indent=2) + "\n")


def _read_manifest(skills_dir: Path) -> dict | None:
    """Read the install manifest if present and well-formed.

    Validates that every entry under ``skills`` is a simple basename
    (no path separators, no ``..``). Manifest values that fail this
    check are dropped with a warning so a malicious manifest cannot
    drive ``shutil.rmtree`` outside ``skills_dir``.
    """
    manifest_path = skills_dir / _MANIFEST_FILENAME
    if not manifest_path.exists():
        return None
    try:
        data = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        logger.warning(
            f"Manifest at {manifest_path} is unreadable; falling back to glob."
        )
        return None
    if not isinstance(data, dict) or data.get("version") != _MANIFEST_VERSION:
        logger.warning(
            f"Manifest at {manifest_path} has unexpected shape; falling back to glob."
        )
        return None

    raw_skills = data.get("skills", [])
    safe_skills: list[str] = []
    skipped: list[str] = []
    for entry in raw_skills if isinstance(raw_skills, list) else []:
        if not isinstance(entry, str):
            skipped.append(repr(entry))
            continue
        if not entry or "/" in entry or "\\" in entry or ".." in entry.split("/"):
            skipped.append(entry)
            continue
        if entry in (".", "..", "") or entry.startswith("."):
            skipped.append(entry)
            continue
        safe_skills.append(entry)

    if skipped:
        logger.warning(
            f"Manifest at {manifest_path} dropped suspicious entries: {skipped}"
        )

    data["skills"] = safe_skills
    return data


def _preflight_check(root: Path, *, dry_run: bool = False) -> None:
    """Print a pre-flight summary and abort if arksim-mcp is missing.

    In ``dry_run`` mode the missing-binary case is downgraded to a
    warning so users can still see the install preview before they pip
    install the [claude] extra.
    """
    arksim_mcp_path = shutil.which("arksim-mcp")
    is_git_repo = (root / ".git").is_dir()
    keys_present = [
        name
        for name in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GOOGLE_API_KEY")
        if os.environ.get(name)
    ]
    logger.info("Pre-flight check:")
    logger.info(
        f"  arksim-mcp on PATH:    "
        f"{'yes (' + arksim_mcp_path + ')' if arksim_mcp_path else 'NO'}"
    )
    logger.info(f"  Project is git repo:   {'yes' if is_git_repo else 'no'}")
    if keys_present:
        logger.info(f"  LLM API key set:       {', '.join(keys_present)}")
    else:
        logger.warning(
            "  LLM API key set:       NONE "
            "(set OPENAI_API_KEY, ANTHROPIC_API_KEY, or GOOGLE_API_KEY)"
        )
    if not arksim_mcp_path:
        msg = "arksim-mcp not found. Run: pip install arksim[claude]"
        if dry_run:
            logger.warning(msg + " (dry-run: continuing for preview)")
        else:
            logger.error(msg)
            sys.exit(EXIT_CONFIG_ERROR)


def _run_setup_claude(
    project_dir: str = ".",
    force: bool = False,
    uninstall: bool = False,
    dry_run: bool = False,
) -> None:
    """Install or uninstall the Claude Code integration for a project.

    Install mode (default):
      - Verifies the target directory looks like a project root.
      - Runs a pre-flight check on environment + tooling.
      - Stages skill directories and atomic-swaps them into place.
      - Writes a manifest of installed paths so uninstall is precise.
      - Writes ``.mcp.json`` atomically.

    Uninstall mode:
      - Reads the manifest and removes only the entries it lists.
      - Falls back to a glob with a warning if the manifest is missing.

    Both modes accept ``dry_run=True`` to print what would change
    without touching the filesystem.
    """
    root = Path(project_dir).expanduser().resolve()
    claude_dir = root / ".claude"
    mcp_config_path = root / ".mcp.json"
    skills_dir = claude_dir / "skills"

    # Sanity-check the project dir for both install and uninstall.
    # Uninstall reads .arksim-manifest.json from skills_dir and
    # deletes paths it lists; running it against $HOME or / would
    # walk into arbitrary deletion via a malicious manifest.
    _sanity_check_project_dir(root, force)

    if uninstall:
        _uninstall_claude(mcp_config_path, skills_dir, dry_run=dry_run)
        return

    _preflight_check(root, dry_run=dry_run)
    integration_dir = _find_integration_dir()
    _install_claude(
        integration_dir,
        claude_dir,
        mcp_config_path,
        skills_dir,
        force,
        dry_run=dry_run,
    )


def _uninstall_claude(
    mcp_config_path: Path,
    skills_dir: Path,
    *,
    dry_run: bool = False,
) -> None:
    """Remove arksim MCP config and skills from a project.

    Uses the install manifest to know exactly what to remove. Falls
    back to a glob over ``arksim-*`` directories with a warning if the
    manifest is missing.
    """
    manifest = _read_manifest(skills_dir)
    skills_to_remove: list[Path] = []
    if manifest is not None:
        skills_root = skills_dir.resolve()
        for rel in manifest.get("skills", []):
            if not isinstance(rel, str) or not rel:
                continue
            candidate = skills_dir / rel
            try:
                resolved_candidate = candidate.resolve()
            except (OSError, RuntimeError):
                continue
            try:
                resolved_candidate.relative_to(skills_root)
            except ValueError:
                logger.warning(
                    f"Refusing to remove {resolved_candidate}: escapes skills directory"
                )
                continue
            if candidate.is_symlink():
                logger.warning(f"Refusing to remove {candidate}: is a symlink")
                continue
            if candidate.is_dir():
                skills_to_remove.append(candidate)
    elif skills_dir.is_dir():
        for skill_dir in sorted(skills_dir.glob("arksim-*")):
            if skill_dir.is_dir():
                skills_to_remove.append(skill_dir)

    if dry_run:
        logger.info("Dry run. Would remove:")
        for p in skills_to_remove:
            logger.info(f"  Skill dir: {p}")
        if mcp_config_path.exists():
            logger.info(f"  arksim entry from {mcp_config_path}")
        manifest_path = skills_dir / _MANIFEST_FILENAME
        if manifest_path.exists():
            logger.info(f"  Manifest:  {manifest_path}")
        return

    if mcp_config_path.exists():
        try:
            mcp_config = json.loads(mcp_config_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            logger.error(
                f"Invalid JSON in {mcp_config_path}. Fix or delete the file and retry."
            )
            sys.exit(EXIT_CONFIG_ERROR)
        mcp_servers = mcp_config.get("mcpServers")
        if not isinstance(mcp_servers, dict):
            mcp_servers = {}
        mcp_servers.pop("arksim", None)
        if mcp_servers:
            mcp_config["mcpServers"] = mcp_servers
            _atomic_write_json(mcp_config_path, mcp_config)
        else:
            mcp_config_path.unlink()

    for skill_dir in skills_to_remove:
        shutil.rmtree(skill_dir)

    manifest_path = skills_dir / _MANIFEST_FILENAME
    if manifest_path.exists():
        manifest_path.unlink()

    logger.info("Removed arksim Claude Code integration.")


def _install_claude(
    integration_dir: Path,
    claude_dir: Path,
    mcp_config_path: Path,
    skills_dir: Path,
    force: bool,
    *,
    dry_run: bool = False,
) -> None:
    """Install arksim MCP config and skills into a project (atomic)."""
    skills_src = integration_dir / "skills"
    if not skills_src.is_dir():
        logger.error(f"Skills directory not found: {skills_src}")
        sys.exit(EXIT_CONFIG_ERROR)

    candidate_skills: list[Path] = []
    for skill_dir in sorted(skills_src.iterdir()):
        if not skill_dir.is_dir() or not skill_dir.name.startswith("arksim-"):
            continue
        if not (skill_dir / "SKILL.md").is_file():
            continue
        candidate_skills.append(skill_dir)

    if not candidate_skills:
        logger.error(
            f"No skill directories found in {skills_src}. "
            "Installation may be incomplete."
        )
        sys.exit(EXIT_CONFIG_ERROR)

    existing_skills = list(skills_dir.glob("arksim-*")) if skills_dir.is_dir() else []
    if existing_skills and not force:
        logger.error(
            "arksim skills already installed. "
            "Use --force to overwrite, or --uninstall to remove first."
        )
        sys.exit(EXIT_CONFIG_ERROR)

    # mcp_config_path lives at the project root.
    server_config = _build_mcp_server_config(mcp_config_path.parent)

    if dry_run:
        logger.info("Dry run. Would install:")
        for skill_dir in candidate_skills:
            logger.info(f"  Skill:    {skills_dir / skill_dir.name}/SKILL.md")
        logger.info(f"  Manifest: {skills_dir / _MANIFEST_FILENAME}")
        logger.info(f"  MCP config: {mcp_config_path} (mcpServers.arksim)")
        return

    claude_dir.mkdir(parents=True, exist_ok=True)
    skills_dir.mkdir(parents=True, exist_ok=True)

    # Atomic skill swap. Rename existing arksim-* directories aside,
    # then move staged skills into place, then write the manifest and
    # .mcp.json, then rmtree the renamed-aside copy. On exception we
    # restore the renamed-aside copy so the user is never left without
    # a working skill set.
    pid = os.getpid()
    staging_dir = skills_dir.parent / f".skills.staging.{pid}"
    aside_dir = skills_dir.parent / f".skills.aside.{pid}"
    if staging_dir.exists():
        shutil.rmtree(staging_dir)
    if aside_dir.exists():
        shutil.rmtree(aside_dir)
    staging_dir.mkdir(parents=True)
    aside_dir.mkdir(parents=True)

    copied: list[str] = []

    def _restore_from_aside() -> None:
        """Move any aside-d skill directories back into place."""
        for entry in aside_dir.iterdir():
            target = skills_dir / entry.name
            if target.exists():
                shutil.rmtree(target)
            shutil.move(str(entry), str(target))

    try:
        for skill_dir in candidate_skills:
            dest_dir = staging_dir / skill_dir.name
            shutil.copytree(skill_dir, dest_dir)
            copied.append(skill_dir.name)

        for old_skill in skills_dir.glob("arksim-*"):
            if old_skill.is_symlink():
                logger.warning(f"Skipping symlinked entry during install: {old_skill}")
                continue
            if old_skill.is_dir():
                shutil.move(str(old_skill), str(aside_dir / old_skill.name))

        try:
            for staged in staging_dir.iterdir():
                target = skills_dir / staged.name
                shutil.move(str(staged), str(target))
        except Exception:
            _restore_from_aside()
            raise
    finally:
        if staging_dir.exists():
            shutil.rmtree(staging_dir, ignore_errors=True)
        if aside_dir.exists():
            shutil.rmtree(aside_dir, ignore_errors=True)

    manifest = {
        "version": _MANIFEST_VERSION,
        "skills": copied,
        "mcp_server_key": "arksim",
    }
    _atomic_write_json(skills_dir / _MANIFEST_FILENAME, manifest)

    mcp_config: dict = {}
    if mcp_config_path.exists():
        try:
            mcp_config = json.loads(mcp_config_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            logger.error(
                f"Invalid JSON in {mcp_config_path}. Fix or delete the file and retry."
            )
            sys.exit(EXIT_CONFIG_ERROR)

    mcp_servers = mcp_config.setdefault("mcpServers", {})
    if not isinstance(mcp_servers, dict):
        mcp_servers = {}
        mcp_config["mcpServers"] = mcp_servers
    mcp_servers["arksim"] = server_config
    _atomic_write_json(mcp_config_path, mcp_config)

    logger.info("Installed arksim Claude Code integration:")
    logger.info(f"  MCP config: {mcp_config_path}")
    for name in copied:
        logger.info(f"  Skill:    {skills_dir / name}/SKILL.md")
    logger.info(f"  Manifest: {skills_dir / _MANIFEST_FILENAME}")
    logger.info('\nStart Claude Code and type "/arksim-test" to begin.')


def _add_config_subparser(
    subparsers: argparse._SubParsersAction,
    name: str,
    help_text: str,
) -> argparse.ArgumentParser:
    """Add a subparser for config-based commands.

    These commands accept a YAML config file and
    arbitrary --key value overrides.
    """
    sp = subparsers.add_parser(name, help=help_text)
    sp.add_argument(
        "config_file",
        type=str,
        nargs="?",
        default=None,
        help="Path to YAML config file",
    )
    sp.add_argument(
        "additional_args",
        nargs=argparse.REMAINDER,
        help="Additional --key value overrides",
    )
    return sp


def build_parser() -> argparse.ArgumentParser:
    """Build the CLI argument parser with subcommands."""
    parser = argparse.ArgumentParser(
        prog="arksim",
        description="⛵️ ArkSim - Know how your agent performs before it goes live.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\

            Examples:
              arksim --version
              arksim simulate config.yaml
              arksim evaluate config.yaml
              arksim simulate-evaluate config.yaml
              arksim show-prompts --category agent_behavior_failure
              arksim examples --list
              arksim examples bank-insurance
              arksim ui --port 9090
              arksim setup-claude
        """),
    )

    parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
        help="Show version and exit",
    )

    sub = parser.add_subparsers(dest="command")

    # Config-based commands (YAML + --key value overrides)
    _add_config_subparser(
        sub,
        "simulate",
        "Run agent simulations",
    )
    _add_config_subparser(
        sub,
        "evaluate",
        "Evaluate simulation results",
    )
    _add_config_subparser(
        sub,
        "simulate-evaluate",
        "Simulate then evaluate",
    )

    # show-prompts
    sp_prompts = sub.add_parser(
        "show-prompts",
        help="Display evaluation prompts",
    )
    sp_prompts.add_argument(
        "--category",
        type=str,
        default=None,
        help="Filter prompts by category",
    )

    # examples
    sp_examples = sub.add_parser(
        "examples",
        help="Download example projects from GitHub",
    )
    sp_examples.add_argument(
        "name",
        nargs="?",
        default=None,
        help="Example to download (e.g. bank-insurance)",
    )
    sp_examples.add_argument(
        "--list",
        action="store_true",
        dest="list_only",
        help="List available examples",
    )

    # init
    sp_init = sub.add_parser(
        "init",
        help="Scaffold starter files for agent testing",
    )
    sp_init.add_argument(
        "--agent-type",
        type=str,
        choices=_INIT_AGENT_TYPES,
        default="custom",
        dest="agent_type",
        help=(
            "Agent connection type (default: custom). "
            "'custom' generates a Python agent file (no server needed). "
            "'chat_completions' for HTTP endpoints. "
            "'a2a' for Agent-to-Agent protocol."
        ),
    )
    sp_init.add_argument(
        "--force",
        action="store_true",
        default=False,
        help="Overwrite existing files",
    )

    # ui
    sp_ui = sub.add_parser(
        "ui",
        help="Launch web UI control plane",
    )
    sp_ui.add_argument(
        "--port",
        type=int,
        default=8080,
        help="Port to serve on (default: 8080)",
    )

    # setup-claude
    sp_setup_claude = sub.add_parser(
        "setup-claude",
        help="Install Claude Code integration (MCP server + skills)",
    )
    sp_setup_claude.add_argument(
        "--project-dir",
        type=str,
        default=".",
        dest="project_dir",
        help="Project root directory (default: current directory)",
    )
    sp_setup_claude.add_argument(
        "--force",
        action="store_true",
        default=False,
        help="Overwrite existing skills",
    )
    sp_setup_claude.add_argument(
        "--uninstall",
        action="store_true",
        default=False,
        help="Remove arksim Claude Code integration",
    )
    sp_setup_claude.add_argument(
        "--dry-run",
        action="store_true",
        default=False,
        dest="dry_run",
        help="Preview changes without writing the filesystem",
    )

    return parser


def _cmd_simulate(
    settings: dict,
    overrides: dict,
    config_path: str | None,
    cli_overrides: set,
    verbose: bool,
) -> None:
    """Validate config and run simulation."""
    try:
        valid_keys = set(SimulationInput.model_fields.keys())
        validate_overrides(overrides, valid_keys)
        _coerce_list_overrides(overrides, SimulationInput)
        settings = _merge_cli_overrides(settings, overrides)
        simulation_input = SimulationInput.model_validate(
            settings,
            context={"config_path": config_path, "cli_overrides": cli_overrides},
        )
        _log_config_summary("Simulation", simulation_input.model_dump())
    except ValidationError as e:
        logger.error(f"Configuration error: {e}")
        logger.debug("Traceback:", exc_info=True)
        sys.exit(EXIT_CONFIG_ERROR)

    asyncio.run(run_simulation(simulation_input, verbose=verbose))


def _cmd_evaluate(
    settings: dict,
    overrides: dict,
    config_path: str | None,
    cli_overrides: set,
) -> None:
    """Validate config and run evaluation."""
    try:
        valid_keys = set(EvaluationInput.model_fields.keys())
        validate_overrides(overrides, valid_keys)
        _coerce_list_overrides(overrides, EvaluationInput)
        settings = _merge_cli_overrides(settings, overrides)
        evaluation_input = EvaluationInput.model_validate(
            settings,
            context={"config_path": config_path, "cli_overrides": cli_overrides},
        )
        if not evaluation_input.simulation_file_path:
            logger.error("simulation_file_path is required.")
            sys.exit(EXIT_CONFIG_ERROR)
        if not os.path.isfile(evaluation_input.simulation_file_path):
            logger.error(
                f"simulation_file_path does not exist: "
                f"{evaluation_input.simulation_file_path}"
            )
            sys.exit(EXIT_CONFIG_ERROR)
        _log_config_summary("Evaluation", evaluation_input.model_dump())
    except ValidationError as e:
        logger.error(f"Configuration error: {e}")
        logger.debug("Traceback:", exc_info=True)
        sys.exit(EXIT_CONFIG_ERROR)

    evaluator_output = run_evaluation(evaluation_input)
    _enforce_thresholds(evaluator_output, evaluation_input)


def _cmd_simulate_evaluate(
    settings: dict,
    overrides: dict,
    config_path: str | None,
    cli_overrides: set,
    verbose: bool,
) -> None:
    """Validate config, run simulation, then run evaluation."""
    try:
        valid_keys = set(SimulationInput.model_fields.keys()) | set(
            EvaluationInput.model_fields.keys()
        )
        validate_overrides(overrides, valid_keys)
        _coerce_list_overrides(overrides, SimulationInput)
        _coerce_list_overrides(overrides, EvaluationInput)
        settings = _merge_cli_overrides(settings, overrides)

        simulation_settings = {
            k: v for k, v in settings.items() if k in SimulationInput.model_fields
        }
        simulation_input = SimulationInput.model_validate(
            simulation_settings,
            context={"config_path": config_path, "cli_overrides": cli_overrides},
        )
        _log_config_summary("Simulation", simulation_input.model_dump())

        evaluation_settings = {
            k: v for k, v in settings.items() if k in EvaluationInput.model_fields
        }
        evaluation_input = EvaluationInput.model_validate(
            evaluation_settings,
            context={"config_path": config_path, "cli_overrides": cli_overrides},
        )
        _log_config_summary("Evaluation", evaluation_input.model_dump())
    except ValidationError as e:
        logger.error(f"Configuration error: {e}")
        logger.debug("Traceback:", exc_info=True)
        sys.exit(EXIT_CONFIG_ERROR)

    sim_start = time.time()
    simulation_output = asyncio.run(run_simulation(simulation_input, verbose=verbose))
    logger.info(f"Simulation completed in {time.time() - sim_start:.2f} seconds")

    eval_start = time.time()
    evaluator_output = run_evaluation(evaluation_input, simulation=simulation_output)
    logger.info(f"Evaluation completed in {time.time() - eval_start:.2f} seconds")

    _enforce_thresholds(evaluator_output, evaluation_input)


def main() -> None:
    """Main entry point for the arksim CLI."""
    s_time = time.time()
    parser = build_parser()
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(EXIT_CONFIG_ERROR)

    if args.command == "show-prompts":
        _run_show_prompts(args.category)
        return

    if args.command == "examples":
        _run_examples(name=args.name, list_only=args.list_only)
        return

    if args.command == "init":
        _run_init(agent_type=args.agent_type, force=args.force)
        return

    if args.command == "ui":
        from arksim.ui.app import launch_ui

        launch_ui(port=args.port)
        return

    if args.command == "setup-claude":
        _run_setup_claude(
            project_dir=args.project_dir,
            force=args.force,
            uninstall=args.uninstall,
            dry_run=args.dry_run,
        )
        return

    use_config_file = (
        args.config_file
        and os.path.exists(args.config_file)
        and args.config_file.endswith((".yaml", ".yml"))
    )

    overrides = parse_extra_args(args.additional_args)

    # Load settings from YAML file if valid, otherwise use empty dict
    settings = {}
    if use_config_file:
        try:
            with open(args.config_file) as f:
                settings = yaml.safe_load(f) or {}
        except Exception as e:
            logger.error(
                f"Could not load config file '{args.config_file}': {e}",
            )
            sys.exit(EXIT_CONFIG_ERROR)
    else:
        logger.warning("No config YAML file provided.")

    # override with the environment variables
    if os.getenv("LLM_PROVIDER"):
        settings["provider"] = os.getenv("LLM_PROVIDER")

    # Resolve log level: env var takes priority over YAML key
    log_level = os.getenv("LOG_LEVEL")
    if log_level:
        logging.getLogger("arksim").setLevel(log_level.upper())

    # Extract verbose flag before building model inputs
    verbose = overrides.pop("verbose", False)

    config_path = os.path.abspath(args.config_file) if use_config_file else None
    cli_overrides = set(overrides.keys())

    try:
        if args.command == "simulate":
            _cmd_simulate(settings, overrides, config_path, cli_overrides, verbose)
        elif args.command == "evaluate":
            _cmd_evaluate(settings, overrides, config_path, cli_overrides)
        elif args.command == "simulate-evaluate":
            _cmd_simulate_evaluate(
                settings, overrides, config_path, cli_overrides, verbose
            )
    except (KeyboardInterrupt, SystemExit):
        raise  # this preserves exit codes for CI
    except Exception as e:
        logger.error(f"Internal error: {e}")
        logger.debug("Traceback:", exc_info=True)
        sys.exit(EXIT_INTERNAL_ERROR)

    logger.info(f"Total elapsed: {time.time() - s_time:.2f} seconds")


if __name__ == "__main__":
    main()
