"""Client-side operations over a unified encrypted channel.

All management and grant traffic uses the same anonymous handshake
(X25519 + Schnorr server identity proof) as authenticated sessions.

Grant flow (admin-initiated, server is trustless):
  1. Admin knows recipient's X25519 issuance public key (out-of-band).
  2. Admin locally prepares a blind request and issues the blind signature.
  3. Admin encrypts the full credential material to the recipient's pk.
  4. Admin posts the opaque ciphertext to the server's mailbox, authenticated
     by a BBS+ admin presentation bound to the session transcript.
  5. Recipient trial-decrypts each pending entry to find matching grants.

The server sees only: recipient_pk, eph_pk, ciphertext, grant_id.
"""

from __future__ import annotations

import base64
import json
import socket

import zkac
from zkac.tcp import FramedSession, client_handshake_anon, client_handshake_managed

from . import store


def _b64(data: bytes) -> str:
    return base64.b64encode(data).decode()


def _unb64(s: str) -> bytes:
    return base64.b64decode(s)


def _parse_server(server: str) -> tuple[str, int]:
    host, _, port = server.rpartition(":")
    return host or "127.0.0.1", int(port)


def parse_spec(spec: str) -> tuple[str, str, str]:
    """Parse 'host:port:registry_id:role' into (server, registry_id, role)."""
    parts = spec.rsplit(":", 2)
    if len(parts) != 3:
        raise ValueError(f"invalid spec {spec!r}, expected host:port:registry_id:role")
    return parts[0], parts[1], parts[2]


# ── Encrypted management session ─────────────────────────────────────

def _resolve_server_pk(server: str) -> zkac.PublicKey:
    """Load pinned server public key or fail."""
    pin = store.load_server_pin(server)
    if pin is None:
        raise RuntimeError(
            f"no pinned key for {server}; run: zkac-node server pin {server} --key <hex>"
        )
    return zkac.PublicKey.from_bytes(_unb64(pin["server_public_key_b64"]))


def _mgmt_connect(server: str) -> tuple[socket.socket, FramedSession]:
    """Open an encrypted management session to the server."""
    host, port = _parse_server(server)
    sock = socket.create_connection((host, port))
    server_pk = _resolve_server_pk(server)
    node = zkac.Node(zkac.Keypair())
    session = client_handshake_anon(sock, node, server_pk)
    framed = FramedSession(sock, session)
    framed.send(json.dumps({"op": "mgmt"}).encode())
    return sock, framed


def _mgmt_cmd(framed: FramedSession, cmd: dict) -> dict:
    framed.send(json.dumps(cmd).encode())
    return json.loads(framed.recv())


def _ok(resp: dict) -> dict:
    if resp.get("error"):
        raise RuntimeError(resp["error"])
    return resp


# ── Registry operations ──────────────────────────────────────────────

def create_registry(server: str, role_names: list[str]) -> str:
    identity = store.load_identity()
    admin_mat = store.new_admin_material()
    bbs_issuer, bbs_pk, admin_cred = store.reconstruct_admin(admin_mat)

    role_entries = [(zkac.role_id(name), bbs_pk, 1) for name in role_names]
    state = zkac.RegistryState.build(
        bbs_pk, identity["issuance_pk"], 1, b"\x00" * 32, role_entries,
    )
    state_bytes = state.serialize()
    state_cert = state.certify(admin_cred)
    registry_id = state.registry_id()

    sock, framed = _mgmt_connect(server)
    try:
        resp = _ok(_mgmt_cmd(framed, {
            "cmd": "create_registry",
            "state_bytes_b64": _b64(state_bytes),
            "state_cert_b64": _b64(bytes(state_cert)),
        }))
    finally:
        sock.close()

    rid_hex = resp["registry_id"]
    store.save_admin(rid_hex, {
        "server": server,
        "roles": role_names,
        **admin_mat,
    })
    return rid_hex


