"""
Chain-as-key payload encryption for LockStock.

Encrypts canonical payloads using keys derived from chain state.
No pre-shared secrets, no PKI. The chain's integrity IS the trust anchor.

Anyone with chain access can derive keys and decrypt -- this is by design.
The encryption provides three guarantees, NOT confidentiality from auditors:

1. **Tamper binding**: AES-GCM auth tag fails if ciphertext is modified.
   Key derivation fails if stamp is modified. Chain link fails if
   previous stamp is modified. Tampering at any layer is detectable.

2. **Position binding**: Encryption key is derived from this specific
   stamp's hash, which incorporates chain position. You can't take an
   encrypted payload from position N and claim it belongs at position M.

3. **Ordering proof**: Each stamp's hash depends on the previous stamp's
   hash. The payload is encrypted with a key derived from its stamp.
   Moving, reordering, or inserting payloads breaks the chain, the keys,
   or both.

Write path (at stamp time):
    1. Receive plaintext JSON payload
    2. Canonicalize via bencode -> deterministic bytes
    3. Hash canonical bytes -> payload_hash (SHA-256)
    4. Submit stamp with payload_hash -> receive stamp_hash back
    5. Derive key: HKDF-SHA256(ikm=stamp_hash, salt=payload_hash, info=context)
    6. Encrypt: AES-256-GCM(key, nonce=derived, plaintext=canonical_bytes)
    7. Store ciphertext blob with reference to stamp position

Read path (audit/verification):
    1. Retrieve stamp from chain -> get payload_hash, stamp_hash
    2. Verify chain integrity (standard chain walk)
    3. Derive key from stamp_hash + payload_hash
    4. Decrypt ciphertext -> canonical_bytes
    5. Verify: SHA256(canonical_bytes) == payload_hash
    6. Parse canonical bytes -> original payload content
"""

import base64
import hashlib
from typing import Optional, Tuple

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes


CONTEXT_INFO = b"lockstock-payload-v1"
NONCE_CONTEXT = b"lockstock-nonce-v1"
PACKAGE_VERSION = "lockstock-encrypted-payload-v1"


def derive_key(stamp_hash: str, payload_hash: str) -> bytes:
    """
    Derive a 256-bit AES key from chain state.

    Args:
        stamp_hash: hex string of the stamp's hash on chain (unique per position)
        payload_hash: hex string of SHA-256(canonical_bytes) (binds key to content)

    Returns:
        32 bytes suitable for AES-256-GCM.
    """
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=bytes.fromhex(payload_hash),
        info=CONTEXT_INFO,
    )
    return hkdf.derive(bytes.fromhex(stamp_hash))


def derive_nonce(stamp_hash: str) -> bytes:
    """
    Derive a 96-bit nonce from the stamp hash.

    Since each stamp_hash is unique (chain property), each nonce is unique.
    Never reuses a nonce with the same key -- guaranteed by the chain's
    hash uniqueness invariant.

    Args:
        stamp_hash: hex string of the stamp's hash on chain

    Returns:
        12 bytes (96-bit nonce for AES-GCM).
    """
    return hashlib.sha256(
        NONCE_CONTEXT + bytes.fromhex(stamp_hash)
    ).digest()[:12]


def encrypt_payload(
    canonical_bytes: bytes,
    stamp_hash: str,
    payload_hash: str,
) -> bytes:
    """
    Encrypt canonical bytes using chain-derived key.

    Args:
        canonical_bytes: bencoded canonical form of the JSON payload
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string of SHA-256(canonical_bytes)

    Returns:
        Ciphertext bytes (includes GCM auth tag).
    """
    key = derive_key(stamp_hash, payload_hash)
    nonce = derive_nonce(stamp_hash)
    aesgcm = AESGCM(key)

    # Associated data: the payload_hash itself, binding ciphertext to stamp
    aad = payload_hash.encode("utf-8")

    return aesgcm.encrypt(nonce, canonical_bytes, aad)


def decrypt_payload(
    ciphertext: bytes,
    stamp_hash: str,
    payload_hash: str,
) -> bytes:
    """
    Decrypt and verify a payload using chain-derived key.

    Args:
        ciphertext: encrypted bytes (with GCM auth tag)
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string from the chain stamp

    Returns:
        Canonical bytes (bencoded).

    Raises:
        cryptography.exceptions.InvalidTag: if tampered or wrong chain position
        ValueError: if decrypted content doesn't match payload_hash
    """
    key = derive_key(stamp_hash, payload_hash)
    nonce = derive_nonce(stamp_hash)
    aesgcm = AESGCM(key)

    aad = payload_hash.encode("utf-8")
    canonical_bytes = aesgcm.decrypt(nonce, ciphertext, aad)

    # Belt and suspenders: verify canonical hash matches stamp
    verify_hash = hashlib.sha256(canonical_bytes).hexdigest()
    if verify_hash != payload_hash:
        raise ValueError(
            f"Decrypted payload hash mismatch: "
            f"got {verify_hash}, expected {payload_hash}. "
            f"Chain integrity compromised."
        )

    return canonical_bytes


