#!/usr/bin/env python3
"""
ZKAC TCP server using client-managed registries.

Loads a registry from managed_registry.json (created by setup_managed_demo.py),
verifies BBS+ state certificates, and authenticates clients against the registry.
Also handles credential issuance requests via the E2E-encrypted relay.

Run setup_managed_demo.py first, then this server, then client_managed.py.
"""

from __future__ import annotations

import argparse
import base64
import json
import socket
import struct
import threading
import traceback
from pathlib import Path

import zkac
from zkac.tcp import (
    FramedSession,
    read_frame,
    write_frame,
    server_handshake_managed,
)


def load_managed_registry(creds_dir: Path, mgr: zkac.RegistryManager) -> bytes:
    """Load the managed registry state + cert into the manager. Returns registry_id."""
    r = json.loads((creds_dir / "managed_registry.json").read_text(encoding="utf-8"))
    state_bytes = base64.b64decode(r["state_bytes_b64"])
    state_cert = base64.b64decode(r["state_cert_b64"])
    rid = mgr.create(state_bytes, state_cert)
    return rid


def _role_label(role_id: bytes) -> str:
    for name in ("analyst", "operator"):
        if role_id == zkac.role_id(name):
            return name
    return role_id.hex()[:16]


def api_body_for_role(role_id: bytes) -> dict:
    if role_id == zkac.role_id("analyst"):
        return {
            "path": "/api",
            "role": "analyst",
            "datasets": ["summary", "aggregated_metrics"],
            "note": "Analyst tier: aggregated data only.",
            "registry": "client-managed",
        }
    if role_id == zkac.role_id("operator"):
        return {
            "path": "/api",
            "role": "operator",
            "datasets": ["summary", "aggregated_metrics", "raw_logs", "pii"],
            "note": "Operator tier: full access.",
            "registry": "client-managed",
        }
    return {"error": "unknown role", "path": "/api"}


def handle_client(
    conn: socket.socket,
    addr: tuple,
    creds_dir: Path,
    mgr: zkac.RegistryManager,
) -> None:
    peer = f"{addr[0]}:{addr[1]}"
    print(f"[zkac-managed] connect peer={peer}")

    try:
        t = json.loads((creds_dir / "transport.json").read_text(encoding="utf-8"))
        sk = base64.b64decode(t["server_secret_key_b64"])
        server_kp = zkac.Keypair.from_secret_key(sk)
        node = zkac.Node(server_kp)

        session, registry_id, role_id = server_handshake_managed(conn, node, mgr)
        label = _role_label(role_id)
        print(
            f"[zkac-managed] auth_ok peer={peer} registry={registry_id.hex()[:16]}… "
            f"role={label!r}"
        )

        framed = FramedSession(conn, session)
        raw = framed.recv()
        req = json.loads(raw.decode("utf-8"))
        print(f"[zkac-managed] request peer={peer} {req!r}")

        path = req.get("path")
        if path != "/api":
            body = {"error": "unsupported path", "got": path}
        else:
            body = api_body_for_role(role_id)

        out = json.dumps(body).encode()
        framed.send(out)
        print(f"[zkac-managed] response peer={peer} {list(body.keys())}")

    except (ConnectionError, BrokenPipeError, OSError) as e:
        print(f"[zkac-managed] peer={peer} connection_error: {e!r}")
    except (json.JSONDecodeError, ValueError) as e:
        print(f"[zkac-managed] peer={peer} protocol_error: {e!r}")
    except Exception as e:
        print(f"[zkac-managed] peer={peer} unexpected_error: {e!r}")
        traceback.print_exc()
    finally:
        conn.close()


def main() -> None:
    ap = argparse.ArgumentParser(description="ZKAC managed-registry TCP server")
    ap.add_argument("--creds-dir", type=Path,
                    default=Path(__file__).resolve().parent / "creds")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=9877)
    args = ap.parse_args()

    creds_dir: Path = args.creds_dir
    if not (creds_dir / "managed_registry.json").is_file():
        raise SystemExit(
            f"Missing {creds_dir}/managed_registry.json — run setup_managed_demo.py first."
        )

    mgr = zkac.RegistryManager()
    rid = load_managed_registry(creds_dir, mgr)
    print(f"Loaded managed registry {rid.hex()[:16]}…")
    print(f"ZKAC managed  {args.host}:{args.port}")

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((args.host, args.port))
    sock.listen(8)
    while True:
        conn, addr = sock.accept()
        threading.Thread(
            target=handle_client,
            args=(conn, addr, creds_dir, mgr),
            daemon=True,
        ).start()


if __name__ == "__main__":
    main()