"""Bundled Antigravity provider catalog."""
from typing import Any, Mapping

from okstra_ctl.domain.provider import (
    LeadLaunchSpec,
    ModelSpec,
    ProviderSpec,
    ServedModelAttestation,
    snapshot_execution_normalizations,
)
from okstra_ctl.domain.role import ALL_ROLE_TOKENS
from okstra_ctl.domain.worker_exec import (
    ExecCommand,
    PolicySupport,
    WorkerExecRequest,
)
from okstra_ctl.domain.worker_presentation import JsonEvents
from okstra_ctl.domain.worker_stream import (
    Result,
    StreamEvent,
    Text,
    ToolCall,
    ToolResult,
)
import okstra_ctl.model_discovery as model_discovery


# `agy` needs a tier suffix at dispatch time. The provider exposes its
# normalizer through ProviderSpec so assignment resolution does not branch on
# provider names.
ANTIGRAVITY = {
    "gemini-3.1-pro": ModelSpec(
        "gemini-3.1-pro", "Gemini 3.1 Pro", "gemini-3.1-pro",
        aliases=("gemini 3.1 pro",),
    ),
    "gemini-3.6-flash": ModelSpec(
        "gemini-3.6-flash", "Gemini 3.6 Flash", "gemini-3.6-flash",
        aliases=("gemini 3.6 flash",),
    ),
    "gemini-3.5-flash": ModelSpec(
        "gemini-3.5-flash", "Gemini 3.5 Flash", "gemini-3.5-flash",
        aliases=("gemini 3.5 flash",),
    ),
}


def _catalog_model_id(observed: str) -> str:
    """Undo the dispatch-time tier suffix so the id is catalog-comparable.

    Dispatch sends `gemini-3.1-pro-low`; agy serves that identity back
    verbatim. The catalog holds the bare `gemini-3.1-pro`, so reading the
    observation as-is made the served-model gate reject a model okstra itself
    had asked for. Only suffixes dispatch can append are stripped, and only
    when what remains is a catalog entry — anything else passes through for
    the gate to judge.
    """
    if observed in ANTIGRAVITY:
        return observed  # identity already carries its tier (agy's Claude models)
    for suffix in model_discovery.dispatch_tier_suffixes():
        bare = observed.removesuffix(f"-{suffix}")
        if bare != observed and bare in ANTIGRAVITY:
            return bare
    return observed


def normalise_served_model(raw_model: str | None) -> ServedModelAttestation:
    if not raw_model or not raw_model.strip():
        return ServedModelAttestation.unknown()
    model_id = raw_model.strip().lower().removeprefix("antigravity/")
    return ServedModelAttestation(
        raw_model,
        f"antigravity/{_catalog_model_id(model_id)}",
        "exact",
        "provider-output",
    )


def observe_served_model(event: Mapping[str, Any]) -> str | None:
    if event.get("event") != "init":
        return None
    init = event.get("init")
    if not isinstance(init, Mapping):
        return None
    model = init.get("model")
    return model if isinstance(model, str) and model.strip() else None


# The wall-clock cap on one print run, carried over from the wrapper. It is
# deliberately not derived from the idle budget: an analyser legitimately works
# for far longer than it goes quiet, so folding the idle budget in here would
# kill a healthy worker mid-run. Idle is the runner's watchdog to enforce.
_PRINT_TIMEOUT = "7200s"

_STEP_UPDATE = "step_update"
_RESULT = "result"
_CALL_STARTED = "ACTIVE"
_CALL_FINISHED = "DONE"


class AntigravityExecutionNormalizer:
    """Capture `agy models` once, then normalize the complete catalog."""

    def snapshot(self, models, roles):
        available = model_discovery.agy_models()
        return snapshot_execution_normalizations(
            models,
            roles,
            lambda execution, role: model_discovery.normalize_antigravity_execution(
                execution=execution,
                role=role,
                available=available,
            ),
        )


def normalize_execution(execution: str, role: str) -> str:
    """Compatibility entry point for direct adapter callers."""
    return model_discovery.normalize_execution_for_dispatch(
        provider="antigravity", execution=execution, role=role
    )


def step_update_events(event: Mapping[str, Any]) -> tuple[StreamEvent, ...]:
    """Normalise this CLI's stream-json, which shares no key with the other four.

    Measured 2026-08-11 against agy 1.1.11. Events are keyed on `event` rather
    than `type`, and progress arrives as one `step_update` per step transition:

    - `step_type: agent_response` carries the worker's prose in `text_delta`
      (absent on the turns that only planned a tool call).
    - `step_type: tool` arrives twice for one call — `ACTIVE` with the
      arguments, then `DONE` with the tool's output.
    - the closing `result` carries `response`, and an `error` string when
      `status` is anything but SUCCESS. A failed run's `response` is empty, so
      without the error text such a run would leave nothing behind at all.
    """
    kind = event.get("event")
    if kind == _RESULT:
        return _result_events(event.get(_RESULT))
    if kind != _STEP_UPDATE:
        return ()
    step = event.get(_STEP_UPDATE)
    if not isinstance(step, Mapping):
        return ()
    if step.get("step_type") == "agent_response":
        text = step.get("text_delta")
        return (Text(body=text),) if isinstance(text, str) and text.strip() else ()
    if step.get("step_type") == "tool":
        return _tool_events(step)
    return ()


