#!/usr/bin/env python3
"""
BAT Runner - Bot acceptance testing for ha-mcp.

Executes MCP test scenarios on real AI agent CLIs (Claude, Gemini, OpenAI-compatible)
against a Home Assistant test instance. The calling agent generates scenarios dynamically
and evaluates results - this script is a dumb executor.

Full results are written to a temp file. Stdout gets a concise summary with
the file path — the calling agent only reads the full file when needed.

Usage:
    echo '{"test_prompt":"Search for light entities."}' | uv run python tests/uat/run_uat.py --agents gemini
    uv run python tests/uat/run_uat.py --scenario-file /tmp/scenario.json --agents claude,gemini
    uv run python tests/uat/run_uat.py --ha-url http://localhost:8123 --ha-token TOKEN --agents gemini
"""

from __future__ import annotations

import argparse
import asyncio
import difflib
import json
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from collections.abc import Callable
from pathlib import Path
from typing import NoReturn

import requests
from testcontainers.core.container import DockerContainer

# Resolve paths relative to repo root
SCRIPT_DIR = Path(__file__).resolve().parent
TESTS_DIR = SCRIPT_DIR.parent
REPO_ROOT = TESTS_DIR.parent

sys.path.insert(0, str(TESTS_DIR))
from test_constants import HA_TEST_IMAGE, TEST_TOKEN  # noqa: E402
from uat._logging import configure_cli_logging  # noqa: E402
from uat.ha_wait import wait_for_ha_ready  # noqa: E402

HA_IMAGE = HA_TEST_IMAGE

DEFAULT_TIMEOUT = 300
DEFAULT_AGENTS = "claude,gemini"


logger = logging.getLogger("uat.run_uat")


class SuggestingArgumentParser(argparse.ArgumentParser):
    """argparse parser that suggests close matches for unknown flags."""

    def error(self, message: str) -> NoReturn:
        match = re.search(r"unrecognized arguments?: (--\S+)", message)
        if match:
            known = [opt for opt in self._option_string_actions if opt.startswith("--")]
            suggestions = difflib.get_close_matches(match.group(1), known, n=1, cutoff=0.6)
            if suggestions:
                message = f"{message} (did you mean {suggestions[0]}?)"
        super().error(message)


# ---------------------------------------------------------------------------
# HA Container Management
# ---------------------------------------------------------------------------
def setup_config_directory() -> Path:
    """Copy initial_test_state to a temp dir for the HA container."""
    config_dir = Path(tempfile.mkdtemp(prefix="ha_bat_"))
    initial_state = TESTS_DIR / "initial_test_state"
    if not initial_state.exists():
        raise FileNotFoundError(f"initial_test_state not found at {initial_state}")

    for item in initial_state.iterdir():
        if item.is_file():
            shutil.copy2(item, config_dir)
        elif item.is_dir():
            shutil.copytree(item, config_dir / item.name)

    # Set permissions
    os.chmod(config_dir, 0o755)
    for item in config_dir.rglob("*"):
        if item.is_file():
            os.chmod(item, 0o644)
        elif item.is_dir():
            os.chmod(item, 0o755)

    return config_dir


class HAContainer:
    """Context manager for a disposable HA test container."""

    def __init__(self) -> None:
        self.container: DockerContainer | None = None
        self.config_dir: Path | None = None
        self.url: str = ""
        self.token: str = TEST_TOKEN

    def __enter__(self) -> HAContainer:
        self.config_dir = setup_config_directory()
        self.container = (
            DockerContainer(HA_IMAGE)
            .with_exposed_ports(8123)
            .with_volume_mapping(str(self.config_dir), "/config", "rw")
            .with_env("TZ", "UTC")
            .with_kwargs(privileged=True)
        )
        self.container.start()
        try:
            port = self.container.get_exposed_port(8123)
            self.url = f"http://localhost:{port}"
            logger.info(f"HA container started on {self.url}")
            wait_for_ha_ready(self.url, self.token)
        except Exception:
            self.__exit__(None, None, None)
            raise
        return self

    def __exit__(self, *exc: object) -> None:
        if self.container:
            logger.info("Stopping HA container...")
            self.container.stop()
        if self.config_dir and self.config_dir.exists():
            shutil.rmtree(self.config_dir, ignore_errors=True)


