#!/usr/bin/env python3
"""P3 broker prototype: peercred-keyed runtime_generation and lease revocation."""

from __future__ import annotations

import argparse
import json
import os
import socket
import struct
from pathlib import Path
from typing import Any


def proc_starttime(pid: int) -> int:
    text = Path(f"/proc/{pid}/stat").read_text()
    close = text.rfind(")")
    return int(text[close + 2 :].split()[19])


def emit(log: Path, value: dict[str, Any]) -> None:
    with log.open("a", encoding="utf-8") as out:
        out.write(json.dumps(value, sort_keys=True) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--socket", required=True)
    parser.add_argument("--log", required=True)
    ns = parser.parse_args()
    socket_path = Path(ns.socket)
    log_path = Path(ns.log)
    socket_path.parent.mkdir(parents=True, exist_ok=True)
    os.chmod(socket_path.parent, 0o700)
    socket_path.unlink(missing_ok=True)
    log_path.unlink(missing_ok=True)

    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(socket_path))
    os.chmod(socket_path, 0o600)
    server.listen(8)
    generations: dict[tuple[int, int], int] = {}
    lease_state: dict[tuple[int, int], str] = {}
    emit(log_path, {"event": "listen", "pid": os.getpid(), "socket": str(socket_path)})

    while True:
        conn, _ = server.accept()
        with conn:
            raw = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
            pid, uid, gid = struct.unpack("3i", raw)
            starttime = proc_starttime(pid)
            request = json.loads(conn.makefile("r", encoding="utf-8").readline())
            if request.get("action") == "shutdown-broker":
                conn.sendall(b'{"ok":true}\n')
                break
            identity = (pid, starttime)
            if request.get("action") == "promote-probe":
                generation = generations.get(identity, 0)
                lease_state[identity] = "VERIFIED"
                record = {
                    "event": "probe_lease_promoted",
                    "peercred": {"pid": pid, "uid": uid, "gid": gid},
                    "starttime_ticks": starttime,
                    "generation": generation,
                    "new_lease_state": "VERIFIED",
                }
                emit(log_path, record)
                conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
                continue
            if request.get("action") == "source-invalid":
                old_lease = lease_state.get(identity, "NONE")
                lease_state[identity] = "REVOKED"
                record = {
                    "event": "source_invalidation_revoke",
                    "peercred": {"pid": pid, "uid": uid, "gid": gid},
                    "starttime_ticks": starttime,
                    "generation": generations.get(identity, 0),
                    "source_reason": request.get("reason"),
                    "prior_lease": old_lease,
                    "new_lease_state": "REVOKED",
                    "promotion": False,
                }
                emit(log_path, record)
                conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
                continue
            if request.get("action") != "lifecycle":
                conn.sendall(b'{"ok":false,"reason":"invalid-action"}\n')
                continue

            old_generation = generations.get(identity, 0)
            old_lease = lease_state.get(identity, "NONE")
            new_generation = old_generation + 1
            generations[identity] = new_generation
            # Every lifecycle boundary revokes first. A start establishes a new
            # UNVERIFIED incarnation; it never inherits prior VERIFIED state.
            lease_state[identity] = "UNVERIFIED" if request.get("phase") == "start" else "REVOKED"
            record = {
                "event": "runtime_generation_bump",
                "peercred": {"pid": pid, "uid": uid, "gid": gid},
                "starttime_ticks": starttime,
                "phase": request.get("phase"),
                "reason": request.get("reason"),
                "old_generation": old_generation,
                "new_generation": new_generation,
                "prior_lease": old_lease,
                "prior_lease_revoked": True,
                "new_lease_state": lease_state[identity],
            }
            emit(log_path, record)
            conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())

    server.close()
    socket_path.unlink(missing_ok=True)


if __name__ == "__main__":
    main()