import secrets

import pytest

from quantumpipe.chab3 import (
    NONCE_LEN,
    Receiver,
    Sender,
    decrypt,
    decrypt_into,
    encrypt,
    encrypt_into,
)

KEY = secrets.token_bytes(32)
NONCE = secrets.token_bytes(12)
AD = secrets.token_bytes(7)
MSG = b"streaming payload, split " + b"into two chunks"


def test_roundtrip():
    ct = encrypt(MSG, KEY, NONCE, ad=AD)
    assert isinstance(ct, bytearray)
    assert len(ct) == len(MSG) + 16
    pt = decrypt(ct, KEY, NONCE, ad=AD)
    assert isinstance(pt, bytearray)
    assert pt == MSG


def test_into_variants():
    ref = encrypt(MSG, KEY, NONCE, ad=AD)
    buf = bytearray(len(MSG) + 16 + 10)  # oversized buffer
    used = encrypt_into(buf, MSG, KEY, NONCE, ad=AD)
    assert bytes(used) == ref
    assert buf[len(used) :] == bytes(10)  # tail untouched
    out = bytearray(len(MSG) + 5)
    pt = decrypt_into(out, used, KEY, NONCE, ad=AD)
    assert len(pt) == len(MSG)
    assert pt == MSG


def test_inplace():
    buf = bytearray(MSG + bytes(16))
    ct = encrypt_into(buf, buf[: len(MSG)], KEY, NONCE, ad=AD)
    assert bytes(ct) == encrypt(MSG, KEY, NONCE, ad=AD)
    pt = decrypt_into(buf, buf, KEY, NONCE, ad=AD)
    assert pt == MSG


def test_maclen_zero():
    ct = encrypt(MSG, KEY, NONCE, ad=AD, maclen=0)
    assert len(ct) == len(MSG)  # no tag appended
    assert decrypt(ct, KEY, NONCE, maclen=0) == MSG
    box = encrypt(MSG, KEY, maclen=0)
    assert len(box) == NONCE_LEN + len(MSG)
    assert decrypt(box, KEY, maclen=0) == MSG


@pytest.mark.parametrize(
    "fn",
    [
        lambda: encrypt_into(bytearray(1), MSG, KEY, NONCE, ad=AD),
        lambda: decrypt_into(
            bytearray(1), encrypt(MSG, KEY, NONCE, ad=AD), KEY, NONCE, ad=AD
        ),
    ],
)
def test_too_small_output_buffer_raises_type_error(fn):
    # a program bug, not bad wire data
    with pytest.raises(TypeError):
        fn()


@pytest.mark.parametrize(
    "fn",
    [
        lambda: decrypt(b"short", KEY, NONCE, ad=AD),
        lambda: decrypt(b"short", KEY, ad=AD),
    ],
)
def test_short_ciphertext_raises_value_error(fn):
    # bad wire data, not a program bug
    with pytest.raises(ValueError):
        fn()


def test_tampering_raises_value_error():
    good = bytes(encrypt(MSG, KEY, NONCE, ad=AD))
    body, tag = good[:-16], good[-16:]
    cases = [
        (bytes([body[0] ^ 1]) + body[1:] + tag, AD),  # tampered ciphertext
        (good, AD + b"x"),  # wrong ad
        (body + bytes([tag[0] ^ 1]) + bytes(tag[1:]), AD),  # wrong tag
    ]
    for ct, ad in cases:
        with pytest.raises(ValueError):
            decrypt(ct, KEY, NONCE, ad=ad)


def test_boxes():
    box1, box2 = encrypt(MSG, KEY, ad=AD), encrypt(MSG, KEY, ad=AD)
    assert box1 != box2  # fresh random nonce each time
    assert len(box1) == NONCE_LEN + len(MSG) + 16
    assert decrypt(box1, KEY, ad=AD) == MSG
    buf = bytearray(len(box1))
    encrypt_into(buf, MSG, KEY, ad=AD)
    out = bytearray(len(MSG))
    decrypt_into(out, buf, KEY, ad=AD)
    assert bytes(out) == MSG
    tampered = bytearray(box1)
    tampered[-1] ^= 1
    with pytest.raises(ValueError):
        decrypt(bytes(tampered), KEY, ad=AD)


def test_decrypt_into_verifies_before_writing():
    ct = bytearray(encrypt(MSG, KEY, NONCE, ad=AD))
    ct[-1] ^= 1  # bad tag
    out = bytearray(b"\xaa" * len(MSG))
    with pytest.raises(ValueError):
        decrypt_into(out, ct, KEY, NONCE, ad=AD)
    assert out == bytearray(b"\xaa" * len(MSG))  # untouched


def test_channel_in_order():
    tx, rx = Sender(KEY), Receiver(KEY)
    for i in range(3):
        ct = tx.encrypt(f"msg {i}".encode(), ad=AD)
        assert rx.decrypt(ct, ad=AD) == f"msg {i}".encode()


def test_channel_into():
    tx, rx = Sender(KEY), Receiver(KEY)
    buf = bytearray(len(MSG) + 16)
    ct = tx.encrypt_into(buf, MSG, ad=AD)
    out = bytearray(len(MSG))
    assert rx.decrypt_into(out, ct, ad=AD) == MSG


def test_channel_desync_rejected():
    tx = Sender(KEY)
    for _ in range(3):
        tx.encrypt(b"x")
    with pytest.raises(ValueError):
        Receiver(KEY).decrypt(tx.encrypt(b"out of order"))  # nonce 0 vs 3


def test_random_reseeds_on_counter_exhaustion():
    from cryptography.hazmat.primitives.ciphers import Cipher
    from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20

    import quantumpipe.chab3 as m

    exhausted = Cipher(
        ChaCha20(bytes(32), (0xFFFFFFFF).to_bytes(4, "little") + bytes(12)),
        mode=None,
    ).encryptor()
    m._random.enc = exhausted
    out = m.random(128)  # needs 2 blocks; the planted stream has 1 left
    assert len(out) == 128
    assert m._random.enc is not exhausted
    buf = bytearray(128)
    m._random.enc = exhausted
    m.random(buf)
    assert m._random.enc is not exhausted
    assert buf != bytes(128)
    del m._random.enc  # don't leak the fixture into other tests