"""Agent CLI stream parsing behind an adapter seam.

Each agent CLI emits a different event stream shape. The shared engine in

analyze_stream handles the lane-level concerns (turn timeline, failure trace,

recovery cost); a StreamAdapter supplies the per-CLI knowledge:

- how turn boundaries appear, and whether a turn's tool calls are carried

  inside the turn event (Claude-style: assistant message holds tool_use

  blocks, results follow) or precede the boundary (Codex-style:

  turn.completed arrives after the turn's command items);

- whether the absence of compaction markers is a measured zero (Claude

  streams emit explicit compact boundaries) or unmeasured (ADR-0002).

Adding a new agent CLI means adding an adapter and a golden transcript

fixture under tests/ — never editing the shared engine heuristics.

All values are None when the stream does not expose the underlying signal.

Never fabricate zeros here: a zero must mean "measured as zero" (ADR-0002).

"""

from __future__ import annotations

from dataclasses import dataclass

from typing import Any, Iterable

from oakbench.classify import (

    UNKNOWN_COMMAND_PATTERN,

    classify_command,

    unwrap_shell_command,

    vcs_info_family,

)

@dataclass

class StreamAnalysis:

    turn_source: str = "unmeasured"

    turns: list[dict[str, Any]] | None = None

    assistant_turns_total: int | None = None

    tool_calls_per_turn_avg: float | None = None

    max_tool_calls_per_turn: int | None = None

    cumulative_input_tokens: int | None = None

    peak_input_tokens: int | None = None

    context_truncation_events: int | None = None

    failed_tool_calls_total: int | None = None

    failed_vcs_commands_total: int | None = None

    failed_test_commands_total: int | None = None

    same_command_retries_total: int | None = None

    unknown_command_failures_total: int | None = None

    turns_to_recovery: int | None = None

    tokens_to_recovery: int | None = None

    vcs_info_followup_calls_total: int | None = None

    failed_command_examples: list[str] | None = None

    parallel_batches_total: int | None = None

    max_concurrency: int | None = None

def walk_json(value: Any) -> Iterable[Any]:

    yield value

    if isinstance(value, dict):

        for child in value.values():

            yield from walk_json(child)

    elif isinstance(value, list):

        for child in value:

            yield from walk_json(child)

def number(value: Any) -> int | None:

    if value is None or isinstance(value, bool):

        return None

    try:

        parsed = int(value)

    except (TypeError, ValueError):

        return None

    return parsed if parsed >= 0 else None

def first_number(mapping: dict[str, Any], keys: tuple[str, ...]) -> int | None:

    for key in keys:

        if key in mapping:

            value = number(mapping[key])

            if value is not None:

                return value

    return None

def usage_from_subtree(value: Any) -> dict[str, int | None] | None:

    input_keys = ("input_tokens", "prompt_tokens", "inputTokenCount", "promptTokens")

    output_keys = ("output_tokens", "completion_tokens", "outputTokenCount", "completionTokens")

    cache_read_keys = ("cache_read_input_tokens", "cache_read_tokens", "cached_tokens", "cached_input_tokens")

    for item in walk_json(value):

        if not isinstance(item, dict):

            continue

        if set(item).intersection(input_keys + output_keys):

            return {

                "input_tokens": first_number(item, input_keys),

                "output_tokens": first_number(item, output_keys),

                "cache_read_tokens": first_number(item, cache_read_keys),

            }

    return None

def is_assistant_turn_event(event: Any) -> bool:

    if not isinstance(event, dict):

        return False

    event_type = str(event.get("type") or "").lower()

    if event_type == "assistant":

        return True

    if event_type in {"turn.completed", "turn_completed"}:

        return True

    if event_type == "item.completed":

        item = event.get("item")

        if isinstance(item, dict) and str(item.get("type") or "") in {"agent_message", "assistant_message"}:

            return True

    return False

def assistant_event_tool_use_count(event: Any) -> int:

    return len(assistant_event_tool_use_refs(event))

def assistant_event_tool_use_refs(event: Any) -> list[str]:

    refs: list[str] = []

    anonymous = 0

    for item in walk_json(event):

        if not isinstance(item, dict) or str(item.get("type") or "") != "tool_use":

            continue

        tool_id = item.get("id")

        if tool_id:

            refs.append(str(tool_id))

        else:

            refs.append(f"anonymous:{anonymous}")

            anonymous += 1

    return refs