# ---------------------------------------------------------------------------
# MCP Config Generation
# ---------------------------------------------------------------------------
def mcp_server_command(branch: str | None) -> list[str]:
    """Build the MCP server command for stdio mode."""
    if branch:
        return [
            "uvx",
            "--from",
            f"git+https://github.com/homeassistant-ai/ha-mcp.git@{branch}",
            "ha-mcp",
        ]
    return ["uv", "run", "--project", str(REPO_ROOT), "ha-mcp"]


def preflight_check_docker(timeout: float = 5.0) -> str | None:
    """Return an error string if the Docker daemon is unreachable, else None."""
    try:
        result = subprocess.run(
            ["docker", "info"],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except FileNotFoundError:
        return "'docker' CLI not found on PATH (install Docker or pass --ha-url)"
    except subprocess.TimeoutExpired:
        return f"Docker daemon did not respond within {timeout:.0f}s"
    if result.returncode != 0:
        stderr = (result.stderr or result.stdout).strip().splitlines()
        hint = stderr[-1] if stderr else f"exit {result.returncode}"
        return f"Docker daemon is not reachable: {hint}"
    return None


def preflight_check_base_url(base_url: str, timeout: float = 5.0) -> str | None:
    """Return an error string if the OpenAI endpoint is unreachable or broken, else None.

    Catches both connection-level failures (ConnectionError, Timeout) and
    HTTP error responses (4xx/5xx from ``raise_for_status``) — the latter
    covers "up but wrong" cases like hitting the wrong port, a bad
    auth token, or an upstream error, which would otherwise only surface
    as an opaque warmup stall.
    """
    url = f"{base_url.rstrip('/')}/models"
    try:
        resp = requests.get(url, timeout=timeout)
        resp.raise_for_status()
    except requests.RequestException as e:
        return f"OpenAI endpoint {base_url} is not reachable ({type(e).__name__}): {e}"
    return None


def _build_mcp_env(
    ha_url: str, ha_token: str, extra_env: dict[str, str] | None
) -> dict[str, str]:
    # Override with --mcp-env LOG_LEVEL=INFO when debugging the server.
    env = {
        "HOMEASSISTANT_URL": ha_url,
        "HOMEASSISTANT_TOKEN": ha_token,
        "LOG_LEVEL": "WARNING",
    }
    if extra_env:
        env.update(extra_env)
    return env


def parse_mcp_env(
    raw_mcp_env: list[str] | None,
    base_url: str | None = None,
    on_default_applied: Callable[[str], None] | None = None,
) -> dict[str, str]:
    """Parse ``KEY=VALUE`` pairs into a dict, with the local-model default.

    A bare ``KEY`` (no ``=``) maps to an empty string — useful for
    boolean-presence flags. When ``base_url`` is set (targeting a local
    OpenAI-compatible endpoint), ``ENABLE_TOOL_SEARCH=true`` is injected
    unless the caller already set it. Local models typically can't prefill
    the full tool catalog, so tool search is the default for that path.
    Override with an explicit ``ENABLE_TOOL_SEARCH=...`` pair.

    If ``on_default_applied`` is provided, it is called with a human-readable
    message each time a default is injected.
    """
    pairs = raw_mcp_env or []
    env = {k: v for pair in pairs for k, _, v in [pair.partition("=")]}
    if base_url and "ENABLE_TOOL_SEARCH" not in env:
        env["ENABLE_TOOL_SEARCH"] = "true"
        if on_default_applied:
            on_default_applied(
                "Defaulting ENABLE_TOOL_SEARCH=true for local model (--base-url set)"
            )
    return env


def build_stdio_mcp_config(
    ha_url: str,
    ha_token: str,
    branch: str | None,
    extra_env: dict[str, str] | None = None,
) -> dict:
    """Build the stdio MCP config dict (format shared with Claude's --mcp-config).

    The dict can be passed directly to ``fastmcp.Client(config)`` for in-process
    use, or serialized to a file via ``write_stdio_mcp_config`` for CLI agents
    that expect a file path.
    """
    cmd = mcp_server_command(branch)
    env = _build_mcp_env(ha_url, ha_token, extra_env)
    return {
        "mcpServers": {
            "home-assistant": {
                "command": cmd[0],
                "args": cmd[1:],
                "env": env,
            }
        }
    }


def write_stdio_mcp_config(
    ha_url: str,
    ha_token: str,
    branch: str | None,
    extra_env: dict[str, str] | None = None,
) -> Path:
    """Write the stdio MCP config to a temporary JSON file, return its path."""
    config = build_stdio_mcp_config(ha_url, ha_token, branch, extra_env)
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", prefix="claude_mcp_", delete=False
    ) as f:
        json.dump(config, f)
    return Path(f.name)


