"""Secret-safe Claude Code adapter for the IOP benchmark lifecycle.

The generic lifecycle deliberately does not know Claude's command line or its
JSONL protocol.  This module converts one immutable benchmark cell and one
prepared workspace into that closed boundary.  It has no network dependency;
tests use a local fake executable.
"""

from __future__ import annotations

import json
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

from scripts.agent_benchmark.connectivity import (
    CallerCapability,
    ConnectivityResult,
    EffectiveBinding,
    RequestedEffectiveBinding,
    make_result,
)
from scripts.agent_benchmark.lifecycle import (
    CALLER_REASON_ERROR,
    CALLER_REASON_SUCCESS,
    CALLER_STATUS_FAILED,
    CALLER_STATUS_SUCCEEDED,
    COMPLETION_EXIT_AFTER_IDLE,
    SUBMISSION_STDIN_ONCE,
    CallerEvent,
    CallerTerminal,
    InvocationSpec,
    LifecycleMetricError,
    ParsedMetric,
    TLS_CA_ENV_KEYS,
    count_metric,
    duration_metric,
    exact_value_redactor,
    inherited_tls_ca_environment,
    is_reported_number,
)
from scripts.agent_benchmark.manifest import MatrixCell, Timeout
from scripts.agent_benchmark.workspace import PreparedWorkspace


REDACTED = "[redacted]"
CLAUDE_ROUTE_KINDS = ("direct", "execution_preset")
# This is lexical order, as required by the closed connectivity capability
# tuple.  The cell's requested effort is still passed through unchanged.
CLAUDE_EFFORTS = ("high", "low", "max", "medium", "xhigh")
_STRUCTURAL_SECRET_KEYS = frozenset(
    {
        "content", "text", "input", "arguments", "tool_input", "prompt", "query",
        "result", "error", "errors", "error_message",
    }
)
# The exact Claude result fields this adapter is allowed to observe.  Anything
# outside these allowlists never becomes timing or usage evidence, and
# ``duration_api_ms`` is marked as overlapping because it is reported inside
# the same wall-clock window as ``duration_ms``.
_CLAUDE_DURATION_FIELDS = (
    ("duration_ms", "total_duration", False),
    ("duration_api_ms", "model_duration", True),
)
_CLAUDE_USAGE_FIELDS = {
    "input_tokens": "input_tokens",
    "output_tokens": "output_tokens",
    "cache_read_input_tokens": "cached_input_tokens",
    "cache_creation_input_tokens": "cache_write_tokens",
}
_CLAUDE_USAGE_BOOKKEEPING_FIELDS = {
    "cache_creation", "inference_geo", "iterations", "output_tokens_details",
    "server_tool_use", "service_tier", "speed",
}
_CLAUDE_USAGE_NESTED_COUNTS = {
    "cache_creation": {"ephemeral_1h_input_tokens", "ephemeral_5m_input_tokens"},
    "output_tokens_details": {"thinking_tokens"},
    "server_tool_use": {"web_fetch_requests", "web_search_requests"},
}


class ClaudeIopError(Exception):
    """Base exception for configuration and protocol failures."""


class ClaudeIopValidationError(ClaudeIopError):
    """Raised when a runtime input or preflight shape is inadmissible."""


class ClaudeIopProtocolError(ClaudeIopError):
    """Raised for a malformed or contradictory claimed Claude terminal."""


@dataclass(frozen=True)
class ClaudeIopRuntime:
    """Runtime-only IOP inputs.  Values are never written to durable evidence."""

    binary: str
    base_url: str
    api_key: str


def _require_string(value: Any, name: str) -> str:
    if not isinstance(value, str) or not value:
        raise ClaudeIopValidationError(f"invalid {name}")
    return value


def _exact_object(raw_line: str) -> dict[str, Any]:
    if not isinstance(raw_line, str):
        raise ClaudeIopProtocolError("invalid Claude JSONL line")
    try:
        value = json.loads(raw_line)
    except (TypeError, ValueError) as exc:
        raise ClaudeIopProtocolError("invalid Claude JSONL line") from exc
    if not isinstance(value, dict):
        raise ClaudeIopProtocolError("invalid Claude JSONL object")
    return value


def _required_string(data: dict[str, Any], name: str) -> str:
    value = data.get(name)
    if not isinstance(value, str) or not value:
        raise ClaudeIopProtocolError(f"missing Claude {name}")
    return value


