#!/usr/bin/env python3
"""
demo_python.py — comprehensive Python demo for connecting to a running n_signer
via Qubes qrexec and performing all three core operations:

  1. get_public_key  — retrieve a Nostr public key by nostr_index
  2. sign_event      — sign a Nostr event (kind 1 text note)
  3. nip44_encrypt   — encrypt a message to a peer (and decrypt it back)

Uses qrexec-client-vm (Qubes OS inter-qube IPC). No auth envelope needed —
identity comes from QREXEC_REMOTE_DOMAIN on the server side.

Prerequisites:
  - n_signer running in the target qube with --bridge-source-trusted
  - qubes.NsignerRpc service installed in the target qube
  - dom0 qrexec policy allowing this qube to call the service
  - Python 3 with no external dependencies (uses only stdlib)

Usage:
  python3 client/demo_python.py <target_qube> [nostr_index]
  python3 client/demo_python.py nostr_signer 1

If no nostr_index is given, defaults to 0.
"""

import json
import struct
import subprocess
import sys
import time


# --- Bech32 encoder (NIP-19 npub conversion, no external dependencies) ---

CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"


def bech32_polymod(values):
    generator = [0x3B6A57B2, 0x26508E6D, 0x1EA119FA, 0x3D4233DD, 0x2A1462B3]
    chk = 1
    for v in values:
        b = chk >> 25
        chk = (chk & 0x1FFFFFF) << 5 ^ v
        for i in range(5):
            chk ^= generator[i] if ((b >> i) & 1) else 0
    return chk


def bech32_hrp_expand(hrp):
    return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]


def bech32_create_checksum(hrp, data):
    values = bech32_hrp_expand(hrp) + data
    polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1
    return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]


def bech32_encode(hrp, data):
    combined = data + bech32_create_checksum(hrp, data)
    return hrp + "1" + "".join([CHARSET[d] for d in combined])


def convertbits(data, frombits, tobits, pad=True):
    acc = 0
    bits = 0
    ret = []
    maxv = (1 << tobits) - 1
    max_acc = (1 << (frombits + tobits - 1)) - 1
    for value in data:
        acc = ((acc << frombits) | value) & max_acc
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad and bits:
        ret.append((acc << (tobits - bits)) & maxv)
    return ret


def hex_to_npub(pubkey_hex):
    """Convert a 32-byte hex pubkey to bech32 npub format (NIP-19)."""
    pubkey_bytes = bytes.fromhex(pubkey_hex)
    data = convertbits(pubkey_bytes, 8, 5)
    return bech32_encode("npub", data)


# --- n_signer qrexec client ---

