"""Closed agy-to-IOP adapter for the comparison benchmark.

The adapter deliberately recognises only a documented agy transport contract.
In particular, it never inherits an ambient Gemini configuration: absent or
unknown transport support is an implementation gap before a caller process is
constructed.  The module is standard-library-only and has no network calls.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
from urllib.parse import urlsplit

from scripts.agent_benchmark.connectivity import (
    ISSUE_RESUME_CODES,
    CallerCapability,
    ConnectivityResult,
    ConnectivityIssue,
    RequestedEffectiveBinding,
    classify_issues,
    make_result,
)
from scripts.agent_benchmark.lifecycle import (
    COMPLETION_EXIT_AFTER_IDLE,
    SUBMISSION_ARGV_TASK,
    InvocationResult,
    InvocationSpec,
    LifecycleMetricError,
    ParsedMetric,
    SupervisorLocator,
    TLS_CA_ENV_KEYS,
    count_metric,
    duration_metric,
    env_pairs,
    exact_value_redactor,
    inherited_tls_ca_environment,
    run_invocation,
)
from scripts.agent_benchmark.manifest import MatrixCell, Timeout
from scripts.agent_benchmark.workspace import PreparedWorkspace


AGY_CALLER = "agy"
AGY_KNOWN_VERSION = "1.1.12"
AGY_ENDPOINT_ENV = "GOOGLE_GEMINI_BASE_URL"
AGY_AUTH_ENV = "GEMINI_API_KEY"
AGY_SETTINGS_RELATIVE_PATH = Path(".gemini/antigravity-cli/settings.json")
_AGY_ISOLATED_SETTINGS = {
    "enableTelemetry": False,
    "modelProvider": "gemini",
    "toolPermission": "always-proceed",
}
AGY_MODEL_LABELS = {
    "gemini-3.6-flash": "Gemini 3.6 Flash",
    "gemini-hybrid": "Gemini 3.6 Flash",
}
_VERSION_RE = re.compile(r"(?:agy\s+)?(\d+\.\d+\.\d+)", re.IGNORECASE)
_SAFE_EVENT_FIELDS = ("event", "state", "step_type", "status")
_DOCUMENTED_OPTIONS = ("--print", "--output-format", "--sandbox", "--model")
_AGY_USAGE_METRICS = {
    "input_tokens": "input_tokens",
    "cache_read_tokens": "cached_input_tokens",
    "output_tokens": "output_tokens",
    "thinking_tokens": "reasoning_tokens",
    "total_tokens": "total_tokens",
}
_IDENTITY_RE = re.compile(r"sha256:[0-9a-f]{64}\Z")


def _exact_token_present(text: str, token: str) -> bool:
    """Match one documented help token, never a prefix or a suffix."""
    return re.search(r"(?<![A-Za-z0-9_-])" + re.escape(token) + r"(?![A-Za-z0-9_-])", text) is not None


@dataclass(frozen=True)
class AgyRuntimeInputs:
    """Private runtime values supplied by a preflight owner, never persisted."""

    binary: str
    endpoint: str
    credential: str


@dataclass(frozen=True)
class AgyRuntimeObservation:
    """Secret-free IOP configuration proof for exactly one benchmark cell."""

    cell_id: str
    route_kind: str
    route_id: str
    endpoint_identity: str
    credential_identity: str
    config_identity: str


@dataclass(frozen=True)
class _ValidatedAgyRuntime:
    """Private launch values admitted only after IOP configuration validation."""

    binary: str
    endpoint: str
    credential: str
    observation: AgyRuntimeObservation


@dataclass(frozen=True)
class AgyDocumentedCapabilities:
    """Exact, known-version tokens parsed from public agy help output."""

    options: tuple[str, ...]
    environment: tuple[str, ...]
    output_formats: tuple[str, ...]


@dataclass(frozen=True)
class AgyCapability:
    """Versioned, documented capability observation; no caller is launched."""

    version: str | None
    iop_transport_supported: bool
    endpoint_supported: bool
    auth_supported: bool
    protocol_supported: bool
    stream_supported: bool
    route_kinds: tuple[str, ...]
    efforts: tuple[str, ...]


@dataclass(frozen=True)
class AgyPreflightResult:
    """Closed outcome before invocation construction."""

    capability: AgyCapability
    binding: RequestedEffectiveBinding
    issues: tuple[ConnectivityIssue, ...]
    status: str
    runtime: _ValidatedAgyRuntime | None


class AgyAdapterError(Exception):
    """Raised when a caller launch is requested without a proven transport."""


def _issue(code: str) -> ConnectivityIssue:
    return ConnectivityIssue(code, ISSUE_RESUME_CODES[code])


def _requested_binding(cell: MatrixCell) -> RequestedEffectiveBinding:
    return RequestedEffectiveBinding(
        cell.id,
        cell.caller,
        cell.iop.route_kind,
        cell.iop.route_id,
        cell.iop.request_model,
        cell.iop.requested_effort,
    )


def parse_documented_agy_capabilities(help_output: str) -> AgyDocumentedCapabilities:
    """Parse only complete documented tokens from a known agy help surface."""
    if not isinstance(help_output, str):
        return AgyDocumentedCapabilities((), (), ())
    return AgyDocumentedCapabilities(
        tuple(token for token in _DOCUMENTED_OPTIONS if _exact_token_present(help_output, token)),
        (),
        ("stream-json",) if _exact_token_present(help_output, "stream-json") else (),
    )


def inspect_agy_iop_capability(version_output: str, help_output: str) -> AgyCapability:
    """Inspect only public, versioned help text for the closed IOP transport.

    The known release is pinned because its Gemini provider environment is not
    printed by ``--help``. The public options and stream format still have to
    match exactly; a changed release fails closed until re-qualified.
    """
    if not isinstance(version_output, str) or not isinstance(help_output, str):
        return AgyCapability(None, False, False, False, False, False, (), ())
    matched = _VERSION_RE.fullmatch(version_output.strip())
    version = matched.group(1) if matched else None
    known_version = version == AGY_KNOWN_VERSION
    documented = parse_documented_agy_capabilities(help_output)
    endpoint_supported = known_version
    auth_supported = known_version
    protocol_supported = known_version and all(
        option in documented.options for option in _DOCUMENTED_OPTIONS
    )
    stream_supported = "stream-json" in documented.output_formats
    supported = endpoint_supported and auth_supported and protocol_supported and stream_supported
    if not supported:
        return AgyCapability(
            version, False, endpoint_supported, auth_supported, protocol_supported, stream_supported, (), ()
        )
    return AgyCapability(
        version,
        True,
        True,
        True,
        True,
        True,
        ("direct", "execution_preset"),
        ("high", "low", "medium"),
    )


def _runtime_identity(label: str, value: str) -> str:
    return "sha256:" + hashlib.sha256(
        b"agy-iop-runtime-v1\0" + label.encode("ascii") + b"\0" + value.encode("utf-8")
    ).hexdigest()


def _validate_config_owner_observation(
    cell: MatrixCell, observation: AgyRuntimeObservation
) -> None:
    """Validate an opaque observation supplied by the independent config owner."""
    if not isinstance(observation, AgyRuntimeObservation):
        raise AgyAdapterError("agy runtime observation is invalid")
    if (observation.cell_id, observation.route_kind, observation.route_id) != (
        cell.id, cell.iop.route_kind, cell.iop.route_id,
    ):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    if not all(_IDENTITY_RE.fullmatch(value) for value in (
        observation.endpoint_identity,
        observation.credential_identity,
        observation.config_identity,
    )):
        raise AgyAdapterError("agy IOP runtime observation identity is invalid")


def validate_agy_iop_runtime(
    cell: MatrixCell, runtime: AgyRuntimeInputs, observation: AgyRuntimeObservation
) -> _ValidatedAgyRuntime:
    """Admit runtime-only launch values only when IOP identity proof is exact."""
    if not isinstance(cell, MatrixCell) or cell.caller != AGY_CALLER:
        raise AgyAdapterError("agy runtime validation requires an agy matrix cell")
    if not isinstance(runtime, AgyRuntimeInputs) or not isinstance(observation, AgyRuntimeObservation):
        raise AgyAdapterError("agy runtime inputs are invalid")
    if not all(isinstance(value, str) and value for value in (runtime.binary, runtime.endpoint, runtime.credential)):
        raise AgyAdapterError("agy runtime values are unavailable")
    endpoint = urlsplit(runtime.endpoint)
    route_path = f"/gemini/{cell.iop.route_id}"
    if (
        endpoint.scheme != "https"
        or not endpoint.netloc
        or endpoint.query
        or endpoint.fragment
        or endpoint.path.rstrip("/") != route_path
    ):
        raise AgyAdapterError("agy IOP endpoint is invalid")
    _validate_config_owner_observation(cell, observation)
    if observation.endpoint_identity != _runtime_identity("endpoint", runtime.endpoint):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    if observation.credential_identity != _runtime_identity("credential", runtime.credential):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    return _ValidatedAgyRuntime(runtime.binary, runtime.endpoint, runtime.credential, observation)


def preflight_agy_iop(
    cell: MatrixCell,
    capability: AgyCapability,
    runtime: AgyRuntimeInputs,
    observation: AgyRuntimeObservation,
) -> AgyPreflightResult:
    """Classify only registration and implementation gaps without launching agy."""
    if not isinstance(cell, MatrixCell) or cell.caller != AGY_CALLER:
        raise AgyAdapterError("agy preflight requires an agy matrix cell")
    if not isinstance(capability, AgyCapability) or not isinstance(runtime, AgyRuntimeInputs):
        raise AgyAdapterError("agy preflight inputs are invalid")

    issues: list[ConnectivityIssue] = []
    validated_runtime: _ValidatedAgyRuntime | None = None
    if runtime.credential:
        try:
            validated_runtime = validate_agy_iop_runtime(cell, runtime, observation)
        except AgyAdapterError:
            issues.append(_issue("endpoint_incompatible"))
    if not runtime.credential:
        issues.append(_issue("credential_missing"))
    if cell.iop.request_model not in AGY_MODEL_LABELS:
        issues.append(_issue("model_missing"))
    if not runtime.endpoint:
        issues.append(_issue("endpoint_incompatible"))
    if not capability.endpoint_supported:
        issues.append(_issue("endpoint_incompatible"))
    if not capability.auth_supported:
        issues.append(_issue("auth_incompatible"))
    if not capability.protocol_supported:
        issues.append(_issue("protocol_incompatible"))
    if not capability.stream_supported:
        issues.append(_issue("stream_incompatible"))
    if capability.iop_transport_supported and cell.iop.route_kind not in capability.route_kinds:
        issues.append(_issue("protocol_incompatible"))
    elif capability.iop_transport_supported and cell.iop.requested_effort not in capability.efforts:
        issues.append(_issue("effort_unsupported"))

    # Preserve connectivity.py's canonical issue order without leaking values.
    unique = {item.code: item for item in issues}
    ordered = tuple(
        unique[code]
        for code in (
            "credential_missing", "model_missing", "route_missing", "effort_unsupported",
            "endpoint_incompatible", "auth_incompatible", "protocol_incompatible", "stream_incompatible",
        )
        if code in unique
    )
    status = classify_issues(ordered)
    return AgyPreflightResult(capability, _requested_binding(cell), ordered, status, validated_runtime)


def build_agy_invocation(
    cell: MatrixCell,
    prepared: PreparedWorkspace,
    task_payload: bytes,
    timeout: Timeout,
    preflight: AgyPreflightResult,
) -> InvocationSpec:
    """Build one isolated official agy print invocation after a ready preflight."""
    if preflight.status != "ready" or not preflight.capability.iop_transport_supported or preflight.runtime is None:
        raise AgyAdapterError("agy IOP transport is not proven")
    runtime = preflight.runtime
    if not runtime.binary or not Path(runtime.binary).is_file():
        raise AgyAdapterError("agy binary is unavailable")
    if (
        not isinstance(prepared, PreparedWorkspace)
        or not Path(prepared.workspace_dir).is_dir()
        or not Path(prepared.session_dir).is_dir()
        or not Path(prepared.attempt_root).is_dir()
    ):
        raise AgyAdapterError("prepared workspace is unavailable")
    if not isinstance(task_payload, bytes) or not task_payload:
        raise AgyAdapterError("agy task payload is unavailable")
    try:
        task_text = task_payload.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise AgyAdapterError("agy task payload must be UTF-8") from exc

    _stage_agy_provider_settings(prepared.session_dir)

    # The child receives a minimal environment and explicit IOP-only provider
    # settings.  No parent agy/Gemini config or session variable is inherited.
    environment = {
        "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "TZ": "UTC",
        "HOME": prepared.session_dir,
        AGY_ENDPOINT_ENV: runtime.endpoint,
        AGY_AUTH_ENV: runtime.credential,
    }
    environment.update(inherited_tls_ca_environment())
    return InvocationSpec(
        argv=(
            runtime.binary,
            "--sandbox",
            "--output-format", "stream-json",
            "--model", AGY_MODEL_LABELS[cell.iop.request_model],
            "--print", task_text,
        ),
        cwd=prepared.workspace_dir,
        env=env_pairs(environment),
        env_allowlist=(AGY_ENDPOINT_ENV, AGY_AUTH_ENV, *TLS_CA_ENV_KEYS),
        submission_mode=SUBMISSION_ARGV_TASK,
        completion_mode=COMPLETION_EXIT_AFTER_IDLE,
        timeout=timeout,
        evidence_dir=prepared.attempt_root,
        control_dir=str(Path(prepared.attempt_root) / "agy-control"),
    )


def _stage_agy_provider_settings(session_dir: str) -> None:
    """Create the exact secret-free provider selector in one isolated HOME."""
    try:
        session = Path(session_dir).resolve(strict=True)
        if not session.is_dir():
            raise OSError("session is not a directory")
        settings = session / AGY_SETTINGS_RELATIVE_PATH
        settings.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        if settings.parent.resolve(strict=True) != session / AGY_SETTINGS_RELATIVE_PATH.parent:
            raise OSError("settings directory escapes session")
        payload = (
            json.dumps(_AGY_ISOLATED_SETTINGS, sort_keys=True, separators=(",", ":")) + "\n"
        ).encode("utf-8")
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
        if hasattr(os, "O_NOFOLLOW"):
            flags |= os.O_NOFOLLOW
        descriptor = os.open(settings, flags, 0o600)
        with os.fdopen(descriptor, "wb") as handle:
            handle.write(payload)
    except (OSError, ValueError) as exc:
        raise AgyAdapterError("agy isolated provider settings are unavailable") from exc


def _safe_identifier(value: Any) -> str | None:
    return value if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:+-]{0,63}", value) else None


def redact_agy_event(raw_line: str, sensitive_values: tuple[str, ...] = ()) -> str:
    """Return a canonical allowlisted event projection, never caller content."""
    try:
        parsed = json.loads(raw_line)
    except (TypeError, json.JSONDecodeError):
        return '{"event":"unparseable"}'
    if not isinstance(parsed, dict):
        return '{"event":"unparseable"}'
    event = _safe_identifier(parsed.get("event"))
    if event is None or event in sensitive_values:
        return '{"event":"unparseable"}'
    safe: dict[str, str] = {"event": event}
    payload = parsed.get(event)
    if not isinstance(payload, dict):
        payload = {}
    for field in _SAFE_EVENT_FIELDS[1:]:
        value = _safe_identifier(payload.get(field))
        if value is not None and value not in sensitive_values:
            safe[field] = value
    return json.dumps(safe, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


class AgyEventParser:
    """Parse the official 1.1.12 stream while trusting binding only from config."""

    def __init__(self, cell: MatrixCell, admitted_binding: RequestedEffectiveBinding) -> None:
        if not isinstance(admitted_binding, RequestedEffectiveBinding):
            raise AgyAdapterError("agy admitted binding is invalid")
        if admitted_binding.cell_id != cell.id or admitted_binding.caller != cell.caller:
            raise AgyAdapterError("agy admitted binding mismatch")
        self._cell = cell
        self._admitted_binding = admitted_binding
        self._init_seen = False
        self._result_seen = False
        self._latest_usage: dict[str, Any] | None = None

    def __call__(self, stream: str, raw_line: str) -> str | ParsedMetric | tuple[Any, ...] | None:
        return self.parse(stream, raw_line)

    def parse(self, stream: str, raw_line: str) -> str | ParsedMetric | tuple[Any, ...] | None:
        if stream != "stdout":
            return None
        try:
            item = json.loads(raw_line)
        except (TypeError, json.JSONDecodeError):
            return "malformed"
        if not isinstance(item, dict):
            return "malformed"
        event = item.get("event")
        payload = item.get(event) if isinstance(event, str) else None
        if not isinstance(payload, dict):
            return "malformed"
        if event == "init":
            if self._init_seen or self._result_seen:
                return "malformed"
            self._init_seen = True
            return None
        if event == "step_update":
            if not self._init_seen or self._result_seen:
                return "malformed"
            usage = payload.get("usage")
            if usage is not None:
                if self._usage_metrics(usage) is None:
                    return "malformed"
                self._latest_usage = usage
            return None
        if event != "result" or not self._init_seen or self._result_seen:
            return "malformed"
        self._result_seen = True
        if payload.get("status") != "SUCCESS":
            return "malformed"
        metrics: list[ParsedMetric] = []
        usage = payload.get("usage", self._latest_usage)
        if usage is not None:
            parsed_usage = self._usage_metrics(usage)
            if parsed_usage is None:
                return "malformed"
            metrics.extend(parsed_usage)
        try:
            if "duration_seconds" in payload:
                metrics.append(duration_metric(
                    "total_duration", payload["duration_seconds"], reported_unit="s",
                    model=self._cell.iop.request_model,
                ))
            if "num_turns" in payload:
                metrics.append(count_metric(
                    "model_calls", payload["num_turns"], model=self._cell.iop.request_model,
                ))
        except LifecycleMetricError:
            return "malformed"
        return tuple(metrics) + ("finish", "idle")

    def _usage_metrics(self, usage: Any) -> tuple[ParsedMetric, ...] | None:
        if not isinstance(usage, dict) or not set(usage) <= set(_AGY_USAGE_METRICS):
            return None
        metrics: list[ParsedMetric] = []
        try:
            for wire_name, metric_name in _AGY_USAGE_METRICS.items():
                if wire_name in usage:
                    metrics.append(count_metric(
                        metric_name, usage[wire_name], model=self._cell.iop.request_model,
                    ))
        except LifecycleMetricError:
            return None
        return tuple(metrics)

    def observed_result(self, capability: AgyCapability, lifecycle: InvocationResult) -> ConnectivityResult:
        requested = _requested_binding(self._cell)
        closed_gap = (_issue("stream_incompatible"),)
        caller_capability = CallerCapability(AGY_CALLER, capability.route_kinds, capability.efforts)
        if (
            not isinstance(lifecycle, InvocationResult)
            or not lifecycle.success
            or not lifecycle.finish_then_idle_then_quiet
            or not self._result_seen
        ):
            return make_result(self._cell, caller_capability, requested, closed_gap)
        try:
            return make_result(self._cell, caller_capability, self._admitted_binding)
        except Exception:
            return make_result(self._cell, caller_capability, requested, closed_gap)


def run_agy_invocation(
    spec: InvocationSpec,
    parser: AgyEventParser,
    preflight: AgyPreflightResult,
    on_started: Callable[[SupervisorLocator], None],
) -> InvocationResult:
    """Run a prepared agy call with structural output redaction only."""
    if (
        not isinstance(preflight, AgyPreflightResult)
        or preflight.status != "ready"
        or not preflight.capability.iop_transport_supported
        or preflight.runtime is None
    ):
        raise AgyAdapterError("agy IOP transport is not proven")
    runtime = preflight.runtime
    sensitive = (runtime.endpoint, runtime.credential)
    structural = lambda line: redact_agy_event(line, sensitive)
    # exact replacement is retained as a final defence for non-JSON stderr.
    exact = exact_value_redactor(sensitive)
    return run_invocation(
        spec,
        parse_event=parser,
        on_started=on_started,
        redact=lambda line: structural(exact(line)),
    )