"""
MeshBay Spike 5 — Encrypted Transfer CLIENT (runs on meshbay.org)

Acts as the "file requester":
  - Opens TCP server on :19003, waits for the node to connect
  - Sends a chunk request
  - Receives the encrypted chunk
  - Verifies Ed25519 signature
  - Verifies blake3 ciphertext hash
  - Decrypts with GEK (included in POC response — never in production)
  - Verifies plaintext integrity via blake3
  - Reports timing and results
"""

import asyncio
import base64
import json
import struct
import sys
import time

import blake3
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

PORT        = 19003
CHUNK_INDEX = 0
PASS = "✓"
FAIL = "✗"

# ── Wire helpers ───────────────────────────────────────────────────────────────

async def send_msg(writer, obj: dict):
    data = json.dumps(obj).encode()
    writer.write(struct.pack('>I', len(data)) + data)
    await writer.drain()

async def recv_msg(reader) -> dict:
    raw_len = await reader.readexactly(4)
    length  = struct.unpack('>I', raw_len)[0]
    data    = await reader.readexactly(length)
    return json.loads(data)

# ── Handler ────────────────────────────────────────────────────────────────────

async def handle_node(reader, writer):
    peer = writer.get_extra_info('peername')
    print(f"  Node connected from {peer[0]}:{peer[1]}")

    # 1. Send chunk request
    await send_msg(writer, {
        "type":        "chunk_req",
        "file_id":     "testfile.bin",
        "chunk_index": CHUNK_INDEX,
    })
    print(f"  Chunk request sent (file=testfile.bin, chunk={CHUNK_INDEX})")

    # 2. Receive encrypted chunk
    t_recv_start = time.perf_counter()
    chunk = await recv_msg(reader)
    recv_ms = (time.perf_counter() - t_recv_start) * 1000
    writer.close()

    ct       = base64.b64decode(chunk["ct_b64"])
    nonce    = base64.b64decode(chunk["nonce_b64"])
    ct_hash  = base64.b64decode(chunk["ct_hash_b64"])
    sig      = base64.b64decode(chunk["sig_b64"])
    pk_bytes = base64.b64decode(chunk["pk_node_b64"])
    gek_raw  = base64.b64decode(chunk["gek_b64"])         # POC only
    file_hash= base64.b64decode(chunk["file_hash_b64"])
    ci       = chunk["chunk_index"]
    chunk_sz = chunk["plaintext_size"]

    print(f"\n  Received {len(ct)/1024:.1f} KB ciphertext in {recv_ms:.0f}ms")
    print(f"  chunk_index={ci}  plaintext_size={chunk_sz/1024:.1f} KB")

    t_verify = time.perf_counter()

    # 3. Verify Ed25519 signature: sig covers (chunk_index || nonce || ct_hash)
    pk_node = Ed25519PublicKey.from_public_bytes(pk_bytes)
    sig_payload = ci.to_bytes(4, 'big') + nonce + ct_hash
    try:
        pk_node.verify(sig, sig_payload)
        print(f"  {PASS} Ed25519 signature valid")
    except Exception as e:
        print(f"  {FAIL} Signature INVALID: {e}")
        sys.exit(1)

    # 4. Verify ciphertext hash
    computed_hash = blake3.blake3(ct).digest()
    if computed_hash == ct_hash:
        print(f"  {PASS} blake3(ciphertext) matches")
    else:
        print(f"  {FAIL} Ciphertext hash MISMATCH")
        sys.exit(1)

    # 5. Derive per-chunk key and decrypt
    chunk_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None,
        info=b"file:" + file_hash + b":chunk:" + ci.to_bytes(4, 'big')
    ).derive(gek_raw)

    try:
        plaintext = ChaCha20Poly1305(chunk_key).decrypt(nonce, ct, None)
        print(f"  {PASS} Decrypted {len(plaintext)/1024:.1f} KB")
    except Exception as e:
        print(f"  {FAIL} Decryption FAILED: {e}")
        sys.exit(1)

    # 6. Verify plaintext hash
    pt_hash = blake3.blake3(plaintext).digest()
    expected_pt_hash = base64.b64decode(chunk["pt_hash_b64"])
    if pt_hash == expected_pt_hash:
        print(f"  {PASS} blake3(plaintext) matches — data integrity confirmed")
    else:
        print(f"  {FAIL} Plaintext hash MISMATCH — corruption or wrong key")
        sys.exit(1)

    verify_ms = (time.perf_counter() - t_verify) * 1000
    total_ms  = recv_ms + verify_ms

    print(f"\n  {'='*50}")
    print("  Timings:")
    print(f"    Network receive : {recv_ms:.0f} ms")
    print(f"    Verify+decrypt  : {verify_ms:.1f} ms")
    print(f"    Total           : {total_ms:.0f} ms")
    throughput = (chunk_sz / 1024 / 1024) / (recv_ms / 1000)
    print(f"    Throughput      : {throughput:.1f} MB/s")
    print(f"\n  {PASS} Spike 5 CLIENT — all checks passed.")
    sys.exit(0)


async def main():
    print("\n=== MeshBay Spike 5 — Encrypted Transfer CLIENT ===")
    print(f"Waiting for node on TCP :{PORT} ...")

    server = await asyncio.start_server(handle_node, '0.0.0.0', PORT)
    async with server:
        try:
            await asyncio.wait_for(server.serve_forever(), timeout=30)
        except TimeoutError:
            print(f"  {FAIL} Timeout — node did not connect within 30s")
            sys.exit(1)

asyncio.run(main())