#!/usr/bin/env python3
"""Run an external LLM CLI with the shared okstra wrapper contract.

Argument parsing and provider lookup only — the run itself belongs to
``okstra_ctl.worker_runner``, which every provider shares. What stays here is
the request the provider strategies then take on trust: resolved paths, the
write scope in the order the CLIs are told it, and the role's idle budget.
"""
from __future__ import annotations

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

_HERE = Path(__file__).resolve().parent
# ``okstra_ctl`` sits beside this file in the repo (``scripts/``) but under
# ``~/.okstra/lib/python/`` once installed, and the four-line shell entrypoints
# that exec this script set no PYTHONPATH. Offer both, repo first.
_HOME_LIB = (
    Path(os.environ.get("OKSTRA_HOME", str(Path.home() / ".okstra"))) / "lib" / "python"
)
sys.path.insert(0, str(_HERE))
if _HOME_LIB.is_dir() and str(_HOME_LIB) not in sys.path:
    sys.path.append(str(_HOME_LIB))

from okstra_ctl.domain.provider import (  # noqa: E402
    ProviderSpec,
    ServedModelAttestation,
    ServedModelNormalizer,
    UnknownProviderError,
)
from okstra_ctl.domain.role import normalize_role, role_for_duty  # noqa: E402
from okstra_ctl.wrapper_status import log_path_for_prompt  # noqa: E402
from okstra_ctl.domain.worker_exec import (  # noqa: E402
    ExecutionStrategy,
    WorkerExecRequest,
)
from okstra_ctl.registry.provider_registry import (  # noqa: E402
    default_provider_registry,
)
from okstra_ctl.worker_request import build_request, idle_timeout  # noqa: E402
from okstra_ctl.worker_runner import LIVE, QUIET, run_worker  # noqa: E402
from okstra_ctl.dispatch_core import verify_served_model  # noqa: E402
from okstra_ctl.dispatch_state import DispatchError  # noqa: E402
from okstra_ctl.execution_manifest import read_execution_manifest  # noqa: E402
from okstra_ctl.write_policy import (  # noqa: E402
    WriteEnforcement,
    WritePolicy,
    write_enforcement_from_payload,
    write_policy_from_payload,
)
from okstra_ctl.model_pool import ModelPool  # noqa: E402
from okstra_ctl.agent_invocation import (  # noqa: E402
    AgentInvocationError,
    agent_model_assignment_from_payload,
    invocation_metadata_identity,
)

_USAGE = (
    "usage: okstra-provider-exec.py <provider> <project-root> "
    "<model-execution-value> <prompt-path> [worktree-path] [role] "
    "[idle-timeout-seconds] [--presentation live|quiet] [--session-id <uuid>] "
    "[--invocation-metadata path]"
)

_PRESENTATION_FLAG = "--presentation"
_SESSION_ID_FLAG = "--session-id"
_INVOCATION_METADATA_FLAG = "--invocation-metadata"
_PRESENTATIONS = (LIVE, QUIET)


class PreflightError(Exception):
    def __init__(self, exit_code: int, message: str) -> None:
        super().__init__(message)
        self.exit_code = exit_code


@dataclass(frozen=True)
class Invocation:
    strategy: ExecutionStrategy
    request: WorkerExecRequest
    presentation: str
    log_path: Path
    status_path: Path
    status_extra: Mapping[str, Any]
    served_model_normalizer: ServedModelNormalizer | None
    served_model_verifier: Callable[[ServedModelAttestation], str] | None


@dataclass(frozen=True)
class _V2InvocationContext:
    status_extra: Mapping[str, Any]
    verify_attestation: Callable[[ServedModelAttestation], str]
    write_policy: WritePolicy
    write_enforcement: WriteEnforcement