def cursor_tool_call(event: Any) -> tuple[str, str | None, bool | None] | None:

    if not isinstance(event, dict) or str(event.get("type") or "").lower() != "tool_call":

        return None

    tool_call = event.get("tool_call")

    if not isinstance(tool_call, dict):

        return None

    tool_name = next((key for key in tool_call if key.endswith("ToolCall")), None)

    if tool_name is None:

        return None

    payload = tool_call.get(tool_name)

    if not isinstance(payload, dict):

        return tool_name, None, None

    command = None

    args = payload.get("args")

    if isinstance(args, dict) and isinstance(args.get("command"), str):

        command = unwrap_shell_command(args["command"])

    failed: bool | None = None

    result = payload.get("result")

    if isinstance(result, dict):

        if "error" in result or "failure" in result:

            failed = True

        elif "success" in result:

            failed = False

        shell_result = result.get("success") or result.get("failure")

        if isinstance(shell_result, dict):

            if command is None and isinstance(shell_result.get("command"), str):

                command = unwrap_shell_command(shell_result["command"])

            exit_code = shell_result.get("exitCode")

            if isinstance(exit_code, int) and not isinstance(exit_code, bool):

                failed = exit_code != 0

    return tool_name, command, failed

def iter_tool_observations(event: Any) -> Iterable[tuple[bool, str]]:

    """Yield (failed, command_or_name) for adapter-visible tool executions with a known outcome."""

    cursor_call = cursor_tool_call(event)

    if cursor_call is not None:

        tool_name, command, failed = cursor_call

        if str(event.get("subtype") or "").lower() == "completed" and failed is not None:

            yield failed, command or tool_name

        return

    for item in walk_json(event):

        if not isinstance(item, dict):

            continue

        item_type = str(item.get("type") or "").lower()

        if item_type == "tool_result":

            yield bool(item.get("is_error")), str(item.get("tool_use_id") or "tool_result")

            continue

        exit_code = item.get("exit_code")

        if exit_code is None and item_type == "command_execution":

            status = str(item.get("status") or "").lower()

            if status == "failed":

                exit_code = 1

            elif status == "completed":

                exit_code = 0

        if exit_code is not None and not isinstance(exit_code, bool) and isinstance(exit_code, int):

            command = item.get("command") or item.get("cmd") or item_type or "command"

            yield exit_code != 0, unwrap_shell_command(str(command))

def event_output_text(event: Any) -> str:

    texts: list[str] = []

    for item in walk_json(event):

        if not isinstance(item, dict):

            continue

        for key in ("aggregated_output", "output", "stderr", "error", "content", "text"):

            value = item.get(key)

            if isinstance(value, str):

                texts.append(value)

    return "\n".join(texts)

class StreamAdapter:

    """Per-CLI stream knowledge. The base class implements the shared heuristics."""

    name = "generic"

    # Claude streams emit explicit compaction markers, so absence there is a

    # measured zero. Streams without such markers leave truncation unmeasured.

    missing_compaction_is_measured_zero = False

    # OpenAI-style usage reports cached tokens as a SUBSET of input_tokens;

    # Anthropic-style reports them separately. Adding cache reads to input is

    # double counting for the former.

    cache_read_included_in_input = False

    def is_turn_event(self, event: Any) -> bool:

        return is_assistant_turn_event(event)

    def is_session_usage_event(self, event: Any) -> bool:

        """Events that carry session-level usage without marking a model turn."""

        return False

    def turn_event_key(self, event: Any) -> str | None:

        """Stable key for streams that repeat one logical turn across events."""

        return None

    def stream_carries_tool_use_blocks(self, events: list[Any]) -> bool:

        # Claude-style streams carry a turn's tool_use blocks inside the

        # assistant event itself (results follow); Codex-style streams emit

        # turn boundaries after the turn's command items. Decide per stream,

        # not per event — mixing the two heuristics double-counts text-only turns.

        return any(

            self.is_turn_event(event) and assistant_event_tool_use_count(event) > 0 for event in events

        )

    def turn_tool_calls(self, event: Any, tools_since_last_turn: int, carries_tool_use: bool) -> int:

        if carries_tool_use:

            return assistant_event_tool_use_count(event)

        return tools_since_last_turn

class ClaudeStreamAdapter(StreamAdapter):

    name = "claude_print_stream_json"

    missing_compaction_is_measured_zero = True

    def turn_event_key(self, event: Any) -> str | None:

        if not isinstance(event, dict):

            return None

        message = event.get("message")

        if not isinstance(message, dict):

            return None

        message_id = message.get("id")

        return str(message_id) if message_id else None

