#!/usr/bin/env python3
"""
Standalone story runner.

Runs YAML stories against a HA test instance:
- Setup: FastMCP in-memory (sub-second, deterministic)
- Test prompt: AI agent CLI via run_uat.py (gemini/claude/openai)
- Each agent gets a fresh HA container (clean state)

Results are appended to a JSONL file for historical tracking.

Usage:
    # Run a single story
    uv run python tests/uat/stories/run_story.py catalog/s01_automation_sunset_lights.yaml --agents gemini

    # Run all stories
    uv run python tests/uat/stories/run_story.py --all --agents gemini

    # Run with a local OpenAI-compatible LLM (LM Studio, Ollama, etc.)
    uv run python tests/uat/stories/run_story.py --all --agents openai --base-url http://localhost:1234/v1

    # Run against a specific branch/tag
    uv run python tests/uat/stories/run_story.py --all --agents gemini --branch v6.6.1

    # Use an existing HA instance instead of starting a container
    uv run python tests/uat/stories/run_story.py --all --agents gemini --ha-url http://localhost:8123

    # Keep container alive after run (for verification/debugging)
    uv run python tests/uat/stories/run_story.py --all --agents gemini --keep-container

    # Just print the BAT scenario JSON
    uv run python tests/uat/stories/run_story.py catalog/s01_automation_sunset_lights.yaml --dry-run
"""

from __future__ import annotations

import argparse
import asyncio
import contextlib
import json
import logging
import subprocess
import sys
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING

import yaml

if TYPE_CHECKING:
    import openai
    from fastmcp import Client as MCPClient

SCRIPT_DIR = Path(__file__).resolve().parent
CATALOG_DIR = SCRIPT_DIR / "catalog"
REPO_ROOT = SCRIPT_DIR.parent.parent.parent
TESTS_DIR = REPO_ROOT / "tests"
RUN_UAT = SCRIPT_DIR.parent / "run_uat.py"
DEFAULT_RESULTS_FILE = REPO_ROOT / "local" / "uat-results.jsonl"

# Agents that can run in-process with a persistent MCP+OpenAI client across
# stories. Claude and Gemini are external CLIs and take the subprocess path.
INLINE_AGENTS: frozenset[str] = frozenset({"openai"})

# A single request timeout is a per-story failure (the backend may be alive but
# slow). This many *consecutive* timeouts means the backend is effectively gone
# (unresponsive, or dropping packets so connects time out rather than refuse),
# so the suite aborts rather than timing out on every remaining story.
CONSECUTIVE_TIMEOUT_ABORT = 2

# test_phase["error"] marker for a request timeout. Shared by the producer (the
# inline timeout handler) and the consumer (the story loop's streak counter) so
# the two can't drift apart and silently disable the consecutive-timeout abort.
TIMEOUT_ERROR_MARKER = "timeout"

# Add paths for imports
sys.path.insert(0, str(REPO_ROOT / "src"))
sys.path.insert(0, str(TESTS_DIR))
sys.path.insert(0, str(SCRIPT_DIR))  # for scripts/ subdirectory imports

from scripts.verify_story import verify_ha_checks  # noqa: E402
from uat._inprocess import inprocess_mcp_client  # noqa: E402
from uat._logging import configure_cli_logging  # noqa: E402
from uat.ha_wait import wait_for_ha_ready  # noqa: E402
from uat.run_uat import SuggestingArgumentParser  # noqa: E402

logger = logging.getLogger("uat.stories.run_story")


# ---------------------------------------------------------------------------
# Story loading
# ---------------------------------------------------------------------------
def load_story(path: Path) -> dict:
    """Load a story YAML file."""
    with open(path) as f:
        return yaml.safe_load(f)


# ---------------------------------------------------------------------------
# HA Container
# ---------------------------------------------------------------------------
def _start_container(*, keep_alive: bool = False) -> dict:
    """Start a HA test container, return {url, token, container, config_dir}.

    Args:
        keep_alive: If True, disable ryuk so the container survives process exit.
    """
    import os
    import shutil
    import tempfile

    from test_constants import HA_TEST_IMAGE, TEST_TOKEN
    from testcontainers.core.container import DockerContainer

    if keep_alive:
        os.environ["TESTCONTAINERS_RYUK_DISABLED"] = "true"

    HA_IMAGE = HA_TEST_IMAGE

    # Copy initial_test_state
    config_dir = Path(tempfile.mkdtemp(prefix="ha_story_"))
    initial_state = TESTS_DIR / "initial_test_state"
    shutil.copytree(initial_state, config_dir, dirs_exist_ok=True)
    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)

    container = (
        DockerContainer(HA_IMAGE)
        .with_exposed_ports(8123)
        .with_volume_mapping(str(config_dir), "/config", "rw")
        .with_env("TZ", "UTC")
        .with_kwargs(privileged=True)
    )
    container.start()

    try:
        port = container.get_exposed_port(8123)
        url = f"http://localhost:{port}"
        logger.info(f"HA container started on {url}")

        # Wait for HA to be fully ready (API + components + entities)
        wait_for_ha_ready(url, TEST_TOKEN)
    except Exception:
        container.stop()
        shutil.rmtree(config_dir, ignore_errors=True)
        raise

    return {
        "url": url,
        "token": TEST_TOKEN,
        "container": container,
        "config_dir": config_dir,
    }


def _stop_container(ha: dict) -> None:
    """Stop HA container and clean up."""
    import shutil

    logger.info("Stopping HA container...")
    ha["container"].stop()
    shutil.rmtree(ha["config_dir"], ignore_errors=True)