def call_nsigner(target_qube, request):
    """
    Call n_signer via qrexec. Sends one framed JSON-RPC request, receives one
    framed response. Each call spawns a fresh qrexec-client-vm process.

    Framing: 4-byte big-endian length prefix + JSON payload.
    No auth envelope needed for qrexec (identity from QREXEC_REMOTE_DOMAIN).
    """
    payload = json.dumps(request, separators=(",", ":")).encode("utf-8")
    frame = struct.pack(">I", len(payload)) + payload

    proc = subprocess.Popen(
        ["qrexec-client-vm", target_qube, "qubes.NsignerRpc"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    out, err = proc.communicate(frame)

    if proc.returncode != 0:
        raise RuntimeError(
            f"qrexec-client-vm exited with code {proc.returncode}: "
            f"{err.decode('utf-8', 'replace').strip()}"
        )

    if len(out) < 4:
        raise RuntimeError("short response (missing frame header)")

    length = struct.unpack(">I", out[:4])[0]
    body = out[4 : 4 + length]
    if len(body) != length:
        raise RuntimeError(
            f"short response payload: expected {length}, got {len(body)}"
        )

    return json.loads(body.decode("utf-8"))


# --- Demos ---

def demo_get_public_key(target_qube, nostr_index):
    """Demo 1: Get a public key by nostr_index."""
    print(f"\n=== Demo 1: get_public_key (nostr_index={nostr_index}) ===")

    response = call_nsigner(
        target_qube,
        {"id": "1", "method": "get_public_key", "params": [{"nostr_index": nostr_index}]},
    )

    if "error" in response:
        raise RuntimeError(f"get_public_key failed: {json.dumps(response['error'])}")

    pubkey_hex = response["result"]
    npub = hex_to_npub(pubkey_hex)
    print(f"  pubkey hex: {pubkey_hex}")
    print(f"  npub:       {npub}")
    return pubkey_hex


def demo_sign_event(target_qube, nostr_index, pubkey_hex):
    """Demo 2: Sign a Nostr event (kind 1 text note)."""
    print("\n=== Demo 2: sign_event (kind 1 text note) ===")

    unsigned_event = {
        "kind": 1,
        "content": "Hello from n_signer Python demo!",
        "created_at": int(time.time()),
        "tags": [],
        "pubkey": pubkey_hex,
    }

    print("  Unsigned event:")
    print(f"    {json.dumps(unsigned_event, separators=(',', ':'))}")

    response = call_nsigner(
        target_qube,
        {
            "id": "2",
            "method": "sign_event",
            "params": [json.dumps(unsigned_event, separators=(",", ":")), {"nostr_index": nostr_index}],
        },
    )

    if "error" in response:
        raise RuntimeError(f"sign_event failed: {json.dumps(response['error'])}")

    signed_event = json.loads(response["result"])
    print("  Signed event:")
    print(f"    {json.dumps(signed_event, separators=(',', ':'))}")
    print(f"  event id: {signed_event['id']}")
    print(f"  signature: {signed_event['sig']}")
    return signed_event


def demo_nip44(target_qube, nostr_index, pubkey_hex):
    """Demo 3: NIP-44 encrypt and decrypt."""
    plaintext = "Secret message from n_signer Python demo!"

    print("\n=== Demo 3: nip44_encrypt / nip44_decrypt ===")
    print(f'  plaintext:  "{plaintext}"')
    print(f"  peer pubkey: {pubkey_hex} (self)")

    # Encrypt
    enc_response = call_nsigner(
        target_qube,
        {
            "id": "3",
            "method": "nip44_encrypt",
            "params": [pubkey_hex, plaintext, {"nostr_index": nostr_index}],
        },
    )

    if "error" in enc_response:
        raise RuntimeError(f"nip44_encrypt failed: {json.dumps(enc_response['error'])}")

    ciphertext = enc_response["result"]
    print(f"  ciphertext: {ciphertext}")

    # Decrypt
    dec_response = call_nsigner(
        target_qube,
        {
            "id": "4",
            "method": "nip44_decrypt",
            "params": [pubkey_hex, ciphertext, {"nostr_index": nostr_index}],
        },
    )

    if "error" in dec_response:
        raise RuntimeError(f"nip44_decrypt failed: {json.dumps(dec_response['error'])}")

    decrypted = dec_response["result"]
    print(f'  decrypted:  "{decrypted}"')

    if plaintext == decrypted:
        print("  ✓ Round-trip verified: plaintext matches decrypted")
    else:
        raise RuntimeError("Round-trip FAILED: plaintext does not match decrypted")


# --- Main ---

def main():
    target_qube = sys.argv[1] if len(sys.argv) > 1 else "nostr_signer"
    nostr_index = int(sys.argv[2]) if len(sys.argv) > 2 else 0

    print("=== n_signer Python Demo ===")
    print(f"Target qube: {target_qube}")
    print(f"Service:     qubes.NsignerRpc")
    print(f"nostr_index: {nostr_index}")
    print("\nConnecting to n_signer via qrexec...")

    try:
        pubkey_hex = demo_get_public_key(target_qube, nostr_index)
        demo_sign_event(target_qube, nostr_index, pubkey_hex)
        demo_nip44(target_qube, nostr_index, pubkey_hex)

        print("\n=== Summary ===")
        print("All demos completed successfully.")
    except Exception as e:
        print("\n=== Summary ===")
        print(f"Demo failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()