def parse_invocation(argv: list[str]) -> Invocation:
    """Resolve the wrapper's positional contract into one runnable dispatch."""
    positional, presentation = _take_presentation(argv)
    # Empty unless the dispatcher issued one. Without it the CLI picks its own
    # id and nothing downstream can map that session back to this worker.
    positional, session_id = _take_flag(positional, _SESSION_ID_FLAG, "")
    positional, metadata_raw = _take_flag(positional, _INVOCATION_METADATA_FLAG, "")
    if not 4 <= len(positional) <= 7:
        raise PreflightError(64, _USAGE)
    provider_id, project_root_raw, model, prompt_raw = positional[:4]
    worktree_raw = positional[4] if len(positional) >= 5 else ""
    role = positional[5] if len(positional) >= 6 and positional[5] else "worker"
    timeout_raw = positional[6] if len(positional) >= 7 and positional[6] else ""

    project_root = _existing_dir(project_root_raw, 65, "project-root")
    if not model:
        raise PreflightError(66, "model-execution-value is empty")
    prompt_path = _existing_file(prompt_raw, 67, "prompt-path")
    idle_timeout_seconds = _idle_timeout(timeout_raw, role)
    worktree = (
        _existing_dir(worktree_raw, 68, "worktree-path") if worktree_raw else None
    )
    spec = _provider_spec(provider_id)
    status_extra: dict[str, Any] = {"wrapper": spec.wrapper, "role": role}
    served_model_normalizer: ServedModelNormalizer | None = None
    served_model_verifier: Callable[[ServedModelAttestation], str] | None = None
    write_policy: WritePolicy | None = None
    write_enforcement: WriteEnforcement | None = None
    if metadata_raw:
        context = _v2_status_extra(
            metadata_path=_existing_file(
                metadata_raw, 64, "invocation-metadata"
            ),
            project_root=project_root,
            prompt_path=prompt_path,
            provider_id=provider_id,
            model=model,
            role=role,
        )
        status_extra.update(context.status_extra)
        served_model_normalizer = spec.served_model_normalizer
        served_model_verifier = context.verify_attestation
        write_policy = context.write_policy
        write_enforcement = context.write_enforcement
        if not spec.supports_role(role):
            raise PreflightError(
                64, f"provider {provider_id!r} does not support role {role!r}"
            )

    request = build_request(
        prompt_text=prompt_path.read_text(encoding="utf-8"),
        model=model,
        project_root=project_root,
        worktree_path=worktree,
        role=role,
        idle_timeout_seconds=idle_timeout_seconds,
        session_id=session_id,
        write_policy=write_policy,
        write_enforcement=write_enforcement,
    )
    strategy = spec.exec_strategy
    _check_command(strategy, request)
    return Invocation(
        strategy=strategy,
        request=request,
        presentation=presentation,
        log_path=log_path_for_prompt(prompt_path),
        status_path=Path(f"{prompt_path}.status.json"),
        status_extra=status_extra,
        served_model_normalizer=served_model_normalizer,
        served_model_verifier=served_model_verifier,
    )


def _take_flag(argv: list[str], flag: str, default: str) -> tuple[list[str], str]:
    """Split one value-carrying flag out of an otherwise positional argv.

    The wrapper's contract is positional, so every flag it grows has to be
    lifted out before the positions are counted. Written once because the flags
    differ only in the value they carry: a second copy of this loop is where
    their refusals of a flag with no value would drift apart.
    """
    positional: list[str] = []
    value = default
    index = 0
    while index < len(argv):
        if argv[index] != flag:
            positional.append(argv[index])
            index += 1
            continue
        if index + 1 >= len(argv):
            raise PreflightError(64, f"{flag} needs a value: {_USAGE}")
        value = argv[index + 1]
        index += 2
    return positional, value


def _take_presentation(argv: list[str]) -> tuple[list[str], str]:
    """``_take_flag`` plus the allowlist only this flag has.

    Defaults to ``quiet``. ``live`` is only ever right where a screen was
    declared, and the only callers that can declare one are the pane backends —
    which pass the flag explicitly. Defaulting the other way assumed a screen
    that a subagent dispatch does not have, and sent every worker's progress
    into its caller's context window instead.
    """
    positional, presentation = _take_flag(argv, _PRESENTATION_FLAG, QUIET)
    if presentation not in _PRESENTATIONS:
        allowed = " | ".join(_PRESENTATIONS)
        raise PreflightError(
            64, f"unsupported presentation {presentation!r}. Allowed values: {allowed}"
        )
    return positional, presentation