# ---------------------------------------------------------------------------
# Token extraction from session files
# ---------------------------------------------------------------------------
def _extract_tokens(session_file: str | None, agent: str) -> dict | None:
    """Extract token usage from an agent session file.

    Returns dict with keys: input, output, cached, thoughts (all ints).
    Only non-cached tokens matter for cost — cached tokens are free.
    """
    if not session_file or not Path(session_file).exists():
        return None

    try:
        if agent == "gemini":
            data = json.loads(Path(session_file).read_text())
            totals = {"input": 0, "output": 0, "cached": 0, "thoughts": 0}
            for msg in data.get("messages", []):
                t = msg.get("tokens", {})
                totals["input"] += t.get("input", 0)
                totals["output"] += t.get("output", 0)
                totals["cached"] += t.get("cached", 0)
                totals["thoughts"] += t.get("thoughts", 0)
            return totals

        if agent == "claude":
            totals = {"input": 0, "output": 0, "cached": 0, "thoughts": 0}
            for line in Path(session_file).read_text().splitlines():
                entry = json.loads(line)
                if entry.get("type") == "assistant":
                    usage = entry.get("message", {}).get("usage", {})
                    totals["input"] += usage.get("input_tokens", 0)
                    totals["output"] += usage.get("output_tokens", 0)
                    totals["cached"] += usage.get(
                        "cache_read_input_tokens", 0
                    ) + usage.get("cache_creation_input_tokens", 0)
            return totals
    except Exception as exc:
        logger.warning(f"  Token extraction failed: {exc}")
        return None

    return None


def _extract_model(session_file: str | None, agent: str) -> str | None:
    """Extract the resolved model id from a subprocess agent's session file.

    Claude records the resolved id (e.g. ``claude-sonnet-4-6``) as
    ``message.model`` on each assistant entry, so the recorded model is
    accurate even on a bare ``--agents claude`` run with no ``--model``.
    Returns None for agents whose session format does not expose it
    (gemini) or on any error, so the caller falls back to None.
    """
    if not session_file or not Path(session_file).exists():
        return None
    try:
        if agent == "claude":
            for line in Path(session_file).read_text().splitlines():
                entry = json.loads(line)
                if entry.get("type") == "assistant":
                    model = entry.get("message", {}).get("model")
                    if model:
                        return model
    except Exception as exc:
        logger.warning(f"  Model extraction failed: {exc}")
        return None

    return None


def _extract_tool_calls(session_file: str | None, agent: str) -> int | None:
    """Count tool calls from an agent session file."""
    seq = _extract_tool_sequence(session_file, agent)
    return None if seq is None else len(seq)


def _extract_tool_sequence(session_file: str | None, agent: str) -> list[str] | None:
    """Extract the ordered tool-call sequence (names only) from a session file.

    Used by the BAT to count flip-flops and identify pick patterns on the
    native ``claude``/``gemini`` paths (the openai path produces this via
    ``openai_agent.tool_trace_sink`` directly in the test_phase result).
    """
    if not session_file or not Path(session_file).exists():
        return None

    try:
        if agent == "gemini":
            data = json.loads(Path(session_file).read_text())
            seq: list[str] = []
            for msg in data.get("messages", []):
                for call in msg.get("toolCalls", []):
                    name = call.get("name") or call.get("toolName")
                    if name:
                        seq.append(name)
            return seq

        if agent == "claude":
            seq = []
            for line in Path(session_file).read_text().splitlines():
                entry = json.loads(line)
                if entry.get("type") == "assistant":
                    for block in entry.get("message", {}).get("content", []):
                        if block.get("type") == "tool_use":
                            name = block.get("name")
                            if name:
                                seq.append(name)
            return seq
    except Exception as exc:
        logger.warning(f"  Tool sequence extraction failed: {exc}")
        return None

    return None


# ---------------------------------------------------------------------------
# Session file detection
# ---------------------------------------------------------------------------
def _find_session_file_by_id(session_id: str) -> str | None:
    """Find a Claude session file by its session_id (UUID).

    Claude stores sessions at ~/.claude/projects/<dir>/<session_id>.jsonl
    """
    home = Path.home()
    claude_projects = home / ".claude" / "projects"
    if not claude_projects.exists():
        return None
    matches = list(claude_projects.glob(f"*/{session_id}.jsonl"))
    return str(matches[0]) if matches else None


def _find_latest_session_file(agent: str, after: float) -> str | None:
    """Find the most recent session file for an agent created after a timestamp.

    Args:
        agent: Agent name ("gemini" or "claude").
        after: Only consider files modified after this unix timestamp.

    Gemini: ~/.gemini/tmp/<hash>/chats/session-*.json
    Claude: Use session_id from JSON output instead (see _find_session_file_by_id).
    """
    home = Path.home()

    if agent == "gemini":
        gemini_tmp = home / ".gemini" / "tmp"
        if not gemini_tmp.exists():
            return None
        session_files = [
            p
            for p in gemini_tmp.glob("*/chats/session-*.json")
            if p.stat().st_mtime > after
        ]
        if not session_files:
            return None
        return str(max(session_files, key=lambda p: p.stat().st_mtime))

    return None


# ---------------------------------------------------------------------------
async def _run_mcp_steps(mcp_client: MCPClient, steps: list[dict], phase: str) -> None:
    """Execute setup or teardown steps via a shared in-memory MCP client."""
    for step in steps:
        tool_name = step["tool"]
        args = step.get("args", {})
        logger.info(f"  [{phase}] {tool_name}({args})")
        try:
            await mcp_client.call_tool(tool_name, args)
        except Exception:
            if phase == "setup":
                logger.info(f"  [{phase}] {tool_name} FAILED (see server log)")
                raise
            logger.warning(
                f"  [{phase}] {tool_name} failed, ignored (may poison shared client "
                "for next story; see server log)"
            )