class CodexStreamAdapter(StreamAdapter):

    """codex exec --json, verified against codex-cli 0.139 transcripts.

    Codex's `turn.completed` covers the WHOLE session (one per run) and its

    `input_tokens` already includes the cached subset (OpenAI semantics).

    Model round-trips are the `agent_message` items, which carry no usage —

    so turns come from agent_message and usage from the session event.

    """

    name = "codex_exec_json"

    cache_read_included_in_input = True

    def is_turn_event(self, event: Any) -> bool:

        if not isinstance(event, dict):

            return False

        if str(event.get("type") or "").lower() != "item.completed":

            return False

        item = event.get("item")

        return isinstance(item, dict) and str(item.get("type") or "") in {"agent_message", "assistant_message"}

    def is_session_usage_event(self, event: Any) -> bool:

        if not isinstance(event, dict):

            return False

        return str(event.get("type") or "").lower() in {"turn.completed", "turn_completed"}

class CursorStreamAdapter(StreamAdapter):

    name = "cursor_agent_stream_json"

ADAPTERS: dict[str, StreamAdapter] = {

    adapter.name: adapter

    for adapter in (ClaudeStreamAdapter(), CodexStreamAdapter(), CursorStreamAdapter())

}

def adapter_for(adapter_name: str) -> StreamAdapter:

    return ADAPTERS.get(adapter_name, StreamAdapter())

