#!/usr/bin/env python3
"""
MeshBay — Spike 1: Crypto Primitives
Validates the full cryptographic stack needed for MeshBay.
"""

import base64
import os
import sys
import time

import blake3
import jwt
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.argon2 import Argon2id
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

PASS = "✓"
FAIL = "✗"
results = []

def check(name, fn):
    try:
        result = fn()
        msg = result if isinstance(result, str) else PASS
        print(f"  {PASS} {name}: {msg}")
        results.append((name, True))
    except Exception as e:
        print(f"  {FAIL} {name}: {e}")
        results.append((name, False))


print("\n=== Test 1: Ed25519 — Hub keypair, sign, verify ===")

sk_hub = Ed25519PrivateKey.generate()
pk_hub = sk_hub.public_key()
sk_user = Ed25519PrivateKey.generate()
pk_user = sk_user.public_key()

def test_ed25519_sign_verify():
    msg = b"meshbay hub token payload"
    sig = sk_hub.sign(msg)
    pk_hub.verify(sig, msg)
    return f"signature {len(sig)} bytes"

def test_ed25519_wrong_sig():
    msg = b"original"
    sig = sk_hub.sign(msg)
    try:
        pk_hub.verify(sig, b"tampered")
        return f"{FAIL} should have raised"
    except Exception:
        return "tampered message correctly rejected"

def test_ed25519_serialization():
    sk_pem = sk_hub.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption()
    )
    pk_pem = pk_hub.public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.SubjectPublicKeyInfo
    )
    pk_raw = pk_hub.public_bytes(
        serialization.Encoding.Raw,
        serialization.PublicFormat.Raw
    )
    return f"PEM sk={len(sk_pem)}B pk={len(pk_pem)}B raw={len(pk_raw)}B"

check("Sign + verify", test_ed25519_sign_verify)
check("Tampered message rejected", test_ed25519_wrong_sig)
check("PEM serialization", test_ed25519_serialization)


print("\n=== Test 2: X25519 — Two-party key agreement (GEK wrapping) ===")

def test_x25519_agreement():
    sk_a = X25519PrivateKey.generate()
    sk_b = X25519PrivateKey.generate()
    shared_a = sk_a.exchange(sk_b.public_key())
    shared_b = sk_b.exchange(sk_a.public_key())
    assert shared_a == shared_b, "Shared secrets don't match"
    return f"shared secret {len(shared_a)} bytes, both sides match"

def test_x25519_different_pairs():
    sk_a = X25519PrivateKey.generate()
    sk_b = X25519PrivateKey.generate()
    sk_c = X25519PrivateKey.generate()
    shared_ab = sk_a.exchange(sk_b.public_key())
    shared_ac = sk_a.exchange(sk_c.public_key())
    assert shared_ab != shared_ac, "Different pairs should produce different secrets"
    return "different key pairs produce different secrets"

check("Two-party agreement", test_x25519_agreement)
check("Different pairs ≠ same secret", test_x25519_different_pairs)


print("\n=== Test 3: ChaCha20-Poly1305 — GEK encryption ===")

gek_raw = ChaCha20Poly1305.generate_key()
cipher_gek = ChaCha20Poly1305(gek_raw)
CHUNK_SIZE = 1024 * 1024  # 1 MB
chunk_data = os.urandom(CHUNK_SIZE)

def test_chacha_roundtrip():
    nonce = os.urandom(12)
    ct = cipher_gek.encrypt(nonce, chunk_data, None)
    pt = cipher_gek.decrypt(nonce, ct, None)
    assert pt == chunk_data, "Decrypted data doesn't match"
    overhead = len(ct) - len(chunk_data)
    return f"1 MB roundtrip OK, AEAD overhead={overhead}B"

def test_chacha_perf():
    times = []
    for _ in range(5):
        nonce = os.urandom(12)
        t0 = time.perf_counter()
        ct = cipher_gek.encrypt(nonce, chunk_data, None)
        cipher_gek.decrypt(nonce, ct, None)
        times.append(time.perf_counter() - t0)
    avg_ms = sum(times) / len(times) * 1000
    throughput = (CHUNK_SIZE * 2) / (sum(times) / len(times)) / (1024**2)
    return f"avg {avg_ms:.1f}ms/chunk (encrypt+decrypt), {throughput:.0f} MB/s"

def test_chacha_tamper():
    nonce = os.urandom(12)
    ct = bytearray(cipher_gek.encrypt(nonce, b"secret", None))
    ct[0] ^= 0xFF  # flip a bit
    try:
        cipher_gek.decrypt(nonce, bytes(ct), None)
        return f"{FAIL} should have raised"
    except Exception:
        return "tampered ciphertext correctly rejected"

check("Encrypt/decrypt roundtrip (1 MB)", test_chacha_roundtrip)
check("Performance (5 runs)", test_chacha_perf)
check("Tampered ciphertext rejected", test_chacha_tamper)