# ---------------------------------------------------------------------------
# Test prompt execution via agent CLI
# ---------------------------------------------------------------------------
def _run_test_prompt(
    prompt: str,
    agent: str,
    ha_url: str,
    ha_token: str,
    branch: str | None = None,
    extra_args: list[str] | None = None,
    model: str | None = None,
    base_url: str | None = None,
    api_key: str | None = None,
    max_tools: int | None = None,
    no_think: bool = False,
    max_tokens: int | None = None,
    mcp_env: list[str] | None = None,
) -> tuple[int, dict | None]:
    """Run test prompt via run_uat.py for a single agent. Returns (exit_code, parsed_summary)."""
    scenario = {"test_prompt": prompt.strip()}

    cmd = [
        sys.executable,
        str(RUN_UAT),
        "--agents",
        agent,
        "--ha-url",
        ha_url,
        "--ha-token",
        ha_token,
    ]
    if branch:
        cmd.extend(["--branch", branch])
    if model:
        cmd.extend(["--model", model])
    if base_url:
        cmd.extend(["--base-url", base_url])
    if api_key:
        cmd.extend(["--api-key", api_key])
    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)])
    if mcp_env:
        for pair in mcp_env:
            cmd.extend(["--mcp-env", pair])
    if extra_args:
        cmd.extend(extra_args)

    result = subprocess.run(
        cmd,
        input=json.dumps(scenario),
        capture_output=True,
        text=True,
        timeout=600,
    )

    for line in result.stderr.splitlines():
        logger.info(line)

    summary = None
    if result.stdout.strip():
        try:
            summary = json.loads(result.stdout)
        except json.JSONDecodeError:
            # Non-JSON stdout (e.g. partial/garbled output); leave summary as None.
            pass

    return result.returncode, summary


async def _run_test_prompt_inline(
    prompt: str,
    agent_name: str,
    openai_client: openai.AsyncOpenAI,
    mcp_client: MCPClient,
    model: str,
    openai_tools: list[dict],
    *,
    no_think: bool = False,
    max_tokens: int | None = None,
) -> tuple[int, dict]:
    """Run the test prompt via the openai_agent library directly.

    Reuses a persistent ``mcp_client``, ``openai_client``, and pre-fetched
    ``openai_tools`` across stories, so no subprocess spawn, MCP server
    restart, model warmup, or ``list_tools`` round trip happens per
    story. Returns ``(exit_code, summary)`` matching the shape of
    ``_run_test_prompt`` so the surrounding story-loop code is agnostic
    to how the test was executed. A failure always produces a summary
    (never ``None``) so the caller can append a row to the results file.
    """
    import traceback

    import openai
    from uat.openai_agent import (
        DEFAULT_MAX_TOKENS,
        ContextWindowExceededError,
        run_scenario_inline,
    )

    tool_trace: list[str] = []
    start = time.monotonic()
    try:
        result = await run_scenario_inline(
            openai_client,
            mcp_client,
            model,
            prompt.strip(),
            max_tokens=max_tokens or DEFAULT_MAX_TOKENS,
            no_think=no_think,
            tool_trace_sink=tool_trace,
            openai_tools=openai_tools,
        )
    except ContextWindowExceededError as e:
        # The conversation overflowed the backend context. Surface it as a
        # distinct hard failure so a verification-based PASS can't mask a run
        # that crashed mid-conversation (the entity may already exist from an
        # earlier turn). See _compute_passed's completed guard.
        logger.error(f"  [{agent_name}] context window exceeded: {e}")
        duration_ms = int((time.monotonic() - start) * 1000)
        return 1, _inline_failure_summary(
            agent_name,
            error_msg="context_window_exceeded",
            duration_ms=duration_ms,
            tool_trace=tool_trace,
            context_overflow={
                "requested_tokens": e.requested_tokens,
                "context_size": e.context_size,
            },
        )
    except openai.APITimeoutError as e:
        # A request timed out, but the backend may well be alive: a slow or
        # degenerate-loop generation on a local model can outrun even a generous
        # per-request timeout. Treat it as a per-story failure, NOT suite-fatal.
        # completed=False (via _inline_failure_summary) so a state-based verify
        # can't mask it. The story loop aborts only after several *consecutive*
        # timeouts (a genuinely unresponsive or packet-dropping backend, whose
        # connect attempts surface as timeouts rather than refusals). Must be
        # caught before APIConnectionError; APITimeoutError subclasses it.
        logger.warning(f"  [{agent_name}] request timed out: {e}")
        duration_ms = int((time.monotonic() - start) * 1000)
        return 1, _inline_failure_summary(
            agent_name,
            error_msg=TIMEOUT_ERROR_MARKER,
            duration_ms=duration_ms,
            tool_trace=tool_trace,
        )
    except openai.APIConnectionError:
        # The LLM backend is unreachable after the SDK exhausted its own retries:
        # a refused or reset connection, NOT a timeout (handled above). This is
        # suite-fatal, not a per-story failure: propagate so the story loop
        # aborts instead of swallowing it into a FAIL summary and then failing on
        # every remaining story.
        raise
    except Exception as e:
        logger.exception(f"  [{agent_name}] inline run failed")
        tb = traceback.format_exc()
        duration_ms = int((time.monotonic() - start) * 1000)
        return 1, _inline_failure_summary(
            agent_name,
            error_msg=f"{type(e).__name__}: {e}",
            traceback_text=tb,
            duration_ms=duration_ms,
            tool_trace=tool_trace,
        )

    duration_ms = int((time.monotonic() - start) * 1000)
    exit_code = 1 if result.get("hit_iteration_limit") else 0

    # Match the summary shape produced by run_uat.make_summary.
    test_phase = {
        "completed": True,
        "duration_ms": duration_ms,
        "exit_code": exit_code,
        "model": result.get("model"),
        "output": result.get("result", ""),
        "num_turns": result.get("num_turns"),
        "tool_stats": result.get("tool_stats"),
        "tokens_input": result.get("tokens_input"),
        "tokens_output": result.get("tokens_output"),
        "tokens_thoughts": result.get("tokens_thoughts"),
        "tool_trace": tool_trace,
        "cost_usd": result.get("cost_usd", 0),
    }
    if result.get("hit_iteration_limit"):
        test_phase["error"] = "hit_iteration_limit"
    tool_stats = result.get("tool_stats") or {}
    aggregate = {
        "total_duration_ms": duration_ms,
        "total_turns": result.get("num_turns"),
        "total_tool_calls": tool_stats.get("totalCalls"),
        "total_tool_success": tool_stats.get("totalSuccess"),
        "total_tool_fail": tool_stats.get("totalFail"),
        "tokens_first_input": result.get("tokens_first_input"),
    }
    summary = {
        "agents": {
            agent_name: {
                "available": True,
                "all_passed": exit_code == 0,
                "test": test_phase,
                "aggregate": aggregate,
            }
        }
    }
    return exit_code, summary