def _structural_redact(
    value: Any,
    sensitive_values: tuple[str, ...],
    key: str = "",
    *,
    redact_terminal_message: bool = False,
) -> Any:
    if key in _STRUCTURAL_SECRET_KEYS:
        return REDACTED
    if redact_terminal_message and key == "message":
        return REDACTED
    if isinstance(value, str):
        for sensitive in sensitive_values:
            if sensitive:
                value = value.replace(sensitive, REDACTED)
        return value
    if isinstance(value, list):
        return [
            _structural_redact(
                item, sensitive_values, redact_terminal_message=redact_terminal_message
            )
            for item in value
        ]
    if isinstance(value, dict):
        return {
            str(name): _structural_redact(
                item,
                sensitive_values,
                str(name),
                redact_terminal_message=redact_terminal_message,
            )
            for name, item in value.items()
        }
    return value


def redact_claude_event(raw_line: str, sensitive_values: tuple[str, ...]) -> str:
    """Return canonical JSON without task/tool content or runtime secrets.

    A malformed line is represented by a fixed marker so error reporting cannot
    accidentally retain the raw malformed payload.
    """
    try:
        event = _exact_object(raw_line)
    except ClaudeIopProtocolError:
        return '{"type":"invalid_claude_json"}'
    redacted = _structural_redact(
        event,
        sensitive_values,
        redact_terminal_message=event.get("type") in ("error", "result"),
    )
    return json.dumps(redacted, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def parse_preflight_binding(raw_line: str, cell: MatrixCell) -> RequestedEffectiveBinding:
    """Parse an explicit, secret-free IOP binding observation from JSONL.

    The adapter never derives effective values from the requested cell.  A
    caller must provide all requested/effective scalars and every stage, or the
    connectivity contract rejects the observation as incomplete.
    """
    event = _exact_object(raw_line)
    if event.get("type") != "system" or event.get("subtype") != "iop_binding":
        raise ClaudeIopProtocolError("invalid Claude preflight event")
    binding = event.get("binding")
    if not isinstance(binding, dict) or set(binding) != {
        "cell_id", "caller", "requested_route_kind", "requested_route_id",
        "requested_model", "requested_effort", "effective_route_kind",
        "effective_route_id", "effective_model", "effective_effort", "effective_bindings",
    }:
        raise ClaudeIopProtocolError("invalid Claude preflight binding")
    raw_stages = binding["effective_bindings"]
    if not isinstance(raw_stages, list):
        raise ClaudeIopProtocolError("invalid Claude preflight stages")
    stages: list[EffectiveBinding] = []
    for item in raw_stages:
        if not isinstance(item, dict) or set(item) != {"stage", "model", "effort"}:
            raise ClaudeIopProtocolError("invalid Claude preflight stage")
        stages.append(EffectiveBinding(item["stage"], item["model"], item["effort"]))
    try:
        result = RequestedEffectiveBinding(
            binding["cell_id"], binding["caller"], binding["requested_route_kind"],
            binding["requested_route_id"], binding["requested_model"],
            binding["requested_effort"], binding["effective_route_kind"],
            binding["effective_route_id"], binding["effective_model"],
            binding["effective_effort"], tuple(stages),
        )
        # Validation is intentionally delegated to the one shared contract.
        make_result(cell, claude_capability(), result)
    except Exception as exc:
        if isinstance(exc, ClaudeIopProtocolError):
            raise
        raise ClaudeIopProtocolError("invalid Claude preflight binding") from exc
    return result


class ClaudeStreamParser:
    """Parse one fresh Claude Code stream without inventing terminal bindings."""

    def __init__(self, cell: MatrixCell, session_id: str) -> None:
        if not isinstance(cell, MatrixCell) or cell.caller != "claude":
            raise ClaudeIopValidationError("Claude adapter requires a Claude matrix cell")
        self.cell = cell
        # The prepared workspace identity proves that this is a fresh caller
        # invocation.  Claude Code emits its own UUID in system/init, so it
        # must be derived from that event rather than compared to this local
        # opaque label.
        self.prepared_session_id = _require_string(session_id, "session_id")
        self.claude_session_id: str | None = None
        self._phase = "await_init"
        self._assistant_messages = 0
        self._active_message_id: str | None = None
        self._completed_message_ids: set[str] = set()
        self._continuation_pending = False
        self._api_error_seen = False

    def _require_bound_session(self, event: dict[str, Any]) -> None:
        if self.claude_session_id is None:
            raise ClaudeIopProtocolError("missing Claude init")
        if _required_string(event, "session_id") != self.claude_session_id:
            raise ClaudeIopProtocolError("Claude session binding mismatch")

    def _consume_init(self, event: dict[str, Any]) -> None:
        if self._phase != "await_init":
            raise ClaudeIopProtocolError("duplicate or out-of-order Claude init")
        if _required_string(event, "model") != self.cell.iop.request_model:
            raise ClaudeIopProtocolError("Claude model binding mismatch")
        self.claude_session_id = _required_string(event, "session_id")
        self._phase = "await_assistant"

    def _consume_assistant(self, event: dict[str, Any]) -> CallerEvent | None:
        if self._phase != "await_assistant":
            raise ClaudeIopProtocolError("duplicate or out-of-order Claude assistant")
        self._require_bound_session(event)
        message = event.get("message")
        if not isinstance(message, dict):
            raise ClaudeIopProtocolError("invalid Claude assistant event")
        if event.get("is_api_error_message") is True:
            if message.get("model") != "<synthetic>" or message.get("stop_reason") != "stop_sequence":
                raise ClaudeIopProtocolError("invalid Claude API error event")
            self._complete_active_message()
            self._api_error_seen = True
            self._phase = "await_error_result"
            return None
        if _required_string(message, "model") != self.cell.iop.request_model:
            raise ClaudeIopProtocolError("Claude model binding mismatch")
        stop_reason = message.get("stop_reason")
        if stop_reason not in (None, "tool_use", "end_turn"):
            raise ClaudeIopProtocolError("invalid Claude assistant stop reason")
        message_id = message.get("id")
        if message_id is not None and (not isinstance(message_id, str) or not message_id):
            raise ClaudeIopProtocolError("invalid Claude assistant message id")
        if self._continuation_pending:
            if message_id != self._active_message_id:
                self._complete_active_message()
            self._continuation_pending = False
        if stop_reason is None:
            if message_id is None:
                raise ClaudeIopProtocolError("unbound Claude assistant snapshot")
            if message_id in self._completed_message_ids:
                raise ClaudeIopProtocolError("duplicate Claude assistant message")
            if self._active_message_id not in (None, message_id):
                raise ClaudeIopProtocolError("overlapping Claude assistant messages")
            self._active_message_id = message_id
            return None

        if message_id is None:
            # Older fixture-shaped output did not include a message id. It is
            # admissible only for the single final assistant event.
            if stop_reason != "end_turn" or self._active_message_id is not None:
                raise ClaudeIopProtocolError("unbound Claude assistant terminal")
            self._assistant_messages += 1
        else:
            if message_id in self._completed_message_ids:
                raise ClaudeIopProtocolError("duplicate Claude assistant message")
            if self._active_message_id not in (None, message_id):
                raise ClaudeIopProtocolError("overlapping Claude assistant messages")
            self._active_message_id = message_id
            self._complete_active_message()
        if stop_reason == "tool_use":
            self._phase = "await_tool_result"
            return None
        self._phase = "await_result"
        return CallerEvent("finish")

    def _consume_user(self, event: dict[str, Any]) -> None:
        if self._phase not in ("await_assistant", "await_tool_result"):
            raise ClaudeIopProtocolError("out-of-order Claude user event")
        self._require_bound_session(event)
        if self._active_message_id is not None:
            # Claude Code may publish cumulative snapshots with the same
            # assistant message id on both sides of one or more tool results.
            # The next assistant snapshot (or the result terminal) determines
            # whether this message continues or a new model call begins.
            self._continuation_pending = True
        elif self._phase != "await_tool_result":
            raise ClaudeIopProtocolError("unexpected Claude user event")
        self._phase = "await_assistant"

    def _complete_active_message(self) -> None:
        if self._active_message_id is None:
            return
        if self._active_message_id in self._completed_message_ids:
            raise ClaudeIopProtocolError("duplicate Claude assistant message")
        self._completed_message_ids.add(self._active_message_id)
        self._active_message_id = None
        self._continuation_pending = False
        self._assistant_messages += 1

    def _consume_result(self, event: dict[str, Any]) -> tuple[Any, ...] | None:
        if self._phase == "await_error_result":
            self._require_bound_session(event)
            if (
                not self._api_error_seen
                or event.get("subtype") != "success"
                or event.get("is_error") is not True
                or event.get("terminal_reason") != "api_error"
            ):
                raise ClaudeIopProtocolError("invalid Claude API error terminal")
            self._phase = "complete"
            return (
                CallerTerminal(CALLER_STATUS_FAILED, CALLER_REASON_ERROR),
                CallerEvent("finish"),
                CallerEvent("idle"),
            )
        result_completes_active_message = False
        if self._phase == "await_assistant" and self._active_message_id is not None:
            if self._continuation_pending:
                raise ClaudeIopProtocolError("Claude result followed an unresolved tool result")
            self._complete_active_message()
            result_completes_active_message = True
        elif self._phase != "await_result":
            raise ClaudeIopProtocolError("duplicate or out-of-order Claude result")
        self._require_bound_session(event)
        if event.get("subtype") != "success" or event.get("is_error") is True:
            raise ClaudeIopProtocolError("invalid Claude result terminal")
        self._phase = "complete"
        terminal = (
            (CallerEvent("finish"), CallerEvent("idle"))
            if result_completes_active_message
            else (CallerEvent("idle"),)
        )
        return (
            *self._observations(event),
            CallerTerminal(CALLER_STATUS_SUCCEEDED, CALLER_REASON_SUCCESS),
            *terminal,
        )

    def _observations(self, event: dict[str, Any]) -> tuple[ParsedMetric, ...]:
        """Convert only allowlisted reported Claude values into observations."""
        model = self.cell.iop.request_model
        observations: list[ParsedMetric] = []
        try:
            for field, name, overlap in _CLAUDE_DURATION_FIELDS:
                if field in event:
                    if not is_reported_number(event[field]):
                        raise ClaudeIopProtocolError("invalid Claude duration observation")
                    observations.append(duration_metric(
                        name, event[field], reported_unit="ms",
                        model=model, overlap=overlap,
                    ))
            observations.append(count_metric(
                "model_calls", self._assistant_messages, model=model
            ))
            observations.extend(self._usage_observations(event, model))
        except LifecycleMetricError as exc:
            raise ClaudeIopProtocolError("invalid Claude usage observation") from exc
        return tuple(observations)

    @staticmethod
    def _usage_observations(event: dict[str, Any], model: str) -> list[ParsedMetric]:
        usage = event.get("usage")
        if usage is None:
            return []
        allowed = set(_CLAUDE_USAGE_FIELDS) | _CLAUDE_USAGE_BOOKKEEPING_FIELDS
        if not isinstance(usage, dict) or not set(usage) <= allowed:
            raise ClaudeIopProtocolError("invalid Claude usage observation")
        ClaudeStreamParser._validate_usage_bookkeeping(usage)
        return [
            count_metric(_CLAUDE_USAGE_FIELDS[field], value, model=model)
            for field, value in sorted(usage.items()) if field in _CLAUDE_USAGE_FIELDS
        ]

    @staticmethod
    def _validate_usage_bookkeeping(usage: dict[str, Any]) -> None:
        for field, nested_fields in _CLAUDE_USAGE_NESTED_COUNTS.items():
            if field not in usage:
                continue
            value = usage[field]
            if not isinstance(value, dict) or not set(value) <= nested_fields:
                raise ClaudeIopProtocolError("invalid Claude usage observation")
            if any(not is_reported_number(item) for item in value.values()):
                raise ClaudeIopProtocolError("invalid Claude usage observation")
        if "iterations" in usage and usage["iterations"] != []:
            raise ClaudeIopProtocolError("invalid Claude usage observation")
        for field in ("inference_geo", "service_tier", "speed"):
            if field in usage and (
                not isinstance(usage[field], str) or len(usage[field]) > 64
            ):
                raise ClaudeIopProtocolError("invalid Claude usage observation")

    def __call__(self, stream: str, raw_line: str) -> Any:
        if stream != "stdout":
            return None
        event = _exact_object(raw_line)
        event_type = event.get("type")
        if event_type == "system" and event.get("subtype") == "init":
            self._consume_init(event)
            return None
        if event_type == "assistant":
            return self._consume_assistant(event)
        if event_type == "user":
            self._consume_user(event)
            return None
        if event_type == "result":
            return self._consume_result(event)
        # Informational events are deliberately ignored only after they have
        # passed exact JSON-object decoding above.
        return None


def claude_capability() -> CallerCapability:
    return CallerCapability("claude", CLAUDE_ROUTE_KINDS, CLAUDE_EFFORTS)


def resolve_claude_binary(binary: str) -> str:
    """Resolve one executable before lifecycle process creation."""
    candidate = _require_string(binary, "Claude binary")
    path = Path(candidate)
    resolved = str(path.resolve()) if path.parent != Path(".") else shutil.which(candidate)
    if not resolved or not Path(resolved).is_file() or not os.access(resolved, os.X_OK):
        raise ClaudeIopValidationError("Claude binary is unavailable")
    return str(Path(resolved).resolve())


class ClaudeIopAdapter:
    """Build safe Claude invocations and parse their IOP-bound stream evidence."""

    capability = claude_capability()

    def __init__(self, cell: MatrixCell, workspace: PreparedWorkspace, runtime: ClaudeIopRuntime) -> None:
        if not isinstance(cell, MatrixCell) or cell.caller != "claude":
            raise ClaudeIopValidationError("Claude adapter requires a Claude matrix cell")
        if not isinstance(workspace, PreparedWorkspace) or not workspace.session_is_fresh:
            raise ClaudeIopValidationError("Claude adapter requires a fresh prepared workspace")
        if not isinstance(runtime, ClaudeIopRuntime):
            raise ClaudeIopValidationError("invalid Claude runtime")
        self.cell = cell
        self.workspace = workspace
        self.runtime = runtime
        self.binary = resolve_claude_binary(runtime.binary)
        self.base_url = _require_string(runtime.base_url, "IOP base URL")
        self.api_key = _require_string(runtime.api_key, "IOP API key")

    def preflight(self, raw_line: str) -> ConnectivityResult:
        return make_result(self.cell, self.capability, parse_preflight_binding(raw_line, self.cell))

    def parser(self) -> ClaudeStreamParser:
        return ClaudeStreamParser(self.cell, self.workspace.session_id)

    def redactor(self, task: str) -> Callable[[str], str]:
        _require_string(task, "task")
        structural_values = (task, self.base_url, self.api_key)
        exact = exact_value_redactor(structural_values)

        def _redact(line: str) -> str:
            return exact(redact_claude_event(line, structural_values))

        return _redact

    def invocation(self, task: str, evidence_dir: str | Path, timeout: Timeout) -> InvocationSpec:
        _require_string(task, "task")
        if not isinstance(timeout, Timeout):
            raise ClaudeIopValidationError("invalid invocation timeout")
        evidence_path = _require_string(str(evidence_dir), "evidence directory")
        cwd = Path(self.workspace.workspace_dir)
        if not cwd.is_dir():
            raise ClaudeIopValidationError("prepared workspace is unavailable")
        env = (
            ("PATH", os.environ.get("PATH", "/usr/bin:/bin")),
            ("ANTHROPIC_BASE_URL", self.base_url),
            ("ANTHROPIC_API_KEY", self.api_key),
            ("CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1"),
            ("CLAUDE_CODE_DISABLE_AUTOUPDATER", "1"),
        ) + inherited_tls_ca_environment()
        return InvocationSpec(
            argv=(
                self.binary, "--bare", "--print", "--verbose",
                "--input-format", "text", "--output-format", "stream-json",
                "--model", self.cell.iop.request_model, "--effort", self.cell.iop.requested_effort,
                "--no-session-persistence", "--permission-mode", "dontAsk",
                "--tools", "Read,Write,Edit", "--allowedTools", "Read,Write,Edit",
            ),
            cwd=str(cwd),
            env=env,
            env_allowlist=(
                "ANTHROPIC_BASE_URL", "ANTHROPIC_API_KEY",
                "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "CLAUDE_CODE_DISABLE_AUTOUPDATER",
                *TLS_CA_ENV_KEYS,
            ),
            submission_mode=SUBMISSION_STDIN_ONCE,
            completion_mode=COMPLETION_EXIT_AFTER_IDLE,
            timeout=timeout,
            evidence_dir=evidence_path,
            task_payload=task.encode("utf-8"),
        )