#!/usr/bin/env python3
"""Full test suite for the Teensy 4.1 n_signer over USB CDC.

Tests all verbs: get_info, get_public_key (all 6 algorithms), sign (secp256k1,
ed25519, ml-dsa-65, slh-dsa-128s), verify, encrypt/decrypt (OTP pad),
encapsulate/decapsulate (ML-KEM-768), derive_shared_secret (X25519), derive,
nostr_get_public_key, nostr_sign_event, nostr_nip04_encrypt/decrypt,
nostr_nip44_encrypt/decrypt.

Ordering note: the PQ keygens (ml-dsa-65, slh-dsa-128s, ml-kem-768) are
heap-heavy on the Teensy 4.1 (~140 KB free heap). Running all three back to
back can exhaust/fragment the heap and crash the device. To keep the suite
useful, the classical + nostr verbs (the common case) are tested FIRST on a
fresh heap, then the PQ verbs are tested LAST. If a PQ verb crashes the
device, the earlier results still stand.

Sends requests in the canonical n_signer wire format (see api.md §4.3):
  params is a JSON ARRAY of positional arguments, with the options object
  (containing "algorithm", "index", "scheme", ...) as the LAST element.
  Binary payloads (messages, pubkeys, ciphertexts) are hex-encoded.

Usage:
    python3 firmware/teensy41/test_signer.py [--port /dev/ttyACM0]
"""

import serial
import struct
import json
import time
import sys
import argparse
import hashlib
import os
import base64

DEFAULT_PORT = "/dev/ttyACM0"
BAUD = 115200


def send_request(ser, req: dict) -> dict:
    """Send a JSON-RPC request with 4-byte big-endian length prefix, read response."""
    payload = json.dumps(req).encode("utf-8")
    header = struct.pack(">I", len(payload))
    ser.write(header + payload)
    ser.flush()

    resp_header = b""
    deadline = time.time() + 30.0
    while len(resp_header) < 4 and time.time() < deadline:
        chunk = ser.read(4 - len(resp_header))
        if chunk:
            resp_header += chunk
        else:
            time.sleep(0.01)
    if len(resp_header) < 4:
        raise TimeoutError("Timeout reading response header")

    resp_len = struct.unpack(">I", resp_header)[0]
    if resp_len == 0 or resp_len > 65536:
        raise ValueError(f"Invalid response length: {resp_len}")

    resp_payload = b""
    while len(resp_payload) < resp_len:
        chunk = ser.read(resp_len - len(resp_payload))
        if chunk:
            resp_payload += chunk
        else:
            time.sleep(0.01)
    return json.loads(resp_payload.decode("utf-8"))


def test_verb(ser, name, params=None, id_counter=[0]):
    id_counter[0] += 1
    req = {"jsonrpc": "2.0", "id": id_counter[0], "method": name}
    if params is not None:
        req["params"] = params

    print(f"\n--- {name} ---")
    print(f"  params: {json.dumps(params, indent=2) if params else '(none)'}")

    try:
        resp = send_request(ser, req)
    except Exception as e:
        print(f"  ❌ FAIL: {e}")
        return None

    if "error" in resp:
        print(f"  ❌ FAIL: error code={resp['error'].get('code')} message={resp['error'].get('message')}")
        return resp
    elif "result" in resp:
        result = resp["result"]
        display = json.dumps(result, indent=2)
        if len(display) > 500:
            display = display[:500] + "... (truncated)"
        print(f"  ✅ PASS: {display}")
        return resp
    else:
        print(f"  ❌ FAIL: no result or error in response: {resp}")
        return resp