def encrypt_and_package(
    canonical_bytes: bytes,
    stamp_hash: str,
    payload_hash: str,
) -> dict:
    """
    Encrypt and return a self-describing package for storage.

    Returns a dict suitable for JSON serialization and storage
    alongside the chain. Contains everything needed to locate
    the decryption key (chain reference) and verify integrity.

    Args:
        canonical_bytes: bencoded canonical form of the JSON payload
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string of SHA-256(canonical_bytes)

    Returns:
        Dict with version, stamp_hash, payload_hash, ciphertext_b64,
        ciphertext_length, plaintext_length, algorithm, kdf.
    """
    ciphertext = encrypt_payload(canonical_bytes, stamp_hash, payload_hash)

    return {
        "version": PACKAGE_VERSION,
        "stamp_hash": stamp_hash,
        "payload_hash": payload_hash,
        "ciphertext_b64": base64.b64encode(ciphertext).decode("ascii"),
        "ciphertext_length": len(ciphertext),
        "plaintext_length": len(canonical_bytes),
        "algorithm": "AES-256-GCM",
        "kdf": "HKDF-SHA256",
    }


def decrypt_package(package: dict) -> bytes:
    """
    Decrypt a self-describing encrypted payload package.

    Args:
        package: dict from encrypt_and_package() (or loaded from storage)

    Returns:
        Canonical bytes (bencoded).

    Raises:
        cryptography.exceptions.InvalidTag: if tampered
        ValueError: if hash mismatch or unsupported version
        KeyError: if package is missing required fields
    """
    if package.get("version") != PACKAGE_VERSION:
        raise ValueError(
            f"Unsupported package version: {package.get('version')}"
        )

    ciphertext = base64.b64decode(package["ciphertext_b64"])
    return decrypt_payload(
        ciphertext,
        package["stamp_hash"],
        package["payload_hash"],
    )


# ---------------------------------------------------------------------------
# Envelope encryption (Layer 4) — composing chain binding with confidentiality
# ---------------------------------------------------------------------------
#
# The functions below use an operator-held DEK (Data Encryption Key) instead
# of deriving keys solely from chain state. The blob key still incorporates
# stamp_hash (position) and payload_hash (content), preserving all chain-as-key
# binding properties, but adds a confidentiality gate: without the DEK,
# knowing the chain is insufficient to decrypt.
#
# Key derivation:
#     blob_key = HKDF(ikm=DEK, salt=stamp_hash, info=prefix+payload_hash)
#
# The chain-as-key functions above remain available for backward compatibility.

ENVELOPE_PACKAGE_VERSION = "lockstock-envelope-payload-v1"


def encrypt_payload_envelope(
    canonical_bytes: bytes,
    stamp_hash: str,
    payload_hash: str,
    dek: bytes,
) -> bytes:
    """
    Encrypt canonical bytes using envelope encryption (DEK + chain binding).

    Args:
        canonical_bytes: bencoded canonical form of the JSON payload
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string of SHA-256(canonical_bytes)
        dek: 32-byte Data Encryption Key (from keyring)

    Returns:
        Ciphertext bytes (includes GCM auth tag).
    """
    from .keyring import derive_blob_key, derive_blob_nonce

    key = derive_blob_key(dek, stamp_hash, payload_hash)
    nonce = derive_blob_nonce(dek, stamp_hash)
    aesgcm = AESGCM(key)

    aad = payload_hash.encode("utf-8")
    return aesgcm.encrypt(nonce, canonical_bytes, aad)


def decrypt_payload_envelope(
    ciphertext: bytes,
    stamp_hash: str,
    payload_hash: str,
    dek: bytes,
) -> bytes:
    """
    Decrypt and verify a payload using envelope encryption.

    Args:
        ciphertext: encrypted bytes (with GCM auth tag)
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string from the chain stamp
        dek: 32-byte Data Encryption Key

    Returns:
        Canonical bytes (bencoded).

    Raises:
        cryptography.exceptions.InvalidTag: if tampered, wrong position, or wrong DEK
        ValueError: if decrypted content doesn't match payload_hash
    """
    from .keyring import derive_blob_key, derive_blob_nonce

    key = derive_blob_key(dek, stamp_hash, payload_hash)
    nonce = derive_blob_nonce(dek, stamp_hash)
    aesgcm = AESGCM(key)

    aad = payload_hash.encode("utf-8")
    canonical_bytes = aesgcm.decrypt(nonce, ciphertext, aad)

    verify_hash = hashlib.sha256(canonical_bytes).hexdigest()
    if verify_hash != payload_hash:
        raise ValueError(
            f"Decrypted payload hash mismatch: "
            f"got {verify_hash}, expected {payload_hash}. "
            f"Chain integrity compromised."
        )

    return canonical_bytes