def _v2_status_extra(
    *,
    metadata_path: Path,
    project_root: Path,
    prompt_path: Path,
    provider_id: str,
    model: str,
    role: str,
) -> _V2InvocationContext:
    try:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise PreflightError(64, "invocation metadata is invalid") from exc
    if not isinstance(metadata, dict):
        raise PreflightError(64, "invocation metadata must be an object")
    try:
        identity = invocation_metadata_identity(metadata)
        assignment = agent_model_assignment_from_payload(
            metadata.get("modelAssignment")
        )
    except AgentInvocationError as exc:
        raise PreflightError(64, str(exc)) from exc
    if identity is None:
        raise PreflightError(64, "new worker dispatch requires v2 invocation metadata")
    if (
        assignment.provider != provider_id
        or assignment.model_execution_value != model
    ):
        raise PreflightError(64, "invocation metadata model assignment does not match wrapper")
    try:
        expected_role = role_for_duty(identity.duty_id)
        actual_role = normalize_role(role)
    except ValueError as exc:
        raise PreflightError(64, f"invocation metadata role is invalid: {exc}") from exc
    if expected_role != actual_role:
        raise PreflightError(64, "invocation metadata duty does not match role")
    prompt = metadata.get("prompt")
    recorded_path = prompt.get("path") if isinstance(prompt, Mapping) else None
    if not isinstance(recorded_path, str):
        raise PreflightError(64, "invocation metadata prompt path is invalid")
    recorded = Path(recorded_path)
    if not recorded.is_absolute():
        recorded = project_root / recorded
    if recorded.resolve(strict=False) != prompt_path.resolve(strict=True):
        raise PreflightError(64, "invocation metadata prompt does not match wrapper")
    role_execution, invocation, pool = _execution_authority(
        metadata=metadata,
        project_root=project_root,
        role_execution_ref=identity.role_execution_ref,
        invocation_ref=identity.invocation_ref,
    )
    binding = role_execution.binding
    if (
        role_execution.participant_ref != identity.participant_ref
        or role_execution.role != expected_role
        or role_execution.provider != assignment.provider
        or binding is None
        or binding.runner != assignment.runner
        or binding.resolved_execution_value != assignment.model_execution_value
    ):
        raise PreflightError(
            64, "invocation metadata role execution binding does not match run manifest"
        )
    if role_execution.execution_label != identity.execution_label:
        raise PreflightError(
            64, "invocation metadata execution label does not match run manifest"
        )

    def verify_attestation(attestation: ServedModelAttestation) -> str:
        try:
            verify_served_model(role_execution, attestation, pool=pool)
        except DispatchError as exc:
            return str(exc)
        return ""

    try:
        write_policy = write_policy_from_payload(invocation.write_policy)
        write_enforcement = write_enforcement_from_payload(
            invocation.write_enforcement
        )
    except ValueError as exc:
        raise PreflightError(64, f"invocation write contract is invalid: {exc}") from exc

    return _V2InvocationContext({
        "schemaVersion": "2.0",
        "executionIdentityVersion": 2,
        "participantRef": identity.participant_ref,
        "roleExecutionRef": identity.role_execution_ref,
        "executionLabel": identity.execution_label,
        "dutyId": identity.duty_id,
        "invocationRef": identity.invocation_ref,
        "attempt": identity.attempt,
        "servedModelAttestation": {
            "observedModel": None,
            "normalizedModelRef": None,
            "level": "unknown",
            "source": "unavailable",
        },
        "writePolicyDigest": invocation.write_policy_digest,
        "writeEnforcement": write_enforcement.to_payload(),
    }, verify_attestation, write_policy, write_enforcement)