def _inline_failure_summary(
    agent_name: str,
    *,
    error_msg: str,
    traceback_text: str | None = None,
    duration_ms: int = 0,
    tool_trace: list[str] | None = None,
    context_overflow: dict | None = None,
) -> dict:
    """Build a summary dict for a failed inline run (matches make_summary shape)."""
    test_phase: dict = {
        "completed": False,
        "duration_ms": duration_ms,
        "exit_code": 1,
        "output": traceback_text or error_msg,
        "error": error_msg,
        "tool_trace": tool_trace or [],
    }
    if context_overflow is not None:
        test_phase["context_overflow"] = context_overflow
    return {
        "agents": {
            agent_name: {
                "available": True,
                "all_passed": False,
                "test": test_phase,
                "aggregate": {"total_duration_ms": duration_ms},
            }
        }
    }


def _record_setup_failure(
    filtered: list,
    agent: str,
    error_msg: str,
    *,
    all_results: list,
    results_file,
    sha: str,
    describe: str,
    branch: str | None,
) -> None:
    """Write a failure row per story and record it in all_results.

    Called when per-agent setup (OpenAI client, MCP server) fails before
    any story runs — otherwise every filtered story is silently dropped
    from both the JSONL log and the process exit code.
    """
    for _path, story in filtered:
        summary = _inline_failure_summary(agent, error_msg=error_msg)
        all_results.append((agent, story["id"], story, 1, summary, None, False))
        append_result(
            results_file,
            story,
            agent,
            sha,
            describe,
            branch,
            summary,
            None,
            passed=False,
            verify_results=None,
        )


# ---------------------------------------------------------------------------
# Git info
# ---------------------------------------------------------------------------
def get_git_info() -> tuple[str, str]:
    """Get (short SHA, git describe) for the current commit."""
    sha = "unknown"
    describe = "unknown"
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True,
            text=True,
            cwd=str(REPO_ROOT),
        )
        sha = result.stdout.strip()
    except Exception:
        # Best-effort git lookup; keep the "unknown" default if git is unavailable.
        pass
    try:
        result = subprocess.run(
            ["git", "describe", "--tags", "--always"],
            capture_output=True,
            text=True,
            cwd=str(REPO_ROOT),
        )
        describe = result.stdout.strip()
    except Exception:
        # Best-effort git lookup; keep the "unknown" default if git is unavailable.
        pass
    return sha, describe


# ---------------------------------------------------------------------------
# Pass/fail determination
# ---------------------------------------------------------------------------
def _run_completed(test_phase: dict, summary: dict | None, exit_code: int) -> bool:
    """Whether the agent run completed, as opposed to crashing mid-run.

    Fail-closed: a subprocess agent that crashes produces unparseable stdout
    (summary is None → empty test_phase), which must not count as completed
    just because ha_checks pass on state left from an earlier turn. That is the
    same masking class the inline ContextWindowExceededError path closes. Both
    the inline and subprocess summary builders set ``completed`` explicitly, so
    the fallback only fires when the summary is missing entirely.
    """
    return test_phase.get("completed", summary is not None and exit_code == 0)


def _compute_passed(
    exit_code: int,
    tool_calls: int | None,
    verify_results: list[dict] | None,
    completed: bool = True,
) -> bool:
    """Determine passed based on tool usage and verification results.

    Rules:
    - Run did not complete (crashed mid-run, e.g. context overflow) → fail.
      Checked first so a verification-based PASS can't mask a crashed agent
      that happened to leave good HA state from an earlier turn.
    - Zero tool calls → fail (agent did nothing useful)
    - ha_checks present → all checks must pass
    - No checks → fall back to exit code
    """
    if not completed:
        return False
    if tool_calls == 0:
        return False
    if verify_results is not None:
        return all(r["passed"] for r in verify_results)
    return exit_code == 0


def _suite_exit_code(all_results: list, backend_aborted: bool) -> int:
    """Process exit code for a story run.

    An aborted suite (the LLM backend went unreachable mid-run) is never a clean
    pass, even if every story that did run passed, because it never ran to
    completion. Otherwise nonzero iff any completed story run failed.
    """
    if backend_aborted:
        return 1
    return 1 if any(not passed for *_, passed in all_results) else 0