def update_registry(server: str, registry_id_hex: str, add_roles: list[str]):
    admin_data = store.load_admin(registry_id_hex)
    bbs_issuer, bbs_pk, admin_cred = store.reconstruct_admin(admin_data)
    identity = store.load_identity()

    sock, framed = _mgmt_connect(server)
    try:
        cur = _ok(_mgmt_cmd(framed, {
            "cmd": "get_registry", "registry_id": registry_id_hex,
        }))

        old_state = zkac.RegistryState.deserialize(_unb64(cur["state_bytes_b64"]))
        prev_hash = old_state.state_hash()
        new_version = old_state.version() + 1

        old_roles = admin_data.get("roles", [])
        all_roles = list(old_roles) + [r for r in add_roles if r not in old_roles]
        role_entries = [(zkac.role_id(name), bbs_pk, 1) for name in all_roles]

        new_state = zkac.RegistryState.build(
            bbs_pk, identity["issuance_pk"], new_version, bytes(prev_hash), role_entries,
        )
        new_cert = new_state.certify(admin_cred)

        _ok(_mgmt_cmd(framed, {
            "cmd": "update_registry",
            "registry_id": registry_id_hex,
            "state_bytes_b64": _b64(new_state.serialize()),
            "state_cert_b64": _b64(bytes(new_cert)),
        }))
    finally:
        sock.close()

    admin_data["roles"] = all_roles
    store.save_admin(registry_id_hex, admin_data)


def get_registry(server: str, registry_id_hex: str) -> dict:
    sock, framed = _mgmt_connect(server)
    try:
        return _ok(_mgmt_cmd(framed, {
            "cmd": "get_registry", "registry_id": registry_id_hex,
        }))
    finally:
        sock.close()


def list_own_registries() -> list[dict]:
    result = []
    for rid in store.list_admin_registries():
        data = store.load_admin(rid)
        result.append({
            "registry_id": rid,
            "server": data.get("server", "?"),
            "roles": data.get("roles", []),
        })
    return result


# ── Admin-initiated grant ────────────────────────────────────────────

def grant(server: str, registry_id_hex: str, role_name: str,
          recipient_pk_hex: str) -> str:
    """Admin issues a credential and posts it (encrypted) to the recipient's mailbox."""
    admin_data = store.load_admin(registry_id_hex)
    roles = admin_data.get("roles", [])
    if role_name not in roles:
        raise RuntimeError(f"role {role_name!r} not in registry (have: {roles})")

    bbs_issuer, bbs_pk, admin_cred = store.reconstruct_admin(admin_data)
    role_rid = zkac.role_id(role_name)
    epoch = 1

    req = zkac.prepare_blind_request()
    blind_sig = bbs_issuer.issue_blind(req.commitment_with_proof(), role_rid, epoch)

    payload = json.dumps({
        "registry_id": registry_id_hex,
        "role_name": role_name,
        "epoch": epoch,
        "issuer_pk_b64": _b64(bbs_pk.to_bytes()),
        "blind_sig_b64": _b64(blind_sig),
        "member_secret_b64": _b64(req.member_secret()),
        "prover_blind_b64": _b64(req.prover_blind()),
    }).encode()

    recipient_pk = bytes.fromhex(recipient_pk_hex)
    eph_kp = zkac.IssuanceKeypair()
    ciphertext = eph_kp.encrypt(recipient_pk, payload)

    sock, framed = _mgmt_connect(server)
    try:
        transcript_hash = bytes(framed.session.transcript_hash())
        admin_proof = admin_cred.present(transcript_hash)
        resp = _ok(_mgmt_cmd(framed, {
            "cmd": "post_grant",
            "registry_id": registry_id_hex,
            "recipient_pk_hex": recipient_pk_hex,
            "eph_pk_b64": _b64(eph_kp.public_key_bytes()),
            "ciphertext_b64": _b64(ciphertext),
            "admin_proof_b64": _b64(admin_proof),
        }))
    finally:
        sock.close()

    return resp["grant_id"]


# ── Receiver: list + collect (trial-decrypt) ─────────────────────────

def list_pending(server: str) -> list[dict]:
    identity = store.load_identity()
    pk_hex = identity["issuance_pk"].hex()
    sock, framed = _mgmt_connect(server)
    try:
        resp = _ok(_mgmt_cmd(framed, {"cmd": "list_grants", "recipient_pk_hex": pk_hex}))
    finally:
        sock.close()

    receiver_kp = zkac.IssuanceKeypair.from_secret(identity["issuance_sk"])
    results = []
    for entry in resp["grants"]:
        try:
            eph_pk = _unb64(entry["eph_pk_b64"])
            ct = _unb64(entry["ciphertext_b64"])
            plaintext = json.loads(receiver_kp.decrypt(eph_pk, ct))
            results.append({
                "grant_id": entry["grant_id"],
                "registry_id": plaintext.get("registry_id", "?"),
                "role_name": plaintext.get("role_name", "?"),
            })
        except Exception:
            results.append({
                "grant_id": entry["grant_id"],
                "registry_id": "?",
                "role_name": "(undecryptable)",
            })
    return results