print("\n=== Test 4: HKDF — Per-chunk key derivation ===")

def test_hkdf_chunk_keys():
    file_hash = blake3.blake3(chunk_data).digest()
    keys = []
    for i in range(3):
        key = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b"file:" + file_hash + b":chunk:" + i.to_bytes(4, "big")
        ).derive(gek_raw)
        keys.append(key)
    assert keys[0] != keys[1] != keys[2], "Chunk keys must differ"
    return f"3 distinct chunk keys derived, each {len(keys[0])} bytes"

def test_hkdf_same_input_same_output():
    file_hash = blake3.blake3(chunk_data).digest()
    key1 = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None,
        info=b"file:" + file_hash + b":chunk:" + (0).to_bytes(4, "big")
    ).derive(gek_raw)
    key2 = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None,
        info=b"file:" + file_hash + b":chunk:" + (0).to_bytes(4, "big")
    ).derive(gek_raw)
    assert key1 == key2, "Same input must produce same key"
    return "deterministic: same input → same key"

check("3 distinct chunk keys", test_hkdf_chunk_keys)
check("Deterministic derivation", test_hkdf_same_input_same_output)


print("\n=== Test 5: blake3 — Content hashing ===")

def test_blake3_hash():
    data = os.urandom(1024 * 1024)
    t0 = time.perf_counter()
    h = blake3.blake3(data).digest()
    elapsed = (time.perf_counter() - t0) * 1000
    return f"1 MB hashed in {elapsed:.1f}ms, digest={h.hex()[:16]}..."

def test_blake3_deterministic():
    data = b"test content"
    h1 = blake3.blake3(data).digest()
    h2 = blake3.blake3(data).digest()
    assert h1 == h2
    return "deterministic hashing confirmed"

def test_blake3_different_data():
    h1 = blake3.blake3(b"file chunk 0").digest()
    h2 = blake3.blake3(b"file chunk 1").digest()
    assert h1 != h2
    return "different data → different hashes"

check("1 MB hash performance", test_blake3_hash)
check("Deterministic", test_blake3_deterministic)
check("Collision resistance (basic)", test_blake3_different_data)


print("\n=== Test 6: Argon2id — Keystore key derivation ===")

def test_argon2id_derive():
    salt = os.urandom(16)
    t0 = time.perf_counter()
    kdf = Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536)
    key = kdf.derive(b"my_node_password")
    elapsed = (time.perf_counter() - t0) * 1000
    return f"derived {len(key)}-byte key in {elapsed:.0f}ms"

def test_argon2id_different_salts():
    pw = b"same_password"
    salt1, salt2 = os.urandom(16), os.urandom(16)
    k1 = Argon2id(salt=salt1, length=32, iterations=3, lanes=4, memory_cost=65536).derive(pw)
    k2 = Argon2id(salt=salt2, length=32, iterations=3, lanes=4, memory_cost=65536).derive(pw)
    assert k1 != k2
    return "different salts → different keys (no rainbow table attack)"

def test_argon2id_verify():
    salt = os.urandom(16)
    kdf1 = Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536)
    key = kdf1.derive(b"correct_password")
    kdf2 = Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536)
    kdf2.verify(b"correct_password", key)
    try:
        kdf3 = Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536)
        kdf3.verify(b"wrong_password", key)
        return f"{FAIL} wrong password should be rejected"
    except Exception:
        return "correct password accepted, wrong password rejected"

check("Key derivation (~1s target)", test_argon2id_derive)
check("Salt uniqueness", test_argon2id_different_salts)
check("Verify correct/wrong password", test_argon2id_verify)


print("\n=== Test 7: AES-256-GCM — Keystore encryption ===")

def test_aes_gcm_keystore():
    # Derive an AES key from Argon2id (as done for keystore unlock)
    salt = os.urandom(16)
    aes_key = Argon2id(salt=salt, length=32, iterations=3, lanes=4,
                       memory_cost=65536).derive(b"password")

    # Encrypt a mock keystore blob
    keystore_data = b'{"sk_user": "base64...", "sk_group": "base64..."}'
    iv = os.urandom(16)
    encryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv)).encryptor()
    ct = encryptor.update(keystore_data) + encryptor.finalize()
    tag = encryptor.tag

    # Decrypt
    decryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv, tag)).decryptor()
    pt = decryptor.update(ct) + decryptor.finalize()
    assert pt == keystore_data
    return f"keystore encrypt/decrypt OK ({len(keystore_data)}B → {len(ct)}B + {len(tag)}B tag)"

check("AES-256-GCM keystore roundtrip", test_aes_gcm_keystore)


print("\n=== Test 8: PyJWT EdDSA — Hub JWT issuance and offline verification ===")

sk_hub_pem = sk_hub.private_bytes(
    serialization.Encoding.PEM,
    serialization.PrivateFormat.PKCS8,
    serialization.NoEncryption()
)
pk_hub_pem = pk_hub.public_bytes(
    serialization.Encoding.PEM,
    serialization.PublicFormat.SubjectPublicKeyInfo
)
pk_user_raw_b64 = base64.b64encode(
    pk_user.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
).decode()