def write_gemini_mcp_config(
    ha_url: str,
    ha_token: str,
    branch: str | None,
    workdir: Path,
    extra_env: dict[str, str] | None = None,
) -> None:
    """Write .gemini/settings.json in the given workdir."""
    cmd = mcp_server_command(branch)
    gemini_dir = workdir / ".gemini"
    gemini_dir.mkdir(exist_ok=True)
    env = _build_mcp_env(ha_url, ha_token, extra_env)
    config = {
        "mcpServers": {
            "homeassistant": {
                "command": cmd[0],
                "args": cmd[1:],
                "env": env,
            }
        }
    }
    (gemini_dir / "settings.json").write_text(json.dumps(config))


# ---------------------------------------------------------------------------
# Agent Execution
# ---------------------------------------------------------------------------
def check_agent_available(name: str) -> bool:
    """Check if an agent CLI is installed."""
    if name == "openai":
        return True  # Local script, always available
    return shutil.which(name) is not None


async def run_cli(cmd: list[str], timeout: int, cwd: Path | None = None) -> dict:
    """Run a CLI command and capture output."""
    # Strip CLAUDECODE env var to allow nested Claude CLI sessions
    env = {k: v for k, v in os.environ.items() if k != "CLAUDECODE"}

    start = time.monotonic()
    # Default for the response schema; both return paths re-assign before
    # using the value, but the explicit init survives a future exit-path
    # addition that would otherwise omit the field.
    duration_ms = 0
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(cwd) if cwd else None,
            env=env,
        )
        stdout_bytes, stderr_bytes = await asyncio.wait_for(
            proc.communicate(), timeout=timeout
        )
        duration_ms = int((time.monotonic() - start) * 1000)
        stdout_text = stdout_bytes.decode("utf-8", errors="replace")
        stderr_text = stderr_bytes.decode("utf-8", errors="replace")

        # Try to parse JSON output
        raw_json = None
        try:
            raw_json = json.loads(stdout_text)
        except json.JSONDecodeError:
            pass

        # Extract fields from JSON if available
        output_text = stdout_text
        num_turns = None
        tool_stats = None
        session_id = None
        cost_usd = None
        tokens_input = None
        tokens_output = None
        tokens_first_input = None
        if raw_json and isinstance(raw_json, dict):
            # Claude JSON format
            if "result" in raw_json:
                output_text = raw_json.get("result", stdout_text)
            # Gemini JSON format
            if "response" in raw_json:
                output_text = raw_json.get("response", stdout_text)
            num_turns = raw_json.get("num_turns")
            tool_stats = raw_json.get("tool_stats")
            session_id = raw_json.get("session_id")
            cost_usd = raw_json.get("total_cost_usd") or raw_json.get("cost_usd")
            # Gemini stats
            if "stats" in raw_json and isinstance(raw_json["stats"], dict):
                tool_stats = raw_json["stats"].get("tools")
            # OpenAI agent token counts and first-input baseline (included directly in JSON output)
            tokens_input = raw_json.get("tokens_input")
            tokens_output = raw_json.get("tokens_output")
            tokens_first_input = raw_json.get("tokens_first_input")

        result: dict = {
            "completed": proc.returncode == 0,
            "output": output_text,
            "duration_ms": duration_ms,
            "exit_code": proc.returncode,
            "stderr": stderr_text,
        }
        if num_turns is not None:
            result["num_turns"] = num_turns
        if tool_stats is not None:
            result["tool_stats"] = tool_stats
        if session_id is not None:
            result["session_id"] = session_id
        if cost_usd is not None:
            result["cost_usd"] = cost_usd
        if tokens_input is not None:
            result["tokens_input"] = tokens_input
        if tokens_output is not None:
            result["tokens_output"] = tokens_output
        if tokens_first_input is not None:
            result["tokens_first_input"] = tokens_first_input
        if raw_json is not None:
            result["raw_json"] = raw_json
        return result
    except TimeoutError:
        # Capture the timeout-firing instant before the cleanup work so the
        # reported duration reflects "wall-clock until the command timed out"
        # rather than "wall-clock including the orphan-process teardown".
        duration_ms = int((time.monotonic() - start) * 1000)
        # Terminate the orphaned process
        try:
            proc.terminate()
            await asyncio.wait_for(proc.wait(), timeout=5)
        except (TimeoutError, ProcessLookupError):
            proc.kill()
        return {
            "completed": False,
            "output": "",
            "duration_ms": duration_ms,
            "exit_code": -1,
            "stderr": f"Timed out after {timeout}s",
        }