def main():
    parser = argparse.ArgumentParser(description="Test Teensy 4.1 n_signer")
    parser.add_argument("--port", default=DEFAULT_PORT, help="Serial port (default: /dev/ttyACM0)")
    args = parser.parse_args()

    print(f"Connecting to {args.port}...")
    ser = serial.Serial(args.port, BAUD, timeout=2.0)
    time.sleep(5)  # wait for auto-generate boot + key derivation

    # Drain any boot messages
    while ser.in_waiting:
        boot_msg = ser.read(ser.in_waiting).decode("utf-8", errors="replace")
        print(f"Boot: {boot_msg}", end="")
    print()

    passed = 0
    failed = 0
    pubkeys = {}

    # 1. get_info
    r = test_verb(ser, "get_info")
    if r and "result" in r: passed += 1
    else: failed += 1

    # 2. get_public_key for classical algorithms (lightweight, tested first)
    classical_algs = ["secp256k1", "ed25519", "x25519"]
    for alg in classical_algs:
        r = test_verb(ser, "get_public_key", [{"algorithm": alg, "index": 0}])
        if r and "result" in r:
            passed += 1
            pubkeys[alg] = r["result"].get("public_key", "")
        else:
            failed += 1

    # 3. sign + verify (secp256k1 schnorr)
    msg32 = hashlib.sha256(b"test message for signing").digest()
    msg_hex = msg32.hex()
    r = test_verb(ser, "sign", [msg_hex, {"algorithm": "secp256k1", "index": 0, "scheme": "schnorr"}])
    sig_secp = None
    if r and "result" in r:
        passed += 1
        sig_secp = r["result"].get("signature", "")
    else:
        failed += 1

    if sig_secp and "secp256k1" in pubkeys:
        r = test_verb(ser, "verify", [msg_hex, sig_secp, {"algorithm": "secp256k1", "index": 0, "scheme": "schnorr"}])
        if r and "result" in r and r["result"].get("valid") in (True, "true"):
            passed += 1
        else:
            failed += 1
    else:
        print("\n--- verify (secp256k1) ---")
        print("  ⏭️ SKIP: no signature or pubkey")
        failed += 1

    # 4. sign (ed25519)
    msg_raw = b"test ed25519 message"
    msg_raw_hex = msg_raw.hex()
    r = test_verb(ser, "sign", [msg_raw_hex, {"algorithm": "ed25519", "index": 0}])
    sig_ed = None
    if r and "result" in r:
        passed += 1
        sig_ed = r["result"].get("signature", "")
    else:
        failed += 1

    if sig_ed and "ed25519" in pubkeys:
        r = test_verb(ser, "verify", [msg_raw_hex, sig_ed, {"algorithm": "ed25519", "index": 0}])
        if r and "result" in r and r["result"].get("valid") in (True, "true"):
            passed += 1
        else:
            failed += 1
    else:
        print("\n--- verify (ed25519) ---")
        print("  ⏭️ SKIP: no signature or pubkey")
        failed += 1

    # 5. encrypt + decrypt (OTP pad)
    plaintext = b"Hello, OTP pad encryption test!"
    pt_b64 = base64.b64encode(plaintext).decode()
    r = test_verb(ser, "encrypt", [pt_b64, {"algorithm": "otp"}])
    ct_b64 = None
    pad_off_before = None
    if r and "result" in r:
        passed += 1
        ct_b64 = r["result"].get("result", "")
        # The pad offset where this ciphertext's pad slice begins. decrypt
        # must rewind to this offset so the same pad bytes are reused.
        pad_off_before = r["result"].get("pad_offset_before")
    else:
        failed += 1

    if ct_b64 and pad_off_before is not None:
        r = test_verb(ser, "decrypt",
                      [ct_b64, {"algorithm": "otp", "pad_offset": int(pad_off_before)}])
        if r and "result" in r:
            pt_result = base64.b64decode(r["result"].get("result", ""))
            if pt_result == plaintext:
                print(f"  ✅ decrypt matches original plaintext")
                passed += 1
            else:
                print(f"  ❌ decrypt mismatch: expected {plaintext}, got {pt_result}")
                failed += 1
        else:
            failed += 1
    else:
        print("\n--- decrypt (otp) ---")
        print("  ⏭️ SKIP: no ciphertext")
        failed += 1

    # 6. derive_shared_secret (X25519)
    try:
        from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
        peer_priv = X25519PrivateKey.generate()
        peer_pub = peer_priv.public_key()
        peer_pub_hex = peer_pub.public_bytes_raw().hex()
        r = test_verb(ser, "derive_shared_secret", [peer_pub_hex, {"algorithm": "x25519", "index": 0}])
        if r and "result" in r: passed += 1
        else: failed += 1
    except ImportError:
        print("\n--- derive_shared_secret (x25519) ---")
        print("  ⏭️ SKIP: 'cryptography' module not installed")
        failed += 1

    # 7. derive (secp256k1 at index 1)
    r = test_verb(ser, "derive", ["derive-test-data", {"algorithm": "secp256k1", "index": 1}])
    if r and "result" in r: passed += 1
    else: failed += 1

    # 8. nostr_get_public_key
    r = test_verb(ser, "nostr_get_public_key", [{"nostr_index": 0}])
    nostr_pub = None
    if r and "result" in r:
        passed += 1
        nostr_pub = r["result"]
        if isinstance(nostr_pub, dict):
            nostr_pub = nostr_pub.get("public_key", "")
    else:
        failed += 1

    # 9. nostr_sign_event
    if nostr_pub:
        event = {"kind": 1, "created_at": int(time.time()), "tags": [], "content": "hello from test_signer"}
        r = test_verb(ser, "nostr_sign_event", [event, {"nostr_index": 0}])
        if r and "result" in r: passed += 1
        else: failed += 1

    # 10. nostr_nip04_encrypt + decrypt (the bug we fixed)
    if nostr_pub:
        nip04_pt = "hello via nip04"
        r = test_verb(ser, "nostr_nip04_encrypt", [nostr_pub, nip04_pt, {"nostr_index": 0}])
        if r and "result" in r:
            passed += 1
            cipher = r["result"]
            r = test_verb(ser, "nostr_nip04_decrypt", [nostr_pub, cipher, {"nostr_index": 0}])
            if r and "result" in r and r["result"] == nip04_pt:
                print(f"  ✅ nip04 round-trip plaintext recovered")
                passed += 1
            else:
                print(f"  ❌ nip04 round-trip mismatch")
                failed += 1
        else:
            failed += 1

    # 11. nostr_nip44_encrypt + decrypt (the is_nip44 dispatch bug we fixed)
    if nostr_pub:
        nip44_pt = "hello via nip44"
        r = test_verb(ser, "nostr_nip44_encrypt", [nostr_pub, nip44_pt, {"nostr_index": 0}])
        if r and "result" in r:
            passed += 1
            cipher44 = r["result"]
            r = test_verb(ser, "nostr_nip44_decrypt", [nostr_pub, cipher44, {"nostr_index": 0}])
            if r and "result" in r and r["result"] == nip44_pt:
                print(f"  ✅ nip44 round-trip plaintext recovered")
                passed += 1
            else:
                print(f"  ❌ nip44 round-trip mismatch")
                failed += 1
        else:
            failed += 1

    # ---- PQ verbs (tested LAST: heap-heavy, may crash the device) ----
    print("\n=== PQ verbs (heap-heavy; tested last) ===")

    # 12. get_public_key for PQ algorithms
    pq_algs = ["ml-dsa-65", "slh-dsa-128s", "ml-kem-768"]
    for alg in pq_algs:
        r = test_verb(ser, "get_public_key", [{"algorithm": alg, "index": 0}])
        if r and "result" in r:
            passed += 1
            pubkeys[alg] = r["result"].get("public_key", "")
        else:
            failed += 1

    # 13. sign (ml-dsa-65)
    r = test_verb(ser, "sign", [msg_raw_hex, {"algorithm": "ml-dsa-65", "index": 0}])
    if r and "result" in r: passed += 1
    else: failed += 1

    # 14. sign (slh-dsa-128s) — slow (~1-2s on Teensy)
    print("\n--- sign (slh-dsa-128s) — may take 1-2 seconds ---")
    r = test_verb(ser, "sign", [msg_raw_hex, {"algorithm": "slh-dsa-128s", "index": 0}])
    if r and "result" in r: passed += 1
    else: failed += 1

    # 15. encapsulate + decapsulate (ML-KEM-768)
    if "ml-kem-768" in pubkeys:
        r = test_verb(ser, "encapsulate", [pubkeys["ml-kem-768"], {"algorithm": "ml-kem-768", "index": 0}])
        ct_kem = None
        ss_kem = None
        if r and "result" in r:
            passed += 1
            ct_kem = r["result"].get("ciphertext", "")
            ss_kem = r["result"].get("shared_secret", "")
        else:
            failed += 1

        if ct_kem:
            r = test_verb(ser, "decapsulate", [ct_kem, {"algorithm": "ml-kem-768", "index": 0}])
            if r and "result" in r:
                ss_result = r["result"].get("shared_secret", "")
                if ss_result == ss_kem:
                    print(f"  ✅ decapsulate shared secret matches encapsulate")
                    passed += 1
                else:
                    print(f"  ❌ shared secret mismatch")
                    failed += 1
            else:
                failed += 1
        else:
            print("\n--- decapsulate (ml-kem-768) ---")
            print("  ⏭️ SKIP: no encapsulate ciphertext")
            failed += 1
    else:
        print("\n--- encapsulate (ml-kem-768) ---")
        print("  ⏭️ SKIP: no ml-kem-768 pubkey")
        failed += 1

    # Summary
    print(f"\n{'='*60}")
    print(f"RESULTS: {passed} passed, {failed} failed, {passed+failed} total")
    print(f"{'='*60}")

    ser.close()
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())