def test_jwt_issue_and_verify():
    payload = {
        "iss": "meshbay.org",
        "sub": "user-uuid-1234",
        "pk_user": pk_user_raw_b64,
        "hub_id": "meshbay.org",
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
    }
    token = jwt.encode(payload, sk_hub_pem, algorithm="EdDSA")
    decoded = jwt.decode(token, pk_hub_pem, algorithms=["EdDSA"])
    assert decoded["sub"] == payload["sub"]
    assert decoded["hub_id"] == "meshbay.org"
    assert decoded["pk_user"] == pk_user_raw_b64
    return f"token {len(token)} chars, all claims verified offline"

def test_jwt_tampered_rejected():
    payload = {"sub": "user-1", "exp": int(time.time()) + 3600}
    token = jwt.encode(payload, sk_hub_pem, algorithm="EdDSA")
    # Tamper: flip one char in the signature (last segment)
    parts = token.split(".")
    tampered = parts[0] + "." + parts[1] + "." + parts[2][:-4] + "AAAA"
    try:
        jwt.decode(tampered, pk_hub_pem, algorithms=["EdDSA"])
        return f"{FAIL} tampered token should be rejected"
    except Exception:
        return "tampered JWT correctly rejected"

def test_jwt_wrong_key_rejected():
    sk_other = Ed25519PrivateKey.generate()
    sk_other_pem = sk_other.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption()
    )
    payload = {"sub": "attacker", "exp": int(time.time()) + 3600}
    fake_token = jwt.encode(payload, sk_other_pem, algorithm="EdDSA")
    try:
        jwt.decode(fake_token, pk_hub_pem, algorithms=["EdDSA"])
        return f"{FAIL} wrong key should be rejected"
    except Exception:
        return "token signed with wrong key correctly rejected"

def test_jwt_expired_rejected():
    payload = {"sub": "user-1", "exp": int(time.time()) - 10}  # already expired
    token = jwt.encode(payload, sk_hub_pem, algorithm="EdDSA")
    try:
        jwt.decode(token, pk_hub_pem, algorithms=["EdDSA"])
        return f"{FAIL} expired token should be rejected"
    except jwt.ExpiredSignatureError:
        return "expired JWT correctly rejected"

check("Issue and verify offline (no hub call)", test_jwt_issue_and_verify)
check("Tampered JWT rejected", test_jwt_tampered_rejected)
check("Wrong signing key rejected", test_jwt_wrong_key_rejected)
check("Expired JWT rejected", test_jwt_expired_rejected)


print("\n=== Test 9: Full pipeline — Encrypt chunk + sign + verify + decrypt ===")

def test_full_pipeline():
    # Simulate: node encrypts a chunk and signs it; client verifies and decrypts

    # Node side
    file_data = os.urandom(CHUNK_SIZE)
    file_hash = blake3.blake3(file_data).digest()
    chunk_index = 0

    # Derive per-chunk key
    chunk_key_raw = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None,
        info=b"file:" + file_hash + b":chunk:" + chunk_index.to_bytes(4, "big")
    ).derive(gek_raw)
    chunk_cipher = ChaCha20Poly1305(chunk_key_raw)

    nonce = os.urandom(12)
    t0 = time.perf_counter()
    ciphertext = chunk_cipher.encrypt(nonce, file_data, None)
    ct_hash = blake3.blake3(ciphertext).digest()

    # Sign: chunk_index || nonce || ciphertext_hash
    sig_payload = chunk_index.to_bytes(4, "big") + nonce + ct_hash
    sig = sk_user.sign(sig_payload)
    encrypt_ms = (time.perf_counter() - t0) * 1000

    # Client side
    t1 = time.perf_counter()
    # 1. Verify signature
    pk_user.verify(sig, sig_payload)
    # 2. Verify ciphertext hash
    assert blake3.blake3(ciphertext).digest() == ct_hash
    # 3. Decrypt
    plaintext = chunk_cipher.decrypt(nonce, ciphertext, None)
    assert plaintext == file_data
    verify_ms = (time.perf_counter() - t1) * 1000

    return f"1 MB: encrypt+sign={encrypt_ms:.1f}ms, verify+decrypt={verify_ms:.1f}ms"

check("Full encrypt→sign→verify→decrypt pipeline (1 MB)", test_full_pipeline)


print("\n" + "="*55)
passed = sum(1 for _, ok in results if ok)
total = len(results)
print(f"Results: {passed}/{total} passed")
if passed == total:
    print("All crypto primitives validated. Spike 1 COMPLETE.")
else:
    print("Some tests failed — review above.")
    failed = [name for name, ok in results if not ok]
    for name in failed:
        print(f"  {FAIL} {name}")
    sys.exit(1)