def _tool_events(step: Mapping[str, Any]) -> tuple[StreamEvent, ...]:
    info = step.get("tool_info")
    info = info if isinstance(info, Mapping) else {}
    name = str(step.get("tool_name") or info.get("name") or "tool")
    if step.get("state") == _CALL_STARTED:
        return (ToolCall(name=name, detail=_argument(info.get("parameters"))),)
    if step.get("state") != _CALL_FINISHED:
        return ()
    output = str(info.get("output", ""))
    # No outcome field: a `run_command` whose command exited non-zero was
    # measured closing as `DONE` with the failure text in `output` and nothing
    # else to distinguish it, so the outcome is left unreported rather than
    # rendered as success.
    return (ToolResult(body=output, size_bytes=len(output.encode("utf-8"))),)


def _argument(parameters: Any) -> str:
    """The one argument worth a row, without guessing at the parameter names.

    Parameter keys belong to each tool, not to this CLI — `view_file` names its
    path `AbsolutePath` while `run_command` names its command `CommandLine` —
    so an allowlist of key names would cover only the tools that happened to be
    observed. The first string argument is the subject of every tool measured.
    """
    if not isinstance(parameters, Mapping):
        return ""
    return next(
        (str(value) for value in parameters.values() if isinstance(value, str) and value),
        "",
    )


def _result_events(result: Any) -> tuple[StreamEvent, ...]:
    if not isinstance(result, Mapping):
        return ()
    error = result.get("error")
    failure = (Text(body=error),) if isinstance(error, str) and error.strip() else ()
    response = result.get("response")
    if not isinstance(response, str):
        return failure
    return (*failure, Result(text=response))


class AntigravityExecution:
    """agy CLI invocation.

    `--add-dir` defines the workspace rather than widening a default one, so the
    project root is included instead of skipped.

    The prompt is an argument, not stdin: this CLI does not read stdin.

    `--dangerously-skip-permissions` is required rather than convenient. Headless
    `--print` has nobody to answer a permission request, so without it every
    command tool is auto-denied and the run still reports SUCCESS with an empty
    response.

    No flag bounds what this worker may write. `--sandbox` was measured on
    2026-08-10 (agy 1.1.11) and did not block writes outside the `--add-dir`
    workspace through either the shell tool or the file tool — identically to a
    control run without it. It is therefore not passed, and `policy_support`
    reports the gap rather than claiming a boundary that was not observed.
    """

    def build_command(self, request: WorkerExecRequest) -> ExecCommand:
        argv = ["agy", "--print", request.prompt_text, "--model", request.model]
        for directory in request.policy.write_scope:
            argv += ["--add-dir", str(directory)]
        argv += ["--output-format", "stream-json"]
        argv += ["--print-timeout", _PRINT_TIMEOUT]
        if request.policy.auto_approve:
            argv.append("--dangerously-skip-permissions")
        return ExecCommand(
            argv=tuple(argv),
            stdin_text=None,
            cwd=request.project_root,
            presentation=JsonEvents(
                normalise=step_update_events,
                observe=observe_served_model,
            ),
        )

    def policy_support(self) -> PolicySupport:
        return PolicySupport(
            can_auto_approve=True,
            can_bound_write_scope=False,
            note=(
                "measured 2026-08-10 (agy 1.1.11): --sandbox did not block writes "
                "outside the --add-dir workspace, so no flag bounds this provider"
            ),
        )


def create_provider() -> ProviderSpec:
    return ProviderSpec(
        provider="antigravity",
        display_label="Antigravity",
        models=ANTIGRAVITY,
        default_models={
            role: "gemini-3.1-pro"
            for role in (
                "lead", "analyser", "critic", "designer", "planner", "executor",
                "verifier", "report-writer", "translator",
            )
        },
        wrapper="okstra-antigravity-exec.sh",
        supported_roles=ALL_ROLE_TOKENS,
        execution_capabilities=frozenset(
            {
                "lead-session",
                "worker-artifact-io",
                "source-readonly",
                "extended-artifact-authoring",
                "project-mutation",
            }
        ),
        lead_launch=LeadLaunchSpec(
            executable="agy", model_flag="--model", prompt_flag="--prompt-interactive"
        ),
        exec_strategy=AntigravityExecution(),
        execution_normalizer=AntigravityExecutionNormalizer(),
        served_model_normalizer=normalise_served_model,
    )