def analyze_stream(events: list[Any], adapter_name: str) -> StreamAnalysis:

    adapter = adapter_for(adapter_name)

    analysis = StreamAnalysis()

    if not events:

        return analysis

    carries_tool_use = adapter.stream_carries_tool_use_blocks(events)

    turns: list[dict[str, Any]] = []

    session_usages: list[dict[str, int | None]] = []

    # Ordered trace mixing turn boundaries and tool outcomes so recovery cost

    # can be measured as turns between first failure and next success.

    trace: list[tuple[str, Any]] = []

    truncations = 0

    tools_since_last_turn = 0

    tool_use_seen = False

    failed_examples: list[str] = []

    failed_total = 0

    failed_vcs = 0

    failed_test = 0

    unknown_failures = 0

    failed_commands_seen: dict[str, int] = {}

    same_command_retries = 0

    # Compact-output cost: a successful read-only VCS command immediately

    # re-asked in the same family with different arguments means the first

    # output lacked the information the agent needed. Claude-style streams

    # report results by tool_use_id, so the actual command text comes from a

    # pre-pass mapping ids to tool_use input commands; when no observation

    # resolves to real command text the metric stays None (ADR-0002), never a

    # fabricated zero.

    vcs_info_followups = 0

    previous_ok_info: tuple[str, str] | None = None

    commands_resolvable = False

    id_to_command: dict[str, str] = {}

    turn_key_to_index: dict[str, int] = {}

    turn_tool_refs: dict[str, set[str]] = {}

    for event in events:

        for item in walk_json(event):

            if not isinstance(item, dict) or str(item.get("type") or "") != "tool_use":

                continue

            tool_id = item.get("id")

            command = (item.get("input") or {}).get("command") if isinstance(item.get("input"), dict) else None

            if tool_id and isinstance(command, str):

                id_to_command[str(tool_id)] = unwrap_shell_command(command)

    for event in events:

        if isinstance(event, dict):

            event_type = str(event.get("type") or "").lower()

            subtype = str(event.get("subtype") or "").lower()

            if "compact" in event_type or "compact" in subtype or "truncat" in event_type:

                truncations += 1

        observed_tools = list(iter_tool_observations(event))

        for failed, command in observed_tools:

            tool_use_seen = True

            tools_since_last_turn += 1

            trace.append(("tool_fail" if failed else "tool_ok", command))

            resolved_command = id_to_command.get(command, command)

            if " " in resolved_command:

                commands_resolvable = True

            if not failed:

                family = vcs_info_family(resolved_command)

                if (

                    family is not None

                    and previous_ok_info is not None

                    and previous_ok_info[0] == family

                    and previous_ok_info[1] != resolved_command

                ):

                    vcs_info_followups += 1

                previous_ok_info = (family, resolved_command) if family is not None else None

            else:

                previous_ok_info = None

            if failed:

                failed_total += 1

                kinds = classify_command(resolved_command)

                if "vcs" in kinds:

                    failed_vcs += 1

                if "test" in kinds:

                    failed_test += 1

                if UNKNOWN_COMMAND_PATTERN.search(event_output_text(event)):

                    unknown_failures += 1

                if len(failed_examples) < 5:

                    failed_examples.append(resolved_command[:160])

                previous = failed_commands_seen.get(resolved_command, 0)

                if previous:

                    same_command_retries += 1

                failed_commands_seen[resolved_command] = previous + 1

        if adapter.is_session_usage_event(event):

            usage = usage_from_subtree(event)

            if usage is not None:

                session_usages.append(usage)

        if adapter.is_turn_event(event):

            turn_key = adapter.turn_event_key(event)

            tool_refs = assistant_event_tool_use_refs(event)

            if turn_key is not None:

                if turn_key in turn_key_to_index:

                    turn_index = turn_key_to_index[turn_key]

                    if carries_tool_use:

                        refs = turn_tool_refs.setdefault(turn_key, set())

                        refs.update(tool_refs)

                        turns[turn_index]["tool_calls_observed"] = len(refs)

                    usage = usage_from_subtree(event)

                    if usage is not None:

                        turn = turns[turn_index]

                        if turn["input_tokens"] is None:

                            turn["input_tokens"] = usage.get("input_tokens")

                        if turn["output_tokens"] is None:

                            turn["output_tokens"] = usage.get("output_tokens")

                        if turn["cache_read_tokens"] is None:

                            turn["cache_read_tokens"] = usage.get("cache_read_tokens")

                    continue

                turn_key_to_index[turn_key] = len(turns)

                if carries_tool_use:

                    turn_tool_refs[turn_key] = set(tool_refs)

            usage = usage_from_subtree(event)

            turns.append(

                {

                    "turn": len(turns),

                    "input_tokens": usage.get("input_tokens") if usage else None,

                    "output_tokens": usage.get("output_tokens") if usage else None,

                    "cache_read_tokens": usage.get("cache_read_tokens") if usage else None,

                    "tool_calls_observed": len(tool_refs)

                    if carries_tool_use

                    else adapter.turn_tool_calls(event, tools_since_last_turn, carries_tool_use),

                }

            )

            tools_since_last_turn = 0

            trace.append(("turn", len(turns) - 1))

    if turns:

        analysis.turn_source = "adapter_stream"

        analysis.turns = turns

        analysis.assistant_turns_total = len(turns)

        tool_counts = [int(turn["tool_calls_observed"]) for turn in turns]

        if tool_use_seen or any(tool_counts):

            analysis.tool_calls_per_turn_avg = round(sum(tool_counts) / len(turns), 3)

            analysis.max_tool_calls_per_turn = max(tool_counts)

        def input_total(usage: dict[str, Any]) -> int:

            base = usage["input_tokens"] or 0

            if adapter.cache_read_included_in_input:

                return base

            return base + (usage["cache_read_tokens"] or 0)

        input_totals = [input_total(turn) for turn in turns if turn["input_tokens"] is not None]

        if not input_totals and session_usages:

            # Session-level usage (codex-style): already cumulative across the

            # run's model calls; peak per round-trip is not adapter-visible.

            input_totals = [

                input_total(usage) for usage in session_usages if usage.get("input_tokens") is not None

            ]

            if input_totals:

                analysis.cumulative_input_tokens = sum(input_totals)

        elif input_totals:

            analysis.cumulative_input_tokens = sum(input_totals)

            analysis.peak_input_tokens = max(input_totals)

        if truncations or adapter.missing_compaction_is_measured_zero:

            analysis.context_truncation_events = truncations

    if tool_use_seen:

        analysis.failed_tool_calls_total = failed_total

        analysis.failed_vcs_commands_total = failed_vcs

        analysis.failed_test_commands_total = failed_test

        analysis.same_command_retries_total = same_command_retries

        analysis.unknown_command_failures_total = unknown_failures

        analysis.vcs_info_followup_calls_total = vcs_info_followups if commands_resolvable else None

        analysis.failed_command_examples = failed_examples

        first_failure_index = next((i for i, (kind, _) in enumerate(trace) if kind == "tool_fail"), None)

        if first_failure_index is not None:

            recovery_turns = 0

            recovered = False

            recovery_turn_ids: list[int] = []

            for kind, payload in trace[first_failure_index + 1 :]:

                if kind == "turn":

                    recovery_turns += 1

                    recovery_turn_ids.append(int(payload))

                elif kind == "tool_ok":

                    recovered = True

                    break

            if recovered:

                analysis.turns_to_recovery = recovery_turns

                if turns and recovery_turn_ids:

                    outputs = [

                        turns[turn_id]["output_tokens"]

                        for turn_id in recovery_turn_ids

                        if turn_id < len(turns) and turns[turn_id]["output_tokens"] is not None

                    ]

                    if outputs:

                        analysis.tokens_to_recovery = sum(outputs)

    return analysis