def build_claude_cmd(
    prompt: str, mcp_config_path: Path, model: str = "sonnet"
) -> list[str]:
    return [
        "claude",
        "-p",
        prompt,
        "--mcp-config",
        str(mcp_config_path),
        "--strict-mcp-config",
        "--allowedTools",
        "mcp__home-assistant",
        "--output-format",
        "json",
        "--permission-mode",
        "bypassPermissions",
        "--model",
        model,
    ]


def build_gemini_cmd(prompt: str) -> list[str]:
    return [
        "gemini",
        "-p",
        prompt,
        "--approval-mode",
        "yolo",
        "--allowed-mcp-server-names",
        "homeassistant",
        "-o",
        "json",
    ]


def build_openai_cmd(
    prompt: str,
    mcp_config_path: Path,
    base_url: str,
    model: str | None = None,
    api_key: str = "no-key",
    max_tools: int | None = None,
    no_think: bool = False,
    max_tokens: int | None = None,
) -> list[str]:
    cmd = [
        sys.executable,
        str(SCRIPT_DIR / "openai_agent.py"),
        "--prompt",
        prompt,
        "--mcp-config",
        str(mcp_config_path),
        "--base-url",
        base_url,
        "--api-key",
        api_key,
    ]
    if model:
        cmd.extend(["--model", model])
    if max_tools is not None:
        cmd.extend(["--max-tools", str(max_tools)])
    if no_think:
        cmd.append("--no-think")
    if max_tokens is not None:
        cmd.extend(["--max-tokens", str(max_tokens)])
    return cmd