# ---------------------------------------------------------------------------
# JSONL results
# ---------------------------------------------------------------------------
def append_result(
    results_file: Path,
    story: dict,
    agent: str,
    sha: str,
    describe: str,
    branch: str | None,
    bat_summary: dict,
    session_file: str | None = None,
    passed: bool = False,
    verify_results: list[dict] | None = None,
    model: str | None = None,
    quantization: str | None = None,
) -> None:
    """Append a single story result as one JSONL line."""
    agent_data = bat_summary.get("agents", {}).get(agent, {})
    test_phase = agent_data.get("test", {})
    aggregate = agent_data.get("aggregate", {})

    record = {
        "sha": sha,
        "version": describe,
        "branch": branch,
        "timestamp": datetime.now(UTC).isoformat(),
        "agent": agent,
        "model": model
        or test_phase.get("model")
        or _extract_model(session_file, agent),
        "quantization": quantization,
        "story": story["id"],
        "category": story["category"],
        "weight": story["weight"],
        "passed": passed,
        "test_duration_ms": test_phase.get("duration_ms"),
        "total_duration_ms": aggregate.get("total_duration_ms"),
        "tool_calls": aggregate.get("total_tool_calls"),
        "tool_failures": aggregate.get("total_tool_fail"),
        "turns": test_phase.get("num_turns"),
    }
    if session_file:
        record["session_file"] = session_file

    # Surface the failure-reason marker so reports can tell an incomplete run
    # (e.g. context_window_exceeded, hit_iteration_limit) apart from a run that
    # finished but failed its checks.
    error = test_phase.get("error")
    if error:
        record["error"] = error
    context_overflow = test_phase.get("context_overflow")
    if context_overflow:
        record["context_overflow"] = context_overflow

    # Cost from Claude's JSON output (direct, no session file needed)
    cost_usd = test_phase.get("cost_usd")
    if cost_usd is not None:
        record["cost_usd"] = cost_usd

    # Extract tool call count from session file if not in summary.
    # Always extract the ordered sequence too — flip-flop count is the BAT's
    # primary native-path metric (openai path emits the sequence directly).
    seq = _extract_tool_sequence(session_file, agent)
    if record["tool_calls"] is None and seq is not None:
        record["tool_calls"] = len(seq)
    if seq:
        record["tool_sequence"] = seq

    # Extract token usage: from phase summary (openai) or session file (claude/gemini)
    tokens = _extract_tokens(session_file, agent)
    if tokens is None:
        ti = test_phase.get("tokens_input")
        to = test_phase.get("tokens_output")
        if ti is not None or to is not None:
            # reasoning_tokens are a subset of completion_tokens, so tokens_output
            # already counts them; tokens_thoughts is the informational breakdown.
            tokens = {
                "input": ti or 0,
                "output": to or 0,
                "cached": 0,
                "thoughts": test_phase.get("tokens_thoughts") or 0,
            }
    if tokens:
        record["tokens_input"] = tokens["input"]
        record["tokens_output"] = tokens["output"]
        record["tokens_cached"] = tokens["cached"]
        record["tokens_thoughts"] = tokens["thoughts"]
        record["tokens_billable"] = (tokens["input"] - tokens["cached"]) + tokens[
            "output"
        ]

    # First-turn input tokens: most direct measure of idle context size (openai agent only).
    # Comes from aggregate (not test_phase) since make_phase_summary doesn't propagate it.
    tokens_first_input = aggregate.get("tokens_first_input")
    if tokens_first_input is not None:
        record["tokens_first_input"] = tokens_first_input

    # Always include agent_response so passing runs are inspectable.
    agent_response = test_phase.get("output", "")
    if agent_response:
        record["agent_response"] = agent_response

    # Always include tool_trace (openai path: ordered name+args from
    # tool_trace_sink) — the BAT's primary signal for tool-pick analysis,
    # not just for failure debug.
    tool_trace = test_phase.get("tool_trace")
    if tool_trace:
        record["tool_trace"] = tool_trace

    # verify_results stays failure-only — too verbose for the BAT result
    # file otherwise, and verify_results.passed is already captured in
    # record["passed"].
    if verify_results is not None and not all(r["passed"] for r in verify_results):
        record["verify_results"] = verify_results

    results_file.parent.mkdir(parents=True, exist_ok=True)
    with open(results_file, "a") as f:
        f.write(json.dumps(record, separators=(",", ":")) + "\n")


def append_abort_marker(
    results_file: Path,
    agent: str,
    sha: str,
    describe: str,
    branch: str | None,
    skipped_stories: list[str],
    detail: str,
) -> None:
    """Append a run-level marker recording that the suite aborted early.

    When the LLM backend goes unreachable the remaining stories are skipped
    rather than recorded, so without this marker the results file would be
    indistinguishable from a complete run that just happened to contain fewer
    stories. A reader (or bat-story-eval) could fold a crashed run into a
    baseline as a clean partial pass. The marker carries no ``story``/``passed``
    keys, so it can't be miscounted as a per-story result; it exists so the
    durable artifact is as self-describing as the nonzero exit code.
    """
    record = {
        "sha": sha,
        "version": describe,
        "branch": branch,
        "timestamp": datetime.now(UTC).isoformat(),
        "agent": agent,
        "aborted": True,
        "error": "backend_unreachable",
        "detail": detail,
        "skipped_stories": skipped_stories,
    }
    results_file.parent.mkdir(parents=True, exist_ok=True)
    with open(results_file, "a") as f:
        f.write(json.dumps(record, separators=(",", ":")) + "\n")