def _execution_authority(
    *,
    metadata: Mapping[str, Any],
    project_root: Path,
    role_execution_ref: str,
    invocation_ref: str,
):
    source = metadata.get("contractSource")
    manifest_raw = source.get("runManifestPath") if isinstance(source, Mapping) else None
    if not isinstance(manifest_raw, str) or not manifest_raw:
        raise PreflightError(64, "invocation metadata run manifest is invalid")
    manifest_path = Path(manifest_raw)
    if not manifest_path.is_absolute():
        manifest_path = project_root / manifest_path
    try:
        manifest_path = manifest_path.resolve(strict=True)
        manifest_path.relative_to(project_root.resolve(strict=True))
        manifest = read_execution_manifest(manifest_path)
    except (OSError, ValueError) as exc:
        raise PreflightError(64, "invocation metadata run manifest is invalid") from exc
    role_execution = next((
        row for row in manifest.role_executions
        if row.role_execution_ref == role_execution_ref
    ), None)
    if role_execution is None:
        raise PreflightError(
            64, "invocation metadata role execution does not match run manifest"
        )
    invocation = next((
        row for row in manifest.invocations
        if row.invocation_ref == invocation_ref
    ), None)
    if invocation is None or invocation.role_execution_ref != role_execution_ref:
        raise PreflightError(
            64, "invocation metadata invocation does not match run manifest"
        )
    return (
        role_execution,
        invocation,
        ModelPool.from_registry(default_provider_registry()),
    )


def _provider_spec(provider_id: str) -> ProviderSpec:
    try:
        spec = default_provider_registry().resolve(provider_id)
    except UnknownProviderError as exc:
        raise PreflightError(64, str(exc)) from exc
    if spec.exec_strategy is None:
        raise PreflightError(
            64, f"provider {provider_id!r} has no execution strategy to run"
        )
    return spec


def _existing_dir(raw: str, exit_code: int, label: str) -> Path:
    """Existence only — `build_request` owns the resolving."""
    path = Path(raw) if raw else None
    if path is None or not path.is_dir():
        raise PreflightError(
            exit_code, f"{label} is missing or not a directory: {raw!r}"
        )
    return path


def _existing_file(raw: str, exit_code: int, label: str) -> Path:
    path = Path(raw) if raw else None
    if path is None or not path.is_file():
        raise PreflightError(exit_code, f"{label} is missing or not a file: {raw!r}")
    return path.resolve()


def _idle_timeout(raw: str, role: str) -> int:
    try:
        return idle_timeout(raw, role)
    except ValueError as exc:
        raise PreflightError(69, str(exc)) from exc


def _check_command(strategy: ExecutionStrategy, request: WorkerExecRequest) -> None:
    """Refuse a missing CLI before the run leaves any artifact behind.

    Building the command is the only truthful way to learn which binary this
    provider runs, and it is a pure call. The check stays out of the runner so a
    refused dispatch leaves no `started` status sidecar for the liveness probe to
    read as a worker that launched.
    """
    binary = strategy.build_command(request).argv[0]
    if shutil.which(binary) is None:
        raise PreflightError(127, f"{binary} CLI is not installed on PATH")


def main(argv: list[str]) -> int:
    try:
        invocation = parse_invocation(argv[1:])
        return run_worker(
            invocation.strategy,
            invocation.request,
            presentation=invocation.presentation,
            log_path=invocation.log_path,
            status_path=invocation.status_path,
            status_extra=invocation.status_extra,
            served_model_normalizer=invocation.served_model_normalizer,
            served_model_verifier=invocation.served_model_verifier,
        )
    except PreflightError as exc:
        print(f"okstra-provider-exec: {exc}", file=sys.stderr)
        return exc.exit_code
    except OSError as exc:
        print(f"okstra-provider-exec: execution failed: {exc}", file=sys.stderr)
        return 127 if isinstance(exc, FileNotFoundError) else 1


if __name__ == "__main__":
    sys.exit(main(sys.argv))