#!/usr/bin/env python3
"""
ZKAC TCP client: load a member credential, complete handshake, request /api (encrypted JSON).

The "API" is not HTTP on port 8765 — it is one JSON request inside the ZKAC session on --port (default 9876).
"""

from __future__ import annotations

import argparse
import base64
import json
import socket
from pathlib import Path

import zkac
from zkac.tcp import FramedSession, client_handshake


def load_credential(member_json: Path) -> zkac.Credential:
    """Rebuild Credential from setup_demo.py output (same fields as zkac.Credential.finalize)."""
    m = json.loads(member_json.read_text(encoding="utf-8"))
    pk = zkac.BbsPublicKey.from_bytes(base64.b64decode(m["issuer_public_key_b64"]))
    rid = bytes.fromhex(m["role_id_hex"])
    return zkac.Credential.finalize(
        base64.b64decode(m["blind_sig_b64"]),
        base64.b64decode(m["member_secret_b64"]),
        base64.b64decode(m["prover_blind_b64"]),
        rid,
        int(m["epoch"]),
        pk,
    )


def load_server_pk(creds_dir: Path) -> zkac.PublicKey:
    """Pinned server identity: must match the Keypair used by server.py (from transport.json)."""
    t = json.loads((creds_dir / "transport.json").read_text(encoding="utf-8"))
    raw = base64.b64decode(t["server_public_key_b64"])
    return zkac.PublicKey.from_bytes(raw)


def main() -> None:
    ap = argparse.ArgumentParser(description="ZKAC demo client (TCP + credential)")
    ap.add_argument(
        "--creds-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "creds",
        help="Directory with transport.json and member_*.json",
    )
    ap.add_argument(
        "--member",
        type=Path,
        help="Path to member_*.json (default: creds-dir/member_analyst.json)",
    )
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=9876)
    args = ap.parse_args()

    creds_dir: Path = args.creds_dir
    member_path = args.member or (creds_dir / "member_analyst.json")
    if not member_path.is_file():
        raise SystemExit(f"Missing member file: {member_path}")

    credential = load_credential(member_path)
    server_pk = load_server_pk(creds_dir)

    # Ephemeral client transport identity (not the BBS+ member secret — that is inside credential).
    client_kp = zkac.Keypair()
    node = zkac.Node(client_kp)

    sock = socket.create_connection((args.host, args.port))
    try:
        # X25519 + server Schnorr + BBS+ auth; returns symmetric Session.
        session = client_handshake(sock, node, server_pk, credential)
        framed = FramedSession(sock, session)

        # Logical GET /api: path is checked by server after decrypt.
        request_obj = {"path": "/api"}
        payload = json.dumps(request_obj).encode("utf-8")
        framed.send(payload)

        reply = framed.recv().decode("utf-8")
        print(json.dumps(json.loads(reply), indent=2))
    finally:
        sock.close()


if __name__ == "__main__":
    main()