#!/usr/bin/env python3
"""Revoke a runtime lease from a compaction or lifecycle observer."""

from __future__ import annotations

import argparse
import json
import os
import socket
import sys
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import Final

from lease_generation import bump_runtime_generation, read_runtime_generation

MAX_FRAME: Final = 64 * 1024
BROKER_TIMEOUT_SECONDS: Final = 1.5


def broker_request(socket_path: Path, request: dict[str, object]) -> dict[str, object]:
    payload = (json.dumps(request, separators=(",", ":")) + "\n").encode()
    if len(payload) > MAX_FRAME:
        raise ValueError("invalid revoke request")
    response = bytearray()
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(BROKER_TIMEOUT_SECONDS)
        connection.connect(str(socket_path))
        connection.sendall(payload)
        connection.shutdown(socket.SHUT_WR)
        while len(response) <= MAX_FRAME:
            chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(response)))
            if not chunk:
                break
            response.extend(chunk)
    if len(response) > MAX_FRAME or not response.endswith(b"\n"):
        raise ValueError("invalid broker reply")
    value = json.loads(response)
    if not isinstance(value, dict):
        raise ValueError("invalid broker reply")
    return value


def main(
    argv: Sequence[str] | None = None,
    *,
    environ: Mapping[str, str] | None = None,
    request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
    parser.add_argument("--reason", required=True)
    parser.add_argument("--bump-generation", action="store_true")
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ

    try:
        if not arguments.reason or len(arguments.reason) > 128:
            raise ValueError("invalid revoke reason")
        socket_path = Path(source_environment["MOSAIC_LEASE_BROKER_SOCKET"])
        session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
        if (
            len(session_id) != 64
            or any(character not in "0123456789abcdef" for character in session_id)
        ):
            raise ValueError("invalid broker session")
        generation = (
            bump_runtime_generation(source_environment)
            if arguments.bump_generation
            else read_runtime_generation(source_environment)
        )
        reply = request(
            socket_path,
            {
                "action": "revoke_lease",
                "session_id": session_id,
                "runtime_generation": generation,
                "reason": arguments.reason,
                "runtime": arguments.runtime,
            },
        )
        if reply.get("ok") is not True or reply.get("state") != "UNVERIFIED":
            raise ValueError("revocation refused")
    except (KeyError, ValueError, OSError, json.JSONDecodeError):
        # A fired observer must remain fail-closed even if the broker transport
        # is unavailable. Advancing the private local generation fences every
        # later tool check; the broker revokes the old lease when it next sees
        # that higher generation. Explicit rollover already advanced it above.
        if not arguments.bump_generation:
            try:
                bump_runtime_generation(source_environment)
            except (KeyError, ValueError, OSError):
                pass
        print("Mosaic lease revocation failed; lifecycle transition denied.", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())