"""Tests for AES-256-GCM webcrypto variant."""

import os

import blake3
import pytest
from meshbay_common.crypto import generate_gek
from meshbay_common.webcrypto import chunk_key_aes, decrypt_chunk_aes, encrypt_chunk_aes


def test_aes_roundtrip():
    gek  = generate_gek()
    data = os.urandom(1024 * 1024)   # 1 MB
    fh   = blake3.blake3(data).digest()
    key  = chunk_key_aes(gek, fh, 0)
    nonce, ct = encrypt_chunk_aes(key, data)
    assert decrypt_chunk_aes(key, nonce, ct) == data


def test_aes_key_distinct_from_chacha_key():
    """AES and ChaCha20 keys for the same chunk must differ."""
    from meshbay_common.crypto import chunk_key as chacha_key
    gek  = generate_gek()
    data = os.urandom(100)
    fh   = blake3.blake3(data).digest()
    aes_k    = chunk_key_aes(gek, fh, 0)
    chacha_k = chacha_key(gek, fh, 0)
    assert aes_k != chacha_k


def test_aes_wrong_key_rejected():
    gek  = generate_gek()
    data = b"private content"
    fh   = blake3.blake3(data).digest()
    key  = chunk_key_aes(gek, fh, 0)
    nonce, ct = encrypt_chunk_aes(key, data)
    wrong_key = chunk_key_aes(generate_gek(), fh, 0)
    with pytest.raises(Exception):
        decrypt_chunk_aes(wrong_key, nonce, ct)


def test_aes_chunk_keys_unique_per_chunk():
    gek  = generate_gek()
    data = os.urandom(32)
    fh   = blake3.blake3(data).digest()
    keys = {chunk_key_aes(gek, fh, i) for i in range(5)}
    assert len(keys) == 5   # all distinct


def test_aes_gek_wrap_unwrap_roundtrip():
    from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
    from meshbay_common.crypto import (
        pk_to_raw,
        sk_to_raw,
        unwrap_gek_aes,
        wrap_gek_aes,
    )
    gek = generate_gek()
    sk = X25519PrivateKey.generate()
    pk_raw = pk_to_raw(sk.public_key())
    sk_raw = sk_to_raw(sk)

    bundle = wrap_gek_aes(gek, pk_raw)
    recovered = unwrap_gek_aes(bundle, sk_raw, pk_raw)
    assert recovered == gek


def test_aes_gek_wrap_wrong_key_rejected():
    from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
    from meshbay_common.crypto import pk_to_raw, sk_to_raw, unwrap_gek_aes, wrap_gek_aes

    gek = generate_gek()
    sk_a = X25519PrivateKey.generate()
    sk_b = X25519PrivateKey.generate()

    bundle = wrap_gek_aes(gek, pk_to_raw(sk_a.public_key()))
    with pytest.raises(Exception):
        unwrap_gek_aes(bundle, sk_to_raw(sk_b), pk_to_raw(sk_b.public_key()))


def test_aes_gek_wrap_differs_from_chacha_wrap():
    from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
    from meshbay_common.crypto import (
        pk_to_raw,
        wrap_gek,
        wrap_gek_aes,
    )
    gek = generate_gek()
    sk = X25519PrivateKey.generate()
    pk_raw = pk_to_raw(sk.public_key())

    bundle_aes = wrap_gek_aes(gek, pk_raw)
    bundle_chacha = wrap_gek(gek, pk_raw)
    assert bundle_aes["wrapped_b64"] != bundle_chacha["wrapped_b64"]