def collect(spec: str) -> dict:
    """Fetch and finalize a pending credential described by 'host:port:registry_id:role'."""
    server, registry_id_hex, role_name = parse_spec(spec)
    identity = store.load_identity()
    pk_hex = identity["issuance_pk"].hex()
    receiver_kp = zkac.IssuanceKeypair.from_secret(identity["issuance_sk"])

    sock, framed = _mgmt_connect(server)
    try:
        pending = _ok(_mgmt_cmd(framed, {
            "cmd": "list_grants", "recipient_pk_hex": pk_hex,
        }))["grants"]

        # Trial-decrypt to find matching grant
        target_grant_id = None
        target_payload = None
        for entry in pending:
            try:
                eph_pk = _unb64(entry["eph_pk_b64"])
                ct = _unb64(entry["ciphertext_b64"])
                plaintext = json.loads(receiver_kp.decrypt(eph_pk, ct))
                if (plaintext.get("registry_id") == registry_id_hex and
                        plaintext.get("role_name") == role_name):
                    target_grant_id = entry["grant_id"]
                    target_payload = plaintext
                    break
            except Exception:
                continue

        if target_grant_id is None:
            raise RuntimeError(f"no pending grant for {spec}")

        _ok(_mgmt_cmd(framed, {
            "cmd": "claim_grant",
            "recipient_pk_hex": pk_hex,
            "grant_id": target_grant_id,
        }))

        reg_info = _ok(_mgmt_cmd(framed, {
            "cmd": "get_registry", "registry_id": registry_id_hex,
        }))
    finally:
        sock.close()

    cred_data = {
        "blind_sig_b64": target_payload["blind_sig_b64"],
        "member_secret_b64": target_payload["member_secret_b64"],
        "prover_blind_b64": target_payload["prover_blind_b64"],
        "role_name": role_name,
        "epoch": target_payload["epoch"],
        "issuer_pk_b64": target_payload["issuer_pk_b64"],
    }
    cred = store.reconstruct_credential(cred_data)
    cred.present(b"self-test")

    store.save_credential(registry_id_hex, role_name, cred_data)

    # Pin server + stash registry metadata locally
    pin = store.load_server_pin(server)
    if pin:
        server_pk_b64 = pin["server_public_key_b64"]
    else:
        server_pk_b64 = _b64(b"\x00" * 32)
    store.pin_server(server, server_pk_b64)

    return {"registry_id": registry_id_hex, "role": role_name, "server": server}


# ── Authenticated session ─────────────────────────────────────────────

def authenticate(registry_id_hex: str, role_name: str,
                 server: str | None = None) -> dict:
    admin_data = None
    try:
        admin_data = store.load_admin(registry_id_hex)
    except FileNotFoundError:
        pass

    if server is None:
        if admin_data and admin_data.get("server"):
            server = admin_data["server"]
        else:
            raise RuntimeError("server address required (--server host:port)")

    cred_data = store.load_credential_data(registry_id_hex, role_name)
    cred = store.reconstruct_credential(cred_data)

    server_pk = _resolve_server_pk(server)
    node = zkac.Node(zkac.Keypair())
    host, port = _parse_server(server)

    sock = socket.create_connection((host, port))
    try:
        session = client_handshake_anon(sock, node, server_pk)
        framed = FramedSession(sock, session)

        transcript_hash = bytes(session.transcript_hash())
        auth_proof = cred.present(transcript_hash)
        role_rid = zkac.role_id(role_name)

        framed.send(json.dumps({
            "op": "auth",
            "registry_id": registry_id_hex,
            "role_id": role_rid.hex(),
            "bbs_auth_b64": _b64(auth_proof),
        }).encode())

        return json.loads(framed.recv())
    finally:
        sock.close()