def encrypt_and_package_envelope(
    canonical_bytes: bytes,
    stamp_hash: str,
    payload_hash: str,
    dek: bytes,
    key_version: int = 1,
) -> dict:
    """
    Encrypt with envelope encryption and return a self-describing package.

    The package includes a key_version field to support DEK rotation.
    The version string distinguishes envelope packages from chain-as-key
    packages, enabling decrypt_package_auto to dispatch correctly.

    Args:
        canonical_bytes: bencoded canonical form of the JSON payload
        stamp_hash: hex string of this stamp's hash on chain
        payload_hash: hex string of SHA-256(canonical_bytes)
        dek: 32-byte Data Encryption Key
        key_version: DEK version number (for rotation tracking)

    Returns:
        Dict with version, stamp_hash, payload_hash, ciphertext_b64,
        key_version, ciphertext_length, plaintext_length, algorithm, kdf.
    """
    ciphertext = encrypt_payload_envelope(canonical_bytes, stamp_hash, payload_hash, dek)

    return {
        "version": ENVELOPE_PACKAGE_VERSION,
        "stamp_hash": stamp_hash,
        "payload_hash": payload_hash,
        "ciphertext_b64": base64.b64encode(ciphertext).decode("ascii"),
        "key_version": key_version,
        "ciphertext_length": len(ciphertext),
        "plaintext_length": len(canonical_bytes),
        "algorithm": "AES-256-GCM",
        "kdf": "HKDF-SHA256-envelope",
    }


def decrypt_package_envelope(package: dict, dek: bytes) -> bytes:
    """
    Decrypt an envelope-encrypted payload package.

    Args:
        package: dict from encrypt_and_package_envelope()
        dek: 32-byte Data Encryption Key

    Returns:
        Canonical bytes (bencoded).

    Raises:
        cryptography.exceptions.InvalidTag: if tampered or wrong DEK
        ValueError: if hash mismatch or unsupported version
    """
    if package.get("version") != ENVELOPE_PACKAGE_VERSION:
        raise ValueError(
            f"Unsupported envelope package version: {package.get('version')}"
        )

    ciphertext = base64.b64decode(package["ciphertext_b64"])
    return decrypt_payload_envelope(
        ciphertext,
        package["stamp_hash"],
        package["payload_hash"],
        dek,
    )


def decrypt_package_auto(
    package: dict,
    dek: Optional[bytes] = None,
    strict_envelope: bool = False,
) -> bytes:
    """
    Auto-detect package version and decrypt with the appropriate method.

    Dispatches based on the 'version' field:
    - "lockstock-encrypted-payload-v1"  → chain-as-key (no DEK needed)
    - "lockstock-envelope-payload-v1"   → envelope (DEK required)

    Args:
        package: dict from encrypt_and_package() or encrypt_and_package_envelope()
        dek: 32-byte DEK (required for envelope packages, ignored for chain-as-key)
        strict_envelope: If True, reject chain-as-key packages with ValueError.
            Auditor tools use this to flag blobs encrypted without envelope
            encryption (weaker key derivation, no confidentiality from server).

    Returns:
        Canonical bytes (bencoded).

    Raises:
        ValueError: if version is unknown, or strict_envelope and package is chain-as-key,
                    or envelope package but dek is None
        cryptography.exceptions.InvalidTag: if tampered
    """
    version = package.get("version")

    if version == PACKAGE_VERSION:
        # Chain-as-key package
        if strict_envelope:
            raise ValueError(
                f"Package version '{version}' uses chain-as-key encryption "
                f"(key_version=0). Envelope encryption required when "
                f"strict_envelope=True."
            )
        return decrypt_package(package)

    if version == ENVELOPE_PACKAGE_VERSION:
        if dek is None:
            raise ValueError(
                "Envelope-encrypted package requires a DEK for decryption. "
                "Pass the agent's Data Encryption Key as the dek parameter."
            )
        return decrypt_package_envelope(package, dek)

    raise ValueError(f"Unknown package version: {version}")