async def run_agent_scenario(
    agent_name: str,
    scenario: dict,
    ha_url: str,
    ha_token: str,
    branch: str | None,
    timeout: int,
    model: str | None = None,
    base_url: str | None = None,
    api_key: str = "no-key",
    max_tools: int | None = None,
    no_think: bool = False,
    max_tokens: int | None = None,
    extra_env: dict[str, str] | None = None,
) -> dict:
    """Run a full scenario (setup/test/teardown) for one agent."""
    results: dict = {"available": True}

    # Prepare MCP config
    stdio_config_path: Path | None = None
    gemini_workdir: Path | None = None

    if agent_name in ("claude", "openai"):
        stdio_config_path = write_stdio_mcp_config(ha_url, ha_token, branch, extra_env)
    elif agent_name == "gemini":
        gemini_workdir = Path(tempfile.mkdtemp(prefix="gemini_bat_"))
        write_gemini_mcp_config(ha_url, ha_token, branch, gemini_workdir, extra_env)

    try:
        for phase in ("setup_prompt", "test_prompt", "teardown_prompt"):
            prompt = scenario.get(phase)
            if not prompt:
                continue

            phase_key = phase.replace("_prompt", "")
            logger.info(f"  [{agent_name}] Running {phase_key}...")

            if agent_name == "claude":
                assert stdio_config_path is not None
                cmd = build_claude_cmd(
                    prompt, stdio_config_path, model=model or "sonnet"
                )
                result = await run_cli(cmd, timeout)
            elif agent_name == "gemini":
                cmd = build_gemini_cmd(prompt)
                result = await run_cli(cmd, timeout, cwd=gemini_workdir)
            elif agent_name == "openai":
                assert stdio_config_path is not None
                assert base_url is not None  # validated in run()
                cmd = build_openai_cmd(
                    prompt,
                    stdio_config_path,
                    base_url=base_url,
                    model=model,
                    api_key=api_key,
                    max_tools=max_tools,
                    no_think=no_think,
                    max_tokens=max_tokens,
                )
                result = await run_cli(cmd, timeout)
            else:
                result = {
                    "completed": False,
                    "output": f"Unknown agent: {agent_name}",
                    "duration_ms": 0,
                    "exit_code": -1,
                    "stderr": "",
                }

            results[phase_key] = result
            logger.info(
                f"  [{agent_name}] {phase_key} completed (exit={result['exit_code']}, {result['duration_ms']}ms)"
            )
            # Forward agent stderr on failure so the error is visible to the user
            if result["exit_code"] != 0 and result.get("stderr"):
                _BOX_CHARS = frozenset("│╭╰╮─▄█▀ \t")
                for line in result["stderr"].splitlines():
                    if "error" in line.lower():
                        logger.info(f"  [{agent_name}] !! {line.strip()}")
                    elif not all(c in _BOX_CHARS for c in line):
                        logger.info(f"  [{agent_name}] stderr: {line}")
    finally:
        # Cleanup temp files
        if stdio_config_path and stdio_config_path.exists():
            stdio_config_path.unlink()
        if gemini_workdir and gemini_workdir.exists():
            shutil.rmtree(gemini_workdir, ignore_errors=True)

    return results


# ---------------------------------------------------------------------------
# Summary Generation
# ---------------------------------------------------------------------------
def make_phase_summary(phase_key: str, phase_result: dict) -> dict:
    """Extract concise summary from a phase result (no raw_json)."""
    summary: dict = {
        "completed": phase_result["completed"],
        "duration_ms": phase_result["duration_ms"],
        "exit_code": phase_result["exit_code"],
    }
    # Always include output — needed for response verification checks in run_story.py
    summary["output"] = phase_result.get("output", "")
    # Always include tool call trace from stderr (lines containing "[tool]")
    stderr = phase_result.get("stderr", "")
    if stderr:
        tool_lines = [line for line in stderr.splitlines() if "[tool]" in line]
        if tool_lines:
            summary["tool_trace"] = tool_lines
    if not phase_result["completed"] and stderr:
        summary["stderr"] = stderr
    # Always include stats (for comparison between branches)
    if phase_result.get("num_turns") is not None:
        summary["num_turns"] = phase_result["num_turns"]
    if phase_result.get("session_id") is not None:
        summary["session_id"] = phase_result["session_id"]
    if phase_result.get("cost_usd") is not None:
        summary["cost_usd"] = phase_result["cost_usd"]
    if phase_result.get("tool_stats") is not None:
        summary["tool_stats"] = phase_result["tool_stats"]
    if phase_result.get("tokens_input") is not None:
        summary["tokens_input"] = phase_result["tokens_input"]
    if phase_result.get("tokens_output") is not None:
        summary["tokens_output"] = phase_result["tokens_output"]
    return summary


def aggregate_agent_stats(agent_data: dict) -> dict:
    """Calculate aggregate stats across all phases for an agent."""
    total_duration = 0
    total_turns = 0
    total_tool_calls = 0
    total_tool_success = 0
    total_tool_fail = 0
    has_turn_data = False
    has_tool_stats = False
    # First non-None value across phases: measures the idle context size before
    # any tool calls happen, regardless of which phase the first LLM call lands in.
    tokens_first_input: int | None = None

    for phase_key in ("setup", "test", "teardown"):
        if phase_key not in agent_data:
            continue
        phase = agent_data[phase_key]
        total_duration += phase.get("duration_ms", 0)
        if "num_turns" in phase:
            has_turn_data = True
            total_turns += phase["num_turns"]

        # Extract tool call counts from tool_stats
        tool_stats = phase.get("tool_stats")
        if tool_stats:
            has_tool_stats = True
            # Gemini and OpenAI format: {totalCalls, totalSuccess, totalFail, ...}
            if "totalCalls" in tool_stats:
                total_tool_calls += tool_stats.get("totalCalls", 0)
                total_tool_success += tool_stats.get("totalSuccess", 0)
                total_tool_fail += tool_stats.get("totalFail", 0)
            # Claude format might differ - handle if needed

        if tokens_first_input is None and phase.get("tokens_first_input") is not None:
            tokens_first_input = phase["tokens_first_input"]

    return {
        "total_duration_ms": total_duration,
        "total_turns": total_turns if has_turn_data else None,
        # Use has_tool_stats to distinguish "no data" (None) from "0 calls" (0).
        # A local model that answers without calling tools should show 0, not null.
        "total_tool_calls": total_tool_calls if has_tool_stats else None,
        "total_tool_success": total_tool_success if has_tool_stats else None,
        "total_tool_fail": total_tool_fail if has_tool_stats else None,
        "tokens_first_input": tokens_first_input,
    }


