"""
Tests for lockstock_core.encrypt — Chain-as-key payload encryption.

Tests verify:
1. Round-trip encrypt/decrypt
2. Tamper binding (modified ciphertext fails)
3. Position binding (wrong stamp_hash fails)
4. Ordering proof (swapped payloads fail)
5. Package serialization round-trip
6. Real chain data integration (using actual stamp from Trinity)
7. Key/nonce derivation determinism
"""

import hashlib
import json
import pytest

from lockstock_core.canonicalize import bencode, canonical_hash_from_bytes
from lockstock_core.encrypt import (
    derive_key,
    derive_nonce,
    encrypt_payload,
    decrypt_payload,
    encrypt_and_package,
    decrypt_package,
    encrypt_payload_envelope,
    decrypt_payload_envelope,
    encrypt_and_package_envelope,
    decrypt_package_envelope,
    decrypt_package_auto,
    PACKAGE_VERSION,
    ENVELOPE_PACKAGE_VERSION,
)


# Fixtures: simulated chain state

@pytest.fixture
def chat_request():
    """A chat completion request body."""
    return b'{"model":"trinity","messages":[{"role":"user","content":"Hello, what is 2+2?"}]}'


@pytest.fixture
def chat_request_canonical(chat_request):
    """Bencoded canonical form of the chat request."""
    parsed = json.loads(chat_request)
    return bencode(parsed)


@pytest.fixture
def chat_request_hash(chat_request):
    """SHA-256 of bencoded canonical form."""
    return canonical_hash_from_bytes(chat_request)


@pytest.fixture
def stamp_hash_seq21():
    """Stamp hash from chain position 21 (from real Trinity test)."""
    return "52d9e6bb163ecd57d2ec3feb440901dda37d8a564f9f3881988d869927c68e26"


@pytest.fixture
def stamp_hash_seq22():
    """Stamp hash from chain position 22 (response)."""
    return "b3fc9283fd83539c35effb972eeb9d5990cc0f758480dd8dfa6e5e8385725ef2"


# Tests

class TestKeyDerivation:
    """Test HKDF key and nonce derivation."""

    def test_deterministic_key(self, stamp_hash_seq21, chat_request_hash):
        """Same inputs always produce same key."""
        key1 = derive_key(stamp_hash_seq21, chat_request_hash)
        key2 = derive_key(stamp_hash_seq21, chat_request_hash)
        assert key1 == key2
        assert len(key1) == 32  # AES-256

    def test_different_stamp_different_key(self, stamp_hash_seq21, stamp_hash_seq22, chat_request_hash):
        """Different chain positions produce different keys."""
        key1 = derive_key(stamp_hash_seq21, chat_request_hash)
        key2 = derive_key(stamp_hash_seq22, chat_request_hash)
        assert key1 != key2

    def test_different_payload_different_key(self, stamp_hash_seq21):
        """Different payload hashes produce different keys."""
        hash_a = "a" * 64
        hash_b = "b" * 64
        key1 = derive_key(stamp_hash_seq21, hash_a)
        key2 = derive_key(stamp_hash_seq21, hash_b)
        assert key1 != key2

    def test_deterministic_nonce(self, stamp_hash_seq21):
        """Same stamp always produces same nonce."""
        nonce1 = derive_nonce(stamp_hash_seq21)
        nonce2 = derive_nonce(stamp_hash_seq21)
        assert nonce1 == nonce2
        assert len(nonce1) == 12  # 96-bit

    def test_different_stamp_different_nonce(self, stamp_hash_seq21, stamp_hash_seq22):
        """Different stamps produce different nonces."""
        nonce1 = derive_nonce(stamp_hash_seq21)
        nonce2 = derive_nonce(stamp_hash_seq22)
        assert nonce1 != nonce2