def _update_timeout_streak(consecutive: int, *, was_timeout: bool) -> int:
    """Running count of consecutive per-story request timeouts.

    A timeout increments the streak; any non-timeout story resets it to 0. The
    caller aborts the agent once the streak reaches ``CONSECUTIVE_TIMEOUT_ABORT``
    (the backend has stopped responding, vs one slow story on a live backend).
    """
    return consecutive + 1 if was_timeout else 0


def _stories_after(filtered: list[tuple[Path, dict]], story_index: int) -> list[str]:
    """Story ids after ``story_index``: the skip list for the consecutive-timeout
    abort. The timed-out story at ``story_index`` was already recorded as a
    per-story failure, so only later stories are skipped. Contrast the
    connection-error path, which skips from ``story_index`` itself because that
    story never validly ran.
    """
    return [s["id"] for _p, s in filtered[story_index + 1 :]]


def _record_backend_abort(
    results_file: Path,
    agent: str,
    sid: str,
    skipped_stories: list[str],
    detail: str,
    *,
    sha: str,
    describe: str,
    branch: str | None,
) -> None:
    """Log and record an early suite abort. Shared by the connection-error path
    (a refused/reset connection) and the consecutive-timeout path (a backend
    that has stopped responding). ``skipped_stories`` are NOT recorded as
    per-story rows, so the pass rate isn't corrupted by false fails."""
    logger.critical(
        f"[{agent}/{sid}] LLM backend unreachable ({detail}); aborting agent. "
        f"Skipped (not run): {', '.join(skipped_stories) if skipped_stories else '(none)'}"
    )
    append_abort_marker(
        results_file, agent, sha, describe, branch, skipped_stories, detail
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
async def run_stories(
    args: argparse.Namespace, filtered: list[tuple[Path, dict]]
) -> int:
    """Run stories with a fresh container per agent.

    For each agent: start container -> run all stories -> stop container.
    When --ha-url is provided, all agents share the external instance.
    """
    run_start = time.monotonic()
    sha, describe = get_git_info()
    agent_list = [a.strip() for a in args.agents.split(",")]
    using_external_ha = bool(args.ha_url)

    from uat.run_uat import (
        build_stdio_mcp_config,
        parse_mcp_env,
        preflight_check_base_url,
        preflight_check_docker,
    )

    if not using_external_ha:
        err = preflight_check_docker()
        if err:
            logger.critical(err)
            return 2
    if args.base_url and "openai" in agent_list:
        err = preflight_check_base_url(args.base_url)
        if err:
            logger.critical(err)
            return 2

    all_results: list[tuple[str, str, dict, int, dict | None, str | None, bool]] = []
    # Each entry: (agent, story_id, story, exit_code, summary, session_file, passed)
    backend_aborted = False  # set if any agent aborted on an unreachable backend

    mcp_env_dict = parse_mcp_env(
        getattr(args, "mcp_env", None),
        base_url=args.base_url,
        on_default_applied=logger.info,
    )
    effective_mcp_env: list[str] = [f"{k}={v}" for k, v in mcp_env_dict.items()]

    for agent in agent_list:
        logger.info(f"\n{'#' * 60}")
        logger.info(f"Agent: {agent}")
        logger.info(f"{'#' * 60}")

        ha = None
        ha_url = args.ha_url
        ha_token = args.ha_token
        if not using_external_ha:
            ha = _start_container(keep_alive=args.keep_container)
            ha_url = ha["url"]
            ha_token = ha["token"]

        # Inline agents (see INLINE_AGENTS) hold one MCP server + one OpenAI
        # client across all stories for this agent, skipping the per-story
        # MCP server spawn and model warmup. Other agents take the subprocess
        # path via _run_test_prompt.
        use_inline = agent in INLINE_AGENTS
        inline_mcp_client: MCPClient | None = None
        openai_client: openai.AsyncOpenAI | None = None
        resolved_model: str | None = None
        resolved_quant: str | None = None
        openai_tools: list[dict] = []

        async with contextlib.AsyncExitStack() as agent_stack:
            if ha and not args.keep_container:
                agent_stack.callback(_stop_container, ha)

            if use_inline:
                import openai  # runtime ref for APIConnectionError in the story loop
                from fastmcp import Client as _MCPClient
                from uat.openai_agent import (
                    create_and_warm_openai_client,
                    fetch_openai_tools,
                )

                config = build_stdio_mcp_config(
                    ha_url, ha_token, args.branch, mcp_env_dict or None
                )
                try:
                    (
                        openai_client,
                        resolved_model,
                        resolved_quant,
                    ) = await create_and_warm_openai_client(
                        base_url=args.base_url,
                        api_key=args.api_key,
                        model=args.model,
                    )
                    agent_stack.push_async_callback(openai_client.close)
                except Exception as e:
                    error_msg = (
                        f"Failed to initialise OpenAI client: {type(e).__name__}: {e}"
                    )
                    logger.error(f"[{agent}] {error_msg}")
                    _record_setup_failure(
                        filtered,
                        agent,
                        error_msg,
                        all_results=all_results,
                        results_file=args.results_file,
                        sha=sha,
                        describe=describe,
                        branch=args.branch,
                    )
                    continue
                try:
                    source = f"uvx download @ {args.branch}" if args.branch else "local"
                    logger.info(f"[{agent}] Starting MCP server ({source})...")
                    inline_mcp_client = await agent_stack.enter_async_context(
                        _MCPClient(config)
                    )
                    openai_tools = await fetch_openai_tools(
                        inline_mcp_client, max_tools=args.max_tools
                    )
                    logger.info(
                        f"[{agent}] MCP server ready ({len(openai_tools)} tools)"
                    )
                except Exception as e:
                    error_msg = f"Failed to start MCP server: {type(e).__name__}: {e}"
                    logger.error(f"[{agent}] {error_msg}")
                    _record_setup_failure(
                        filtered,
                        agent,
                        error_msg,
                        all_results=all_results,
                        results_file=args.results_file,
                        sha=sha,
                        describe=describe,
                        branch=args.branch,
                    )
                    continue

            if ha and args.keep_container:
                logger.info(f"\n[{agent}] Container kept alive: {ha['url']}")
                logger.info(f"[{agent}] Token: {ha['token']}")
                logger.info(f"[{agent}] Config dir: {ha['config_dir']}")
                logger.info(f"[{agent}] Stop manually: docker stop <container>")

            shared_mcp = await agent_stack.enter_async_context(
                inprocess_mcp_client(ha_url, ha_token)
            )

            # Consecutive per-story timeouts; reset by any non-timeout story.
            # Hitting CONSECUTIVE_TIMEOUT_ABORT aborts the agent (backend gone).
            consecutive_timeouts = 0

            for story_index, (_path, story) in enumerate(filtered):
                sid = story["id"]
                logger.info(f"\n{'=' * 60}")
                logger.info(f"[{agent}] Story {sid}: {story['title']}")
                logger.info(f"{'=' * 60}")

                setup_steps = story.get("setup") or []
                if setup_steps:
                    logger.info(
                        f"[{agent}/{sid}] Setup ({len(setup_steps)} steps via FastMCP)..."
                    )
                    await _run_mcp_steps(shared_mcp, setup_steps, "setup")

                logger.info(f"[{agent}/{sid}] Running test prompt...")
                # ``prompt_start`` is compared to file ``st_mtime`` in
                # ``_find_latest_session_file`` below, so it MUST be a
                # wall-clock Unix epoch (not ``time.monotonic()``).
                prompt_start = time.time()
                summary: dict | None
                if use_inline:
                    assert (
                        openai_client is not None
                        and inline_mcp_client is not None
                        and resolved_model is not None
                    ), (
                        "inline setup invariant: clients/model are populated when use_inline is True"
                    )
                    try:
                        rc, summary = await _run_test_prompt_inline(
                            story["prompt"],
                            agent_name=agent,
                            openai_client=openai_client,
                            mcp_client=inline_mcp_client,
                            model=resolved_model,
                            openai_tools=openai_tools,
                            no_think=args.no_think,
                            max_tokens=getattr(args, "max_tokens", None),
                        )
                    except openai.APIConnectionError as e:
                        # A refused/reset connection (NOT a timeout; those are
                        # handled per-story inside _run_test_prompt_inline). The
                        # backend is gone, so abort this agent's remaining stories
                        # rather than failing each one in turn. The current story
                        # never validly ran, so it is skipped too (not recorded).
                        _record_backend_abort(
                            args.results_file,
                            agent,
                            sid,
                            [s["id"] for _p, s in filtered[story_index:]],
                            f"{type(e).__name__}: {e}",
                            sha=sha,
                            describe=describe,
                            branch=args.branch,
                        )
                        backend_aborted = True
                        break
                else:
                    rc, summary = _run_test_prompt(
                        story["prompt"],
                        agent,
                        ha_url,
                        ha_token,
                        args.branch,
                        args.extra_args or None,
                        model=args.model,
                        base_url=args.base_url,
                        api_key=args.api_key,
                        max_tools=args.max_tools,
                        no_think=args.no_think,
                        max_tokens=getattr(args, "max_tokens", None),
                        mcp_env=effective_mcp_env or None,
                    )

                session_file = None
                test_phase = (
                    (summary or {}).get("agents", {}).get(agent, {}).get("test", {})
                )
                claude_session_id = test_phase.get("session_id")
                if claude_session_id:
                    session_file = _find_session_file_by_id(claude_session_id)
                if not session_file:
                    session_file = _find_latest_session_file(agent, after=prompt_start)

                verify_results = None
                ha_checks = (story.get("verify") or {}).get("ha_checks")
                if ha_checks:
                    agent_output = (
                        (summary or {})
                        .get("agents", {})
                        .get(agent, {})
                        .get("test", {})
                        .get("output", "")
                    )
                    logger.info(
                        f"[{agent}/{sid}] Verifying {len(ha_checks)} ha_check(s)..."
                    )
                    verify_results = await verify_ha_checks(
                        ha_url, ha_token, ha_checks, agent_output, shared_mcp
                    )
                    failed_checks = [r for r in verify_results if not r["passed"]]
                    if failed_checks:
                        logger.warning(
                            f"[{agent}/{sid}] {len(failed_checks)}/{len(ha_checks)} check(s) FAILED"
                        )
                        for r in failed_checks:
                            logger.warning(f"  FAIL [{r['type']}] {r['detail']}")
                    else:
                        logger.info(f"[{agent}/{sid}] All checks passed")

                agg = (
                    (summary or {})
                    .get("agents", {})
                    .get(agent, {})
                    .get("aggregate", {})
                )
                passed = _compute_passed(
                    exit_code=rc,
                    tool_calls=agg.get("total_tool_calls"),
                    verify_results=verify_results,
                    completed=_run_completed(test_phase, summary, rc),
                )
                all_results.append(
                    (agent, sid, story, rc, summary, session_file, passed)
                )

                if summary or verify_results is not None:
                    append_result(
                        args.results_file,
                        story,
                        agent,
                        sha,
                        describe,
                        args.branch,
                        summary or {},
                        session_file,
                        passed=passed,
                        verify_results=verify_results,
                        model=resolved_model or args.model,
                        quantization=resolved_quant,
                    )

                if session_file:
                    logger.info(f"[{agent}/{sid}] Session file: {session_file}")

                # A request timeout is recorded above as a per-story failure, but
                # several in a row mean the backend has stopped responding: abort
                # rather than time out on every remaining story. Unlike the
                # connection-error path (which skips the current story, never
                # validly run), the timed-out story here WAS recorded, so only
                # the stories after it are skipped.
                consecutive_timeouts = _update_timeout_streak(
                    consecutive_timeouts,
                    was_timeout=test_phase.get("error") == TIMEOUT_ERROR_MARKER,
                )
                if consecutive_timeouts >= CONSECUTIVE_TIMEOUT_ABORT:
                    _record_backend_abort(
                        args.results_file,
                        agent,
                        sid,
                        _stories_after(filtered, story_index),
                        f"{consecutive_timeouts} consecutive request timeouts",
                        sha=sha,
                        describe=describe,
                        branch=args.branch,
                    )
                    backend_aborted = True
                    break

    # Summary
    logger.info(f"\n{'=' * 60}")
    logger.info("Summary")
    logger.info(f"{'=' * 60}")
    for agent, sid, story, _rc, _, session_file, passed in all_results:
        status = "PASS" if passed else "FAIL"
        session_info = f" (session: {session_file})" if session_file else ""
        logger.info(f"  [{status}] {agent}/{sid}: {story['title']}{session_info}")

    elapsed = time.monotonic() - run_start
    mins, secs = divmod(int(elapsed), 60)
    logger.info(f"\nTotal time: {mins}m {secs}s")
    logger.info(f"Results appended to {args.results_file}")

    failed = sum(1 for *_, passed in all_results if not passed)
    total = len(all_results)
    if backend_aborted:
        logger.warning(
            f"\nRun aborted: LLM backend became unreachable "
            f"({failed}/{total} completed story runs failed before the abort)"
        )
    elif failed:
        logger.warning(f"\n{failed}/{total} story runs failed")
    else:
        logger.info(f"\nAll {total} story runs passed")
    return _suite_exit_code(all_results, backend_aborted)


def main() -> None:
    parser = SuggestingArgumentParser(
        description="Run user acceptance stories via BAT",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("story_file", nargs="?", help="Path to story YAML file")
    parser.add_argument(
        "--all", action="store_true", help="Run all stories in catalog/"
    )
    parser.add_argument("--agents", default="gemini", help="Comma-separated agent list")
    parser.add_argument("--branch", help="Git branch/tag to install ha-mcp from")
    parser.add_argument("--ha-url", help="Use existing HA instance (skip container)")
    parser.add_argument("--ha-token", help="HA long-lived access token")
    parser.add_argument(
        "--dry-run", action="store_true", help="Print BAT scenario JSON"
    )
    parser.add_argument(
        "--min-weight", type=int, default=1, help="Minimum story weight"
    )
    parser.add_argument(
        "--keep-container",
        action="store_true",
        help="Keep HA container alive after run (for verification/debugging)",
    )
    parser.add_argument(
        "--results-file",
        type=Path,
        default=DEFAULT_RESULTS_FILE,
        help=f"JSONL file to append results to (default: {DEFAULT_RESULTS_FILE})",
    )
    parser.add_argument(
        "--model",
        help="Model name (e.g., haiku, sonnet for Claude; auto-detected for openai)",
    )
    parser.add_argument(
        "--base-url",
        help="OpenAI-compatible API base URL (required for openai agent)",
    )
    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 (reduces context size)",
    )
    parser.add_argument(
        "--no-think",
        action="store_true",
        help=(
            "Disable reasoning mode: prepends /no_think (original Qwen3) and "
            "sends enable_thinking=false chat-template kwarg (Qwen3.5/3.6, "
            "honored by vLLM/llama-server)"
        ),
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=None,
        help="Max output tokens per completion for openai agent (agent default: 8192)",
    )
    parser.add_argument(
        "--mcp-env",
        action="append",
        metavar="KEY=VALUE",
        help=(
            "Extra env var for the MCP server (repeatable). "
            "ENABLE_TOOL_SEARCH defaults to true when --base-url is set "
            "(local model); override with --mcp-env ENABLE_TOOL_SEARCH=false."
        ),
    )
    parser.add_argument("extra_args", nargs="*", help="Extra args passed to run_uat.py")
    args = parser.parse_args()

    # Validate --base-url is provided when using the openai agent
    agent_list = [a.strip() for a in args.agents.split(",")]
    if "openai" in agent_list and not args.base_url:
        parser.error("--base-url is required when using the openai agent")

    configure_cli_logging()

    if args.all:
        stories = sorted(CATALOG_DIR.glob("s*.yaml"))
    elif args.story_file:
        story_path = Path(args.story_file).resolve()
        stories = [story_path]
    else:
        parser.print_help()
        sys.exit(1)

    # Filter by weight
    filtered = []
    for path in stories:
        story = load_story(path)
        if story.get("weight", 1) >= args.min_weight:
            filtered.append((path, story))

    if args.dry_run:
        for _, story in filtered:
            scenario = {"test_prompt": story["prompt"].strip()}
            print(f"# {story['id']}: {story['title']}")
            if story.get("setup"):
                print(f"# Setup: {len(story['setup'])} steps (FastMCP in-memory)")
            print(json.dumps(scenario, indent=2))
            print()
        return

    try:
        exit_code = asyncio.run(run_stories(args, filtered))
    except KeyboardInterrupt:
        logger.info("\nInterrupted")
        sys.exit(130)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