def make_summary(full_results: dict) -> dict:
    """Build a concise summary from full results (for stdout)."""
    summary: dict = {
        "mcp_source": full_results["mcp_source"],
        "branch": full_results.get("branch"),
        "agents": {},
    }
    for agent_name, agent_data in full_results["results"].items():
        if not agent_data.get("available", False):
            summary["agents"][agent_name] = {"available": False}
            continue

        agent_summary: dict = {"available": True, "all_passed": True}
        for phase_key in ("setup", "test", "teardown"):
            if phase_key not in agent_data:
                continue
            phase_summary = make_phase_summary(phase_key, agent_data[phase_key])
            agent_summary[phase_key] = phase_summary
            if not phase_summary["completed"]:
                agent_summary["all_passed"] = False

        # Add aggregate stats
        agg_stats = aggregate_agent_stats(agent_data)
        agent_summary["aggregate"] = agg_stats

        summary["agents"][agent_name] = agent_summary

    return summary


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
async def run(args: argparse.Namespace) -> dict:
    """Execute the BAT scenario and return results."""
    # Read scenario
    if args.scenario_file:
        scenario = json.loads(Path(args.scenario_file).read_text())  # noqa: ASYNC240
    else:
        if sys.stdin.isatty():
            raise ValueError(
                "No scenario provided. Pipe scenario JSON via stdin, or pass --scenario-file.\n"
                "  echo '{\"test_prompt\":\"...\"}' | uv run python tests/uat/run_uat.py --agents gemini\n"
                "  uv run python tests/uat/run_uat.py --scenario-file scenario.json --agents gemini\n"
                "For the pre-built story catalog, use tests/uat/stories/run_story.py --all."
            )
        scenario = json.loads(sys.stdin.read())

    if "test_prompt" not in scenario:
        raise ValueError("scenario must contain 'test_prompt'")

    # Determine agents
    requested_agents = [a.strip() for a in args.agents.split(",")]
    agents: dict[str, bool] = {}
    for name in requested_agents:
        available = check_agent_available(name)
        agents[name] = available
        if not available:
            logger.warning(f"{name} CLI not found, skipping")

    active_agents = [name for name, avail in agents.items() if avail]
    if not active_agents:
        raise ValueError("No agents available")

    if "openai" in active_agents and not getattr(args, "base_url", None):
        raise ValueError(
            "--base-url is required when using the openai agent. "
            "Example: --base-url http://localhost:1234/v1"
        )

    # Preflight: fail fast if Docker or the OpenAI endpoint is unreachable,
    # rather than stalling inside container startup / model warmup.
    if not args.ha_url:
        err = preflight_check_docker()
        if err:
            raise RuntimeError(err)
    if "openai" in active_agents and getattr(args, "base_url", None):
        err = preflight_check_base_url(args.base_url)
        if err:
            raise RuntimeError(err)

    # Start HA (container or external)
    ha_url = args.ha_url
    ha_token = args.ha_token or TEST_TOKEN
    mcp_source = "branch" if args.branch else "local"

    container: HAContainer | None = None
    if not ha_url:
        container = HAContainer()
        container.__enter__()
        ha_url = container.url
        ha_token = container.token

    try:
        logger.info(f"HA: {ha_url}")
        logger.info(f"MCP source: {mcp_source}" + (f" ({args.branch})" if args.branch else ""))
        logger.info(f"Agents: {', '.join(active_agents)}")

        extra_env = parse_mcp_env(
            getattr(args, "mcp_env", None),
            base_url=getattr(args, "base_url", None),
            on_default_applied=logger.info,
        )

        # Run agents sequentially to avoid resource contention
        agent_results = {}
        for name in active_agents:
            agent_results[name] = await run_agent_scenario(
                name,
                scenario,
                ha_url,
                ha_token,
                args.branch,
                args.timeout,
                model=getattr(args, "model", None),
                base_url=getattr(args, "base_url", None),
                api_key=getattr(args, "api_key", "no-key"),
                max_tools=getattr(args, "max_tools", None),
                no_think=getattr(args, "no_think", False),
                max_tokens=getattr(args, "max_tokens", None),
                extra_env=extra_env,
            )

        # Add unavailable agents
        for name, avail in agents.items():
            if not avail:
                agent_results[name] = {"available": False}

        return {
            "scenario": scenario,
            "ha_url": ha_url,
            "mcp_source": mcp_source,
            "branch": args.branch,
            "results": agent_results,
        }
    finally:
        if container:
            container.__exit__(None, None, None)


