#!/usr/bin/env python3
"""
E2E smoke + timing: one server, clients A (admin) and B (recipient).

  1. A creates a registry with 3 roles
  2. A posts one or more grants to B (same role ``beta``)
  3. Time B's mailbox fetch and permission-style checks (same phases as
     ``zkac-node credentials list B --server …``: local creds, list_pending, has_credential)

Default (no args): one grant, asserts one pending ``beta`` grant.

Scaling: ``--sizes 2,5,25,50`` runs a fresh server + pool for each size (all grants
to B), prints a timing table (``list_pending`` = tags + PIR full-row decode per match).

Run from repo root, e.g.:

  uv run python scripts/e2e_two_clients_timing.py
  uv run python scripts/e2e_two_clients_timing.py --sizes 2,5,25,50
"""

from __future__ import annotations

import argparse
import base64
import os
import socket
import sys
import tempfile
import threading
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "python"))
sys.path.insert(0, str(ROOT / "cli"))


def _free_port() -> int:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("127.0.0.1", 0))
    _, port = s.getsockname()
    s.close()
    return port


def _run_scaled_sizes(sizes: list[int]) -> int:
    from zkac_cli import client, store
    from zkac_cli.server import _ServerStore, serve

    def log(msg: str) -> None:
        print(msg, flush=True)

    rows: list[tuple[int, float, float, float, float, float, float]] = []

    for n in sizes:
        td = tempfile.mkdtemp(prefix="zkac-e2e-")
        os.environ["ZKAC_HOME"] = td
        port = _free_port()
        server = f"127.0.0.1:{port}"
        server_dd = Path(td) / "srv"
        server_dd.mkdir(parents=True)

        ss = _ServerStore(server_dd)
        kp = ss.load_or_create_keypair()
        pk_b64 = base64.b64encode(kp.public_key().to_bytes()).decode()

        store.create_user("A")
        store.create_user("B")
        store.pin_server("A", server, pk_b64)
        store.pin_server("B", server, pk_b64)

        t_srv = threading.Thread(
            target=lambda: serve(str(server_dd), "127.0.0.1", port),
            daemon=True,
        )
        t_srv.start()
        time.sleep(0.25)

        t_setup0 = time.perf_counter()
        rid = client.create_registry("A", server, ["alpha", "beta", "gamma"])
        t_after_create = time.perf_counter()
        b_pk = store.load_identity("B")["issuance_pk"].hex()
        for _ in range(n):
            client.grant("A", server, rid, "beta", b_pk)
        t_gr1 = time.perf_counter()

        t_loc0 = time.perf_counter()
        local_creds = store.list_credentials("B")
        t_loc1 = time.perf_counter()

        t_mail0 = time.perf_counter()
        pending = client.list_pending("B", server)
        t_mail1 = time.perf_counter()

        t_perm0 = time.perf_counter()
        for g in pending:
            r = g.get("registry_id")
            role = g.get("role_name")
            if r not in (None, "?") and role not in (None, "?"):
                store.has_credential("B", r, role)
        t_perm1 = time.perf_counter()

        create_ms = (t_after_create - t_setup0) * 1000
        grant_ms = (t_gr1 - t_after_create) * 1000
        local_ms = (t_loc1 - t_loc0) * 1000
        mail_ms = (t_mail1 - t_mail0) * 1000
        perm_ms = (t_perm1 - t_perm0) * 1000
        total_ms = (t_perm1 - t_loc0) * 1000

        rows.append((n, create_ms, grant_ms, local_ms, mail_ms, perm_ms, total_ms))

        log(
            f"n={n}: create_registry={create_ms:.0f} ms  N×grant={grant_ms:.0f} ms  "
            f"list_pending={mail_ms:.0f} ms  pending={len(pending)}  ZKAC_HOME={td}"
        )
        if len(pending) != n:
            print(f"ERROR: expected {n} pending, got {len(pending)}", flush=True)
            return 1
        for p in pending:
            if p.get("role_name") != "beta" or p.get("registry_id") != rid:
                print(f"ERROR: bad pending row {p!r}", flush=True)
                return 1

    print()
    print(
        "pool_n | create_registry (ms) | N×grant (ms) | list_local (ms) | "
        "list_pending mailbox (ms) | has_cred (ms) | cred_list_total (ms)"
    )
    print("-" * 120)
    for n, create_ms, grant_ms, local_ms, mail_ms, perm_ms, total_ms in rows:
        print(
            f"{n:6d} | {create_ms:20.1f} | {grant_ms:12.1f} | {local_ms:15.3f} | "
            f"{mail_ms:25.1f} | {perm_ms:12.3f} | {total_ms:20.1f}"
        )
    print("OK")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="ZKAC e2e timing (mailbox / credentials list)")
    parser.add_argument(
        "--sizes",
        default=None,
        metavar="N,N,...",
        help="comma-separated pool sizes (each run: fresh server, N grants to B, then list_pending)",
    )
    args = parser.parse_args()

    if args.sizes is not None:
        sizes = [int(x.strip()) for x in args.sizes.split(",") if x.strip()]
        if not sizes or any(x < 1 for x in sizes):
            print("error: --sizes must be positive integers", file=sys.stderr)
            return 2
        return _run_scaled_sizes(sizes)

    from zkac_cli import client, store
    from zkac_cli.server import _ServerStore, serve

    def log(msg: str) -> None:
        print(msg, flush=True)

    td = tempfile.mkdtemp(prefix="zkac-e2e-")
    os.environ["ZKAC_HOME"] = td
    port = _free_port()
    server = f"127.0.0.1:{port}"
    server_dd = Path(td) / "srv"
    server_dd.mkdir(parents=True)

    ss = _ServerStore(server_dd)
    kp = ss.load_or_create_keypair()
    pk_b64 = base64.b64encode(kp.public_key().to_bytes()).decode()

    store.create_user("A")
    store.create_user("B")
    store.pin_server("A", server, pk_b64)
    store.pin_server("B", server, pk_b64)

    t_srv = threading.Thread(
        target=lambda: serve(str(server_dd), "127.0.0.1", port),
        daemon=True,
    )
    t_srv.start()
    time.sleep(0.25)
    log("server thread up")

    t0 = time.perf_counter()
    rid = client.create_registry("A", server, ["alpha", "beta", "gamma"])
    t_create = time.perf_counter()
    log(f"registry created ({(t_create - t0) * 1000:.0f} ms)")

    b_pk = store.load_identity("B")["issuance_pk"].hex()
    client.grant("A", server, rid, "beta", b_pk)
    t_grant = time.perf_counter()
    log(f"grant posted ({(t_grant - t_create) * 1000:.0f} ms)")

    log("B mailbox + permission-style checks (same work as `credentials list`) …")
    t_loc0 = time.perf_counter()
    local_creds = store.list_credentials("B")
    t_loc1 = time.perf_counter()

    t_mail0 = time.perf_counter()
    pending = client.list_pending("B", server)
    t_mail1 = time.perf_counter()

    t_perm0 = time.perf_counter()
    for g in pending:
        r = g.get("registry_id")
        role = g.get("role_name")
        if r not in (None, "?") and role not in (None, "?"):
            store.has_credential("B", r, role)
    t_perm1 = time.perf_counter()

    log(f"ZKAC_HOME={td}")
    log(f"server={server} registry={rid[:24]}…")
    print(f"create_registry: {(t_create - t0) * 1000:.1f} ms")
    print(f"grant:           {(t_grant - t_create) * 1000:.1f} ms")
    print(
        f"list_local_creds(B):   {(t_loc1 - t_loc0) * 1000:.3f} ms  ({len(local_creds)} on disk)"
    )
    print(
        f"list_pending(mailbox): {(t_mail1 - t_mail0) * 1000:.1f} ms  "
        f"({len(pending)} match(es); tags + PIR row + decrypt)"
    )
    print(
        f"has_credential checks: {(t_perm1 - t_perm0) * 1000:.3f} ms  "
        f"({len(pending)} grant(s))"
    )
    print(
        f"credentials_list_total: {(t_perm1 - t_loc0) * 1000:.1f} ms  "
        "(local + mailbox + permission flags)"
    )
    for p in pending:
        print(
            f"  pending: registry={p.get('registry_id', '?')[:16]}… "
            f"role={p.get('role_name')} idx={p.get('pool_index')}"
        )

    assert len(pending) == 1, pending
    assert pending[0].get("role_name") == "beta"
    assert pending[0].get("registry_id") == rid
    print("OK")
    return 0


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