#!/usr/bin/env python3
"""Runtime-neutral whole mutator-class gate backed by the Mosaic lease broker."""

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 BinaryIO, Final

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


def deny(code: str) -> int:
    print(f"BLOCKED: Mosaic mutator gate denied this tool ({code}).", file=sys.stderr)
    return 2


def read_tool_name(stream: BinaryIO | None = None) -> str:
    source = sys.stdin.buffer if stream is None else stream
    raw = source.read(MAX_FRAME + 1)
    if len(raw) > MAX_FRAME:
        raise ValueError("INVALID_GATE_INPUT")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("INVALID_GATE_INPUT")
    tool_name = value.get("tool_name")
    if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
        raise ValueError("INVALID_GATE_INPUT")
    return tool_name


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_GATE_INPUT")
    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,
    stream: BinaryIO | 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"))
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ

    try:
        tool_name = read_tool_name(stream)
        socket_value = source_environment["MOSAIC_LEASE_BROKER_SOCKET"]
        session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
        generation = int(source_environment["MOSAIC_RUNTIME_GENERATION"])
        if generation < 0:
            raise ValueError("INVALID_GENERATION")
        reply = request(
            Path(socket_value),
            {
                "action": "authorize_tool",
                "session_id": session_id,
                "runtime_generation": generation,
                "runtime": arguments.runtime,
                "tool_name": tool_name,
            },
        )
    except (KeyError, ValueError, OSError, json.JSONDecodeError):
        return deny("GATE_UNAVAILABLE")

    if reply.get("ok") is True and reply.get("decision") == "allow":
        return 0
    code = reply.get("code")
    return deny(code if isinstance(code, str) else "MUTATOR_UNVERIFIED")


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