def main() -> None:
    configure_cli_logging()

    parser = SuggestingArgumentParser(
        description="BAT Runner - Execute MCP test scenarios on AI agent CLIs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  echo '{"test_prompt":"Search for light entities."}' | uv run python tests/uat/run_uat.py --agents gemini
  uv run python tests/uat/run_uat.py --scenario-file /tmp/scenario.json --agents claude,gemini
  uv run python tests/uat/run_uat.py --ha-url http://localhost:8123 --ha-token TOKEN --agents gemini
  uv run python tests/uat/run_uat.py --branch feat/tool-errors --agents gemini
        """,
    )
    parser.add_argument(
        "--agents",
        default=DEFAULT_AGENTS,
        help=f"Comma-separated list of agents to run (default: {DEFAULT_AGENTS})",
    )
    parser.add_argument(
        "--scenario-file",
        help="Read scenario from file instead of stdin",
    )
    parser.add_argument(
        "--ha-url",
        help="Use an existing HA instance instead of starting a container",
    )
    parser.add_argument(
        "--ha-token",
        help="HA long-lived access token (default: test token)",
    )
    parser.add_argument(
        "--branch",
        help="Git branch/tag to install ha-mcp from (default: local code)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT,
        help=f"Timeout per phase in seconds (default: {DEFAULT_TIMEOUT})",
    )
    parser.add_argument(
        "--model",
        help="Model to use (e.g., haiku/sonnet/opus for Claude, or model name for openai agent)",
    )
    parser.add_argument(
        "--base-url",
        help="OpenAI-compatible API base URL (for --agents openai)",
    )
    parser.add_argument(
        "--api-key",
        default="no-key",
        help="API key for OpenAI-compatible endpoint (default: no-key)",
    )
    parser.add_argument(
        "--max-tools",
        type=int,
        default=None,
        help="Limit MCP tools passed to the openai agent (useful for small context windows)",
    )
    parser.add_argument(
        "--no-think",
        action="store_true",
        help="Prepend /no_think to disable reasoning mode (qwen3 and compatible models)",
    )
    parser.add_argument(
        "--mcp-env",
        action="append",
        metavar="KEY=VALUE",
        help="Extra env var for the MCP server (repeatable, e.g. --mcp-env ENABLE_TOOL_SEARCH=true)",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=None,
        help="Max output tokens per completion for openai agent (agent default: 8192)",
    )
    args = parser.parse_args()

    try:
        full_results = asyncio.run(run(args))
    except ValueError as e:
        logger.error(str(e))
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)

    # Write full results to temp file
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".json", prefix="bat_results_", delete=False
    ) as results_file:
        json.dump(full_results, results_file, indent=2)

    # Output concise summary + file path to stdout
    summary = make_summary(full_results)
    summary["results_file"] = results_file.name
    json.dump(summary, sys.stdout, indent=2)
    print()  # trailing newline


if __name__ == "__main__":
    main()