class TestRoundTrip:
    """Test encrypt/decrypt round-trip."""

    def test_basic_roundtrip(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Encrypt then decrypt returns original bytes."""
        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        plaintext = decrypt_payload(ciphertext, stamp_hash_seq21, chat_request_hash)
        assert plaintext == chat_request_canonical

    def test_ciphertext_differs_from_plaintext(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Ciphertext is not the same as plaintext."""
        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        assert ciphertext != chat_request_canonical

    def test_ciphertext_longer_than_plaintext(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Ciphertext includes GCM auth tag (16 bytes)."""
        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        assert len(ciphertext) == len(chat_request_canonical) + 16

    def test_empty_payload(self, stamp_hash_seq21):
        """Empty canonical bytes encrypt and decrypt correctly."""
        payload_hash = hashlib.sha256(b"").hexdigest()
        ciphertext = encrypt_payload(b"", stamp_hash_seq21, payload_hash)
        plaintext = decrypt_payload(ciphertext, stamp_hash_seq21, payload_hash)
        assert plaintext == b""

    def test_large_payload(self, stamp_hash_seq21):
        """650KB payload (realistic chat response) round-trips."""
        large_content = b"x" * 650_000
        payload_hash = hashlib.sha256(large_content).hexdigest()
        ciphertext = encrypt_payload(large_content, stamp_hash_seq21, payload_hash)
        plaintext = decrypt_payload(ciphertext, stamp_hash_seq21, payload_hash)
        assert plaintext == large_content


class TestTamperBinding:
    """Test tamper detection at every layer."""

    def test_modified_ciphertext_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Flipping a byte in ciphertext raises InvalidTag."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        tampered = bytearray(ciphertext)
        tampered[0] ^= 0xFF  # Flip first byte
        with pytest.raises(InvalidTag):
            decrypt_payload(bytes(tampered), stamp_hash_seq21, chat_request_hash)

    def test_truncated_ciphertext_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Truncated ciphertext raises InvalidTag."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        with pytest.raises(InvalidTag):
            decrypt_payload(ciphertext[:-1], stamp_hash_seq21, chat_request_hash)

    def test_wrong_aad_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Wrong payload_hash (AAD) fails decryption."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        wrong_hash = "a" * 64
        with pytest.raises(InvalidTag):
            decrypt_payload(ciphertext, stamp_hash_seq21, wrong_hash)


class TestPositionBinding:
    """Test that payloads are bound to specific chain positions."""

    def test_wrong_stamp_hash_fails(self, chat_request_canonical, stamp_hash_seq21, stamp_hash_seq22, chat_request_hash):
        """Payload encrypted at position 21 cannot be decrypted at position 22."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        with pytest.raises(InvalidTag):
            decrypt_payload(ciphertext, stamp_hash_seq22, chat_request_hash)

    def test_fabricated_stamp_hash_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Fabricated stamp hash fails."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        fake_stamp = "deadbeef" * 8
        with pytest.raises(InvalidTag):
            decrypt_payload(ciphertext, fake_stamp, chat_request_hash)


class TestOrderingProof:
    """Test that swapping payloads between positions fails."""

    def test_swapped_payloads_fail(self, stamp_hash_seq21, stamp_hash_seq22):
        """Payload A at position 21 and payload B at position 22 cannot be swapped."""
        from cryptography.exceptions import InvalidTag

        payload_a = b'{"role":"user","content":"Hello"}'
        payload_b = b'{"role":"assistant","content":"Hi there"}'

        hash_a = hashlib.sha256(bencode(json.loads(payload_a))).hexdigest()
        hash_b = hashlib.sha256(bencode(json.loads(payload_b))).hexdigest()

        canonical_a = bencode(json.loads(payload_a))
        canonical_b = bencode(json.loads(payload_b))

        ct_a = encrypt_payload(canonical_a, stamp_hash_seq21, hash_a)
        ct_b = encrypt_payload(canonical_b, stamp_hash_seq22, hash_b)

        # Correct positions work
        assert decrypt_payload(ct_a, stamp_hash_seq21, hash_a) == canonical_a
        assert decrypt_payload(ct_b, stamp_hash_seq22, hash_b) == canonical_b

        # Swapped positions fail
        with pytest.raises(InvalidTag):
            decrypt_payload(ct_a, stamp_hash_seq22, hash_a)  # ct_a at position 22
        with pytest.raises(InvalidTag):
            decrypt_payload(ct_b, stamp_hash_seq21, hash_b)  # ct_b at position 21


class TestPackaging:
    """Test encrypt_and_package / decrypt_package round-trip."""

    def test_package_roundtrip(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Package encrypts, serializes, and decrypts correctly."""
        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)

        assert package["version"] == "lockstock-encrypted-payload-v1"
        assert package["stamp_hash"] == stamp_hash_seq21
        assert package["payload_hash"] == chat_request_hash
        assert package["algorithm"] == "AES-256-GCM"
        assert package["kdf"] == "HKDF-SHA256"
        assert package["plaintext_length"] == len(chat_request_canonical)
        assert package["ciphertext_length"] == len(chat_request_canonical) + 16

        plaintext = decrypt_package(package)
        assert plaintext == chat_request_canonical

    def test_package_json_serializable(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Package can be JSON serialized and deserialized."""
        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        serialized = json.dumps(package)
        deserialized = json.loads(serialized)
        plaintext = decrypt_package(deserialized)
        assert plaintext == chat_request_canonical

    def test_wrong_version_rejected(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Wrong package version raises ValueError."""
        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        package["version"] = "wrong-version"
        with pytest.raises(ValueError, match="Unsupported package version"):
            decrypt_package(package)

    def test_tampered_package_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Modifying ciphertext_b64 in package fails."""
        import base64
        from cryptography.exceptions import InvalidTag

        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)

        # Tamper with ciphertext
        ct = base64.b64decode(package["ciphertext_b64"])
        tampered = bytearray(ct)
        tampered[0] ^= 0xFF
        package["ciphertext_b64"] = base64.b64encode(bytes(tampered)).decode("ascii")

        with pytest.raises(InvalidTag):
            decrypt_package(package)


class TestRealChainIntegration:
    """Test with actual chain data from the Trinity deployment."""

    def test_real_request_hash_matches(self):
        """The payload_hash from seq 21 matches our independent computation."""
        request_body = b'{"model":"trinity","messages":[{"role":"user","content":"Hello, what is 2+2?"}]}'
        computed = canonical_hash_from_bytes(request_body)
        # This is the payload_hash recorded in audit log seq 21
        assert computed == "fade8d4214d42417d4309c878cda679e7763336230f44f25a03322f1fa959c06"

    def test_encrypt_real_request(self):
        """Encrypt the actual Trinity request with the real stamp hash."""
        request_body = b'{"model":"trinity","messages":[{"role":"user","content":"Hello, what is 2+2?"}]}'
        canonical = bencode(json.loads(request_body))
        payload_hash = canonical_hash_from_bytes(request_body)
        stamp_hash = "52d9e6bb163ecd57d2ec3feb440901dda37d8a564f9f3881988d869927c68e26"

        package = encrypt_and_package(canonical, stamp_hash, payload_hash)
        recovered = decrypt_package(package)

        # Recovered canonical bytes decode back to original content
        assert recovered == canonical

        # Parse the bencoded content to verify it's the right data
        # (bencode dict keys are sorted, so we can verify structure)
        assert b"trinity" in recovered
        assert b"Hello, what is 2+2?" in recovered


# --- Envelope encryption tests (Layer 4) ---

# Fixed DEK for envelope tests (not a real key — test only)
ENVELOPE_DEK = bytes.fromhex(
    "1122334455667788990011223344556677889900112233445566778899001122"
)
ALT_DEK = bytes.fromhex(
    "ffeeddccbbaa99887766554433221100ffeeddccbbaa99887766554433221100"
)


class TestEnvelopeRoundTrip:
    """Test envelope encrypt/decrypt round-trip."""

    def test_basic_roundtrip(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Envelope encrypt then decrypt returns original bytes."""
        ciphertext = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        plaintext = decrypt_payload_envelope(
            ciphertext, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        assert plaintext == chat_request_canonical

    def test_ciphertext_differs_from_chain_as_key(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Envelope ciphertext differs from chain-as-key ciphertext."""
        chain_ct = encrypt_payload(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        envelope_ct = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        assert chain_ct != envelope_ct

    def test_empty_payload(self, stamp_hash_seq21):
        """Empty canonical bytes encrypt and decrypt correctly."""
        payload_hash = hashlib.sha256(b"").hexdigest()
        ciphertext = encrypt_payload_envelope(b"", stamp_hash_seq21, payload_hash, ENVELOPE_DEK)
        plaintext = decrypt_payload_envelope(ciphertext, stamp_hash_seq21, payload_hash, ENVELOPE_DEK)
        assert plaintext == b""

    def test_large_payload(self, stamp_hash_seq21):
        """650KB payload round-trips with envelope encryption."""
        large_content = b"x" * 650_000
        payload_hash = hashlib.sha256(large_content).hexdigest()
        ciphertext = encrypt_payload_envelope(large_content, stamp_hash_seq21, payload_hash, ENVELOPE_DEK)
        plaintext = decrypt_payload_envelope(ciphertext, stamp_hash_seq21, payload_hash, ENVELOPE_DEK)
        assert plaintext == large_content


class TestEnvelopeTamperBinding:
    """Test envelope tamper detection."""

    def test_wrong_dek_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Wrong DEK fails decryption (confidentiality gate)."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        with pytest.raises(InvalidTag):
            decrypt_payload_envelope(ciphertext, stamp_hash_seq21, chat_request_hash, ALT_DEK)

    def test_wrong_stamp_hash_fails(self, chat_request_canonical, stamp_hash_seq21, stamp_hash_seq22, chat_request_hash):
        """Wrong stamp_hash fails (position binding preserved)."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        with pytest.raises(InvalidTag):
            decrypt_payload_envelope(ciphertext, stamp_hash_seq22, chat_request_hash, ENVELOPE_DEK)

    def test_wrong_payload_hash_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Wrong payload_hash fails (content binding preserved)."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        wrong_hash = "a" * 64
        with pytest.raises(InvalidTag):
            decrypt_payload_envelope(ciphertext, stamp_hash_seq21, wrong_hash, ENVELOPE_DEK)

    def test_modified_ciphertext_fails(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Flipping a byte in ciphertext raises InvalidTag."""
        from cryptography.exceptions import InvalidTag

        ciphertext = encrypt_payload_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        tampered = bytearray(ciphertext)
        tampered[0] ^= 0xFF
        with pytest.raises(InvalidTag):
            decrypt_payload_envelope(bytes(tampered), stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK)


class TestEnvelopePackaging:
    """Test envelope package format and round-trip."""

    def test_package_roundtrip(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Envelope package encrypts, serializes, and decrypts correctly."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )

        assert package["version"] == ENVELOPE_PACKAGE_VERSION
        assert package["stamp_hash"] == stamp_hash_seq21
        assert package["payload_hash"] == chat_request_hash
        assert package["algorithm"] == "AES-256-GCM"
        assert package["kdf"] == "HKDF-SHA256-envelope"
        assert package["key_version"] == 1
        assert package["plaintext_length"] == len(chat_request_canonical)
        assert package["ciphertext_length"] == len(chat_request_canonical) + 16

        plaintext = decrypt_package_envelope(package, ENVELOPE_DEK)
        assert plaintext == chat_request_canonical

    def test_custom_key_version(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """key_version parameter flows through to package."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash,
            ENVELOPE_DEK, key_version=3
        )
        assert package["key_version"] == 3

    def test_json_serializable(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Envelope package can be JSON serialized and deserialized."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        serialized = json.dumps(package)
        deserialized = json.loads(serialized)
        plaintext = decrypt_package_envelope(deserialized, ENVELOPE_DEK)
        assert plaintext == chat_request_canonical

    def test_wrong_version_rejected(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Wrong version string raises ValueError."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        package["version"] = "wrong-version"
        with pytest.raises(ValueError, match="Unsupported envelope package version"):
            decrypt_package_envelope(package, ENVELOPE_DEK)


class TestDecryptPackageAuto:
    """Test auto-dispatch between chain-as-key and envelope packages."""

    def test_chain_as_key_dispatch(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Auto-detect chain-as-key package and decrypt without DEK."""
        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        plaintext = decrypt_package_auto(package)
        assert plaintext == chat_request_canonical

    def test_envelope_dispatch(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Auto-detect envelope package and decrypt with DEK."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        plaintext = decrypt_package_auto(package, dek=ENVELOPE_DEK)
        assert plaintext == chat_request_canonical

    def test_envelope_without_dek_raises(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """Envelope package without DEK raises ValueError."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        with pytest.raises(ValueError, match="requires a DEK"):
            decrypt_package_auto(package, dek=None)

    def test_strict_envelope_rejects_chain_as_key(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """strict_envelope=True rejects chain-as-key packages."""
        package = encrypt_and_package(chat_request_canonical, stamp_hash_seq21, chat_request_hash)
        with pytest.raises(ValueError, match="chain-as-key encryption"):
            decrypt_package_auto(package, strict_envelope=True)

    def test_strict_envelope_accepts_envelope(self, chat_request_canonical, stamp_hash_seq21, chat_request_hash):
        """strict_envelope=True passes envelope packages through."""
        package = encrypt_and_package_envelope(
            chat_request_canonical, stamp_hash_seq21, chat_request_hash, ENVELOPE_DEK
        )
        plaintext = decrypt_package_auto(package, dek=ENVELOPE_DEK, strict_envelope=True)
        assert plaintext == chat_request_canonical

    def test_unknown_version_raises(self):
        """Unknown version string raises ValueError."""
        package = {"version": "unknown-v99"}
        with pytest.raises(ValueError, match="Unknown package version"):
            decrypt_package_auto(package)
