"""Cryptographic utilities: Ed25519, X25519, AES-256-GCM, Double Ratchet, Sender Keys.

RSA functions retained for login challenge-response only.
"""

import hashlib
import hmac
import json
import os
import struct
import uuid
from dataclasses import dataclass, field

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


# ---------------------------------------------------------------------------
# Password-based key encryption (M3: PBKDF2 600k iterations + AES-256-GCM)
# ---------------------------------------------------------------------------

PBKDF2_ITERATIONS = 600_000
_ECP1_MAGIC = b"ECP1"  # Encrypted Chat PBKDF v1 format marker


def _encrypt_private_key(raw_bytes: bytes, password: bytes) -> bytes:
    """Encrypt raw key bytes with PBKDF2-HMAC-SHA256 (600k iterations) + AES-256-GCM.

    Output format: MAGIC(4) + salt(16) + nonce(12) + ciphertext_with_tag(N+16)
    """
    salt = os.urandom(16)
    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32,
                      salt=salt, iterations=PBKDF2_ITERATIONS)
    derived = kdf.derive(password)
    nonce = os.urandom(12)
    aesgcm = AESGCM(derived)
    ct = aesgcm.encrypt(nonce, raw_bytes, _ECP1_MAGIC)  # AAD = magic bytes
    return _ECP1_MAGIC + salt + nonce + ct


def _decrypt_private_key(data: bytes, password: bytes) -> bytes:
    """Decrypt key bytes encrypted with _encrypt_private_key."""
    if not data.startswith(_ECP1_MAGIC):
        raise ValueError("Not ECP1 format")
    salt = data[4:20]
    nonce = data[20:32]
    ct = data[32:]
    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32,
                      salt=salt, iterations=PBKDF2_ITERATIONS)
    derived = kdf.derive(password)
    aesgcm = AESGCM(derived)
    return aesgcm.decrypt(nonce, ct, _ECP1_MAGIC)


# ---------------------------------------------------------------------------
# RSA (login challenge-response ONLY)
# ---------------------------------------------------------------------------

def generate_rsa_keypair(key_size: int = 4096) -> tuple[rsa.RSAPrivateKey, rsa.RSAPublicKey]:
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)
    return private_key, private_key.public_key()


def serialize_private_key(key: rsa.RSAPrivateKey, password: bytes | None = None) -> bytes:
    if password:
        raw = key.private_bytes(serialization.Encoding.DER, serialization.PrivateFormat.PKCS8,
                                serialization.NoEncryption())
        return _encrypt_private_key(raw, password)
    return key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8,
                             serialization.NoEncryption())


def serialize_public_key(key: rsa.RSAPublicKey) -> bytes:
    return key.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)


def load_private_key(data: bytes, password: bytes | None = None) -> rsa.RSAPrivateKey:
    if data.startswith(_ECP1_MAGIC):
        raw = _decrypt_private_key(data, password)
        return serialization.load_der_private_key(raw, password=None)
    # Legacy PEM format (old BestAvailableEncryption or unencrypted)
    return serialization.load_pem_private_key(data, password=password)


def load_public_key(pem: bytes) -> rsa.RSAPublicKey:
    return serialization.load_pem_public_key(pem)


def rsa_sign(private_key: rsa.RSAPrivateKey, data: bytes) -> bytes:
    return private_key.sign(
        data,
        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
        hashes.SHA256(),
    )


def rsa_verify(public_key: rsa.RSAPublicKey, signature: bytes, data: bytes) -> bool:
    try:
        public_key.verify(
            signature, data,
            padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.AUTO),
            hashes.SHA256(),
        )
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# AES-256-GCM (symmetric encryption — used by ratchet message keys & images)
# ---------------------------------------------------------------------------

def aes_encrypt(plaintext: bytes, key: bytes | None = None) -> tuple[bytes, bytes, bytes, bytes]:
    """Encrypt with AES-256-GCM. Returns (key, nonce, ciphertext, tag)."""
    if key is None:
        key = AESGCM.generate_key(bit_length=256)
    nonce = os.urandom(12)
    aesgcm = AESGCM(key)
    ct_with_tag = aesgcm.encrypt(nonce, plaintext, None)
    ciphertext = ct_with_tag[:-16]
    tag = ct_with_tag[-16:]
    return key, nonce, ciphertext, tag


def aes_decrypt(key: bytes, nonce: bytes, ciphertext: bytes, tag: bytes) -> bytes:
    """Decrypt with AES-256-GCM."""
    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, ciphertext + tag, None)


# ---------------------------------------------------------------------------
# Ed25519 Identity Keys
# ---------------------------------------------------------------------------

def generate_identity_keypair() -> tuple[Ed25519PrivateKey, Ed25519PublicKey]:
    priv = Ed25519PrivateKey.generate()
    return priv, priv.public_key()


def serialize_ed25519_private(key: Ed25519PrivateKey, password: bytes | None = None) -> bytes:
    if password:
        raw = serialize_ed25519_private_raw(key)  # 32 bytes
        return _encrypt_private_key(raw, password)
    return serialize_ed25519_private_raw(key)  # 32 bytes, no password


def serialize_ed25519_private_raw(key: Ed25519PrivateKey) -> bytes:
    """Serialize Ed25519 private key to 32 raw bytes (unencrypted)."""
    return key.private_bytes(serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption())


def serialize_ed25519_public(key: Ed25519PublicKey) -> bytes:
    """Serialize Ed25519 public key to 32 raw bytes."""
    return key.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)


def load_ed25519_private(data: bytes, password: bytes | None = None) -> Ed25519PrivateKey:
    if data.startswith(_ECP1_MAGIC):
        raw = _decrypt_private_key(data, password)
        return Ed25519PrivateKey.from_private_bytes(raw)
    # Legacy formats: PEM (old BestAvailableEncryption) or 32-byte raw
    if password:
        return serialization.load_pem_private_key(data, password=password)
    if len(data) == 32:
        return Ed25519PrivateKey.from_private_bytes(data)
    return serialization.load_pem_private_key(data, password=None)


def load_ed25519_public(data: bytes) -> Ed25519PublicKey:
    if len(data) == 32:
        return Ed25519PublicKey.from_public_bytes(data)
    return serialization.load_pem_public_key(data)


def ed25519_sign(private_key: Ed25519PrivateKey, data: bytes) -> bytes:
    """Sign data with Ed25519. Returns 64-byte signature."""
    return private_key.sign(data)


def ed25519_verify(public_key: Ed25519PublicKey, signature: bytes, data: bytes) -> bool:
    """Verify Ed25519 signature."""
    try:
        public_key.verify(signature, data)
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# X25519 Key Exchange
# ---------------------------------------------------------------------------

def generate_x25519_keypair() -> tuple[X25519PrivateKey, X25519PublicKey]:
    priv = X25519PrivateKey.generate()
    return priv, priv.public_key()


def serialize_x25519_private(key: X25519PrivateKey) -> bytes:
    """Serialize X25519 private key to 32 raw bytes."""
    return key.private_bytes(serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption())


def serialize_x25519_public(key: X25519PublicKey) -> bytes:
    """Serialize X25519 public key to 32 raw bytes."""
    return key.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)


def load_x25519_private(data: bytes) -> X25519PrivateKey:
    return X25519PrivateKey.from_private_bytes(data)


def load_x25519_public(data: bytes) -> X25519PublicKey:
    return X25519PublicKey.from_public_bytes(data)


def x25519_dh(private_key: X25519PrivateKey, public_key: X25519PublicKey) -> bytes:
    """Perform X25519 Diffie-Hellman. Returns 32-byte shared secret."""
    return private_key.exchange(public_key)


# ---------------------------------------------------------------------------
# Ed25519 <-> X25519 conversion (for Identity Key dual use)
# ---------------------------------------------------------------------------

def ed25519_private_to_x25519(ed_private: Ed25519PrivateKey) -> X25519PrivateKey:
    """Derive X25519 private key from Ed25519 private key via RFC 7748 clamping."""
    raw = ed_private.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption()
    )
    # SHA-512 hash of the seed, take first 32 bytes, clamp per RFC 7748
    h = hashlib.sha512(raw).digest()[:32]
    clamped = bytearray(h)
    clamped[0] &= 248
    clamped[31] &= 127
    clamped[31] |= 64
    return X25519PrivateKey.from_private_bytes(bytes(clamped))


def ed25519_public_to_x25519(ed_public: Ed25519PublicKey) -> X25519PublicKey:
    """Derive X25519 public key from Ed25519 public key.

    Uses the cryptography library's internal conversion. For production use,
    we compute the X25519 public key from the converted private key when possible.
    For remote keys (where we don't have the private key), we use a pure-Python
    implementation of the Ed25519->X25519 point conversion.
    """
    # Montgomery u = (1 + y) / (1 - y) mod p, where p = 2^255 - 19
    raw = ed_public.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    y = int.from_bytes(raw, "little")
    # Clear the sign bit
    y &= (1 << 255) - 1
    p = (1 << 255) - 19
    # u = (1 + y) * inverse(1 - y) mod p
    one_plus_y = (1 + y) % p
    one_minus_y = (1 - y) % p
    inv = pow(one_minus_y, p - 2, p)
    u = (one_plus_y * inv) % p
    x25519_bytes = u.to_bytes(32, "little")
    return X25519PublicKey.from_public_bytes(x25519_bytes)


# ---------------------------------------------------------------------------
# HKDF
# ---------------------------------------------------------------------------

_HKDF_INFO_SELF = b"EncryptedChat_SelfKey"
_HKDF_INFO_RK = b"EncryptedChat_RootKey"


def derive_self_encryption_key(identity_private: Ed25519PrivateKey) -> bytes:
    """Derive a static AES-256 key from identity key for encrypting own sent messages.

    This is NOT a ratchet — it's a static key. Safe because only the owner
    has the identity private key, and self-copies don't need forward secrecy.
    """
    raw = identity_private.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption()
    )
    return hkdf_derive(raw, salt=b"self_encryption", info=_HKDF_INFO_SELF, length=32)


_HKDF_INFO_LOCAL = b"EncryptedChat_LocalStorage"


def derive_local_storage_key(identity_private: Ed25519PrivateKey) -> bytes:
    """Derive AES-256 key for encrypting local session/sender key files."""
    raw = identity_private.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption()
    )
    return hkdf_derive(raw, salt=b"local_storage", info=_HKDF_INFO_LOCAL, length=32)


_HKDF_INFO_CK_MSG = b"\x01"  # chain key -> message key
_HKDF_INFO_CK_NEXT = b"\x02"  # chain key -> next chain key


def hkdf_derive(input_key: bytes, salt: bytes, info: bytes, length: int = 32) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=length, salt=salt, info=info).derive(input_key)


def kdf_rk(root_key: bytes, dh_output: bytes) -> tuple[bytes, bytes]:
    """Root key KDF. Returns (new_root_key, chain_key).

    Uses HKDF with the root key as salt and DH output as input key material.
    Derives 64 bytes: first 32 = new root key, last 32 = chain key.
    """
    derived = hkdf_derive(dh_output, salt=root_key, info=_HKDF_INFO_RK, length=64)
    return derived[:32], derived[32:]


def kdf_ck(chain_key: bytes) -> tuple[bytes, bytes]:
    """Chain key KDF. Returns (new_chain_key, message_key).

    Uses HMAC-SHA256:
      message_key = HMAC(chain_key, 0x01)
      new_chain_key = HMAC(chain_key, 0x02)
    """
    message_key = hmac.new(chain_key, _HKDF_INFO_CK_MSG, hashlib.sha256).digest()
    new_chain_key = hmac.new(chain_key, _HKDF_INFO_CK_NEXT, hashlib.sha256).digest()
    return new_chain_key, message_key


# ---------------------------------------------------------------------------
# X3DH
# ---------------------------------------------------------------------------

_X3DH_INFO = b"EncryptedChat_X3DH"


def generate_signed_prekey(identity_private: Ed25519PrivateKey) -> dict:
    """Generate a signed pre-key (SPK).

    Returns {private: X25519PrivateKey, public: X25519PublicKey, signature: bytes, id: str}.
    """
    spk_priv, spk_pub = generate_x25519_keypair()
    spk_pub_bytes = serialize_x25519_public(spk_pub)
    signature = ed25519_sign(identity_private, spk_pub_bytes)
    return {
        "private": spk_priv,
        "public": spk_pub,
        "signature": signature,
        "id": str(uuid.uuid4()),
    }


def generate_one_time_prekeys(count: int = 50) -> list[dict]:
    """Generate a batch of one-time pre-keys.

    Returns [{private: X25519PrivateKey, public: X25519PublicKey, id: str}, ...].
    """
    result = []
    for _ in range(count):
        priv, pub = generate_x25519_keypair()
        result.append({"private": priv, "public": pub, "id": str(uuid.uuid4())})
    return result


def x3dh_initiate(
    ik_private_ed: Ed25519PrivateKey,
    ik_public_remote_ed: Ed25519PublicKey,
    spk_remote: X25519PublicKey,
    spk_signature: bytes,
    opk_remote: X25519PublicKey | None = None,
) -> tuple[bytes, X25519PrivateKey, X25519PublicKey]:
    """Initiator side of X3DH.

    Args:
        ik_private_ed: Our Ed25519 identity private key
        ik_public_remote_ed: Remote Ed25519 identity public key
        spk_remote: Remote signed pre-key (X25519 public)
        spk_signature: Ed25519 signature of spk_remote by ik_public_remote_ed
        opk_remote: Optional one-time pre-key (X25519 public)

    Returns:
        (shared_secret, ephemeral_private, ephemeral_public)
    """
    # Verify SPK signature
    spk_remote_bytes = serialize_x25519_public(spk_remote)
    if not ed25519_verify(ik_public_remote_ed, spk_signature, spk_remote_bytes):
        raise ValueError("Invalid SPK signature")

    # Convert identity keys to X25519
    ik_x25519_private = ed25519_private_to_x25519(ik_private_ed)
    ik_x25519_remote = ed25519_public_to_x25519(ik_public_remote_ed)

    # Generate ephemeral keypair
    ek_priv, ek_pub = generate_x25519_keypair()

    # DH computations
    dh1 = x25519_dh(ik_x25519_private, spk_remote)   # IK_A, SPK_B
    dh2 = x25519_dh(ek_priv, ik_x25519_remote)        # EK_A, IK_B
    dh3 = x25519_dh(ek_priv, spk_remote)               # EK_A, SPK_B

    dh_concat = dh1 + dh2 + dh3
    if opk_remote is not None:
        dh4 = x25519_dh(ek_priv, opk_remote)            # EK_A, OPK_B
        dh_concat += dh4

    # Derive shared secret
    shared_secret = hkdf_derive(dh_concat, salt=b"\x00" * 32, info=_X3DH_INFO, length=32)
    return shared_secret, ek_priv, ek_pub


def x3dh_respond(
    ik_private_ed: Ed25519PrivateKey,
    spk_private: X25519PrivateKey,
    ik_remote_ed: Ed25519PublicKey,
    ek_remote: X25519PublicKey,
    opk_private: X25519PrivateKey | None = None,
) -> bytes:
    """Responder side of X3DH.

    Args:
        ik_private_ed: Our Ed25519 identity private key
        spk_private: Our signed pre-key private (X25519)
        ik_remote_ed: Remote Ed25519 identity public key
        ek_remote: Remote ephemeral key (X25519 public)
        opk_private: Our one-time pre-key private (X25519), if used

    Returns:
        shared_secret (32 bytes)
    """
    ik_x25519_private = ed25519_private_to_x25519(ik_private_ed)
    ik_x25519_remote = ed25519_public_to_x25519(ik_remote_ed)

    dh1 = x25519_dh(spk_private, ik_x25519_remote)     # SPK_B, IK_A
    dh2 = x25519_dh(ik_x25519_private, ek_remote)       # IK_B, EK_A
    dh3 = x25519_dh(spk_private, ek_remote)              # SPK_B, EK_A

    dh_concat = dh1 + dh2 + dh3
    if opk_private is not None:
        dh4 = x25519_dh(opk_private, ek_remote)          # OPK_B, EK_A
        dh_concat += dh4

    shared_secret = hkdf_derive(dh_concat, salt=b"\x00" * 32, info=_X3DH_INFO, length=32)
    return shared_secret


# ---------------------------------------------------------------------------
# Double Ratchet
# ---------------------------------------------------------------------------

MAX_SKIP = 256  # max messages to skip in a single chain (out-of-order tolerance)


@dataclass
class RatchetHeader:
    """Header sent with each ratchet message."""
    dh_pub: bytes    # sender's current ratchet public key (32 bytes)
    n: int           # message number in current sending chain
    pn: int          # number of messages in previous sending chain

    def serialize(self) -> bytes:
        return json.dumps({
            "dh_pub": serialize_x25519_public(load_x25519_public(self.dh_pub)).hex()
                if isinstance(self.dh_pub, bytes) else serialize_x25519_public(self.dh_pub).hex(),
            "n": self.n,
            "pn": self.pn,
        }).encode()

    def to_dict(self) -> dict:
        pub_hex = self.dh_pub.hex() if isinstance(self.dh_pub, bytes) else \
            serialize_x25519_public(self.dh_pub).hex()
        return {"dh_pub": pub_hex, "n": self.n, "pn": self.pn}

    @classmethod
    def from_dict(cls, d: dict) -> "RatchetHeader":
        return cls(dh_pub=bytes.fromhex(d["dh_pub"]), n=d["n"], pn=d["pn"])


class DoubleRatchet:
    """Signal Double Ratchet implementation."""

    def __init__(self):
        self.dh_pair: tuple[X25519PrivateKey, X25519PublicKey] | None = None
        self.dh_remote: X25519PublicKey | None = None
        self.root_key: bytes = b""
        self.send_chain_key: bytes | None = None
        self.recv_chain_key: bytes | None = None
        self.send_n: int = 0
        self.recv_n: int = 0
        self.prev_send_n: int = 0
        # (dh_pub_hex, n) -> message_key for out-of-order messages
        self.skipped: dict[tuple[str, int], bytes] = {}

    @classmethod
    def init_alice(cls, shared_secret: bytes, bob_spk_pub: X25519PublicKey) -> "DoubleRatchet":
        """Initialize as initiator (Alice) after X3DH.

        Alice performs the first DH ratchet step immediately.
        """
        ratchet = cls()
        ratchet.dh_pair = generate_x25519_keypair()
        ratchet.dh_remote = bob_spk_pub

        # Perform DH ratchet to derive send chain
        dh_output = x25519_dh(ratchet.dh_pair[0], ratchet.dh_remote)
        ratchet.root_key, ratchet.send_chain_key = kdf_rk(shared_secret, dh_output)
        ratchet.recv_chain_key = None
        ratchet.send_n = 0
        ratchet.recv_n = 0
        ratchet.prev_send_n = 0
        return ratchet

    @classmethod
    def init_bob(cls, shared_secret: bytes, spk_pair: tuple[X25519PrivateKey, X25519PublicKey]) -> "DoubleRatchet":
        """Initialize as responder (Bob) after X3DH.

        Bob uses his SPK as the initial ratchet key pair.
        """
        ratchet = cls()
        ratchet.dh_pair = spk_pair
        ratchet.root_key = shared_secret
        ratchet.send_chain_key = None
        ratchet.recv_chain_key = None
        ratchet.send_n = 0
        ratchet.recv_n = 0
        ratchet.prev_send_n = 0
        return ratchet

    def encrypt(self, plaintext: bytes) -> dict:
        """Encrypt a message.

        Returns {header: {dh_pub, n, pn}, ciphertext: bytes, nonce: bytes}.
        """
        if self.send_chain_key is None:
            raise RuntimeError("Send chain not initialized")

        self.send_chain_key, message_key = kdf_ck(self.send_chain_key)

        header = RatchetHeader(
            dh_pub=serialize_x25519_public(self.dh_pair[1]),
            n=self.send_n,
            pn=self.prev_send_n,
        )

        # Encrypt with AES-256-GCM using the message key
        nonce = os.urandom(12)
        aesgcm = AESGCM(message_key)
        # Include header as AAD to bind ciphertext to header
        aad = header.serialize()
        ct_with_tag = aesgcm.encrypt(nonce, plaintext, aad)

        self.send_n += 1

        return {
            "header": header.to_dict(),
            "ciphertext": ct_with_tag,  # includes 16-byte tag
            "nonce": nonce,
        }

    def decrypt(self, header_dict: dict, ciphertext: bytes, nonce: bytes) -> bytes:
        """Decrypt a message. Handles DH ratchet step if new dh_pub.

        State is snapshotted before modification and restored on failure (M9 fix).
        """
        header = RatchetHeader.from_dict(header_dict)
        remote_dh_pub_bytes = header.dh_pub

        # Check if this is from a skipped message (no state modification needed)
        skip_key = (remote_dh_pub_bytes.hex(), header.n)
        if skip_key in self.skipped:
            mk = self.skipped.pop(skip_key)
            aad = header.serialize()
            aesgcm = AESGCM(mk)
            try:
                return aesgcm.decrypt(nonce, ciphertext, aad)
            except Exception:
                self.skipped[skip_key] = mk  # restore skipped key
                raise

        # Snapshot state before modifications
        snap = self._snapshot()

        try:
            remote_dh_pub = load_x25519_public(remote_dh_pub_bytes)
            current_remote_bytes = serialize_x25519_public(self.dh_remote) if self.dh_remote else None

            if current_remote_bytes is None or remote_dh_pub_bytes != current_remote_bytes:
                # New DH ratchet step
                self._skip_messages(header.pn)
                self._dh_ratchet(remote_dh_pub)

            self._skip_messages(header.n)

            # Derive message key from receive chain
            self.recv_chain_key, mk = kdf_ck(self.recv_chain_key)
            self.recv_n += 1

            aad = header.serialize()
            aesgcm = AESGCM(mk)
            return aesgcm.decrypt(nonce, ciphertext, aad)
        except Exception:
            self._restore(snap)
            raise

    def _snapshot(self) -> dict:
        """Capture mutable state for rollback on decrypt failure."""
        return {
            "dh_pair": self.dh_pair,
            "dh_remote": self.dh_remote,
            "root_key": self.root_key,
            "send_chain_key": self.send_chain_key,
            "recv_chain_key": self.recv_chain_key,
            "send_n": self.send_n,
            "recv_n": self.recv_n,
            "prev_send_n": self.prev_send_n,
            "skipped": dict(self.skipped),
        }

    def _restore(self, snap: dict):
        """Restore state from snapshot."""
        self.dh_pair = snap["dh_pair"]
        self.dh_remote = snap["dh_remote"]
        self.root_key = snap["root_key"]
        self.send_chain_key = snap["send_chain_key"]
        self.recv_chain_key = snap["recv_chain_key"]
        self.send_n = snap["send_n"]
        self.recv_n = snap["recv_n"]
        self.prev_send_n = snap["prev_send_n"]
        self.skipped = snap["skipped"]

    def _skip_messages(self, until: int):
        """Skip ahead in the receive chain, storing message keys for out-of-order delivery."""
        if self.recv_chain_key is None:
            return
        if until - self.recv_n > MAX_SKIP:
            raise RuntimeError(f"Too many skipped messages ({until - self.recv_n} > {MAX_SKIP})")
        while self.recv_n < until:
            self.recv_chain_key, mk = kdf_ck(self.recv_chain_key)
            remote_hex = serialize_x25519_public(self.dh_remote).hex() if self.dh_remote else ""
            self.skipped[(remote_hex, self.recv_n)] = mk
            self.recv_n += 1

    def _dh_ratchet(self, remote_dh_pub: X25519PublicKey):
        """Perform a DH ratchet step: update receive chain, generate new DH pair, update send chain."""
        self.prev_send_n = self.send_n
        self.send_n = 0
        self.recv_n = 0
        self.dh_remote = remote_dh_pub

        # Derive new receive chain key
        dh_output = x25519_dh(self.dh_pair[0], self.dh_remote)
        self.root_key, self.recv_chain_key = kdf_rk(self.root_key, dh_output)

        # Generate new DH pair and derive new send chain key
        self.dh_pair = generate_x25519_keypair()
        dh_output = x25519_dh(self.dh_pair[0], self.dh_remote)
        self.root_key, self.send_chain_key = kdf_rk(self.root_key, dh_output)

    def export_state(self) -> bytes:
        """Serialize full ratchet state for persistent storage."""
        state = {
            "dh_priv": serialize_x25519_private(self.dh_pair[0]).hex() if self.dh_pair else None,
            "dh_pub": serialize_x25519_public(self.dh_pair[1]).hex() if self.dh_pair else None,
            "dh_remote": serialize_x25519_public(self.dh_remote).hex() if self.dh_remote else None,
            "root_key": self.root_key.hex(),
            "send_ck": self.send_chain_key.hex() if self.send_chain_key else None,
            "recv_ck": self.recv_chain_key.hex() if self.recv_chain_key else None,
            "send_n": self.send_n,
            "recv_n": self.recv_n,
            "prev_send_n": self.prev_send_n,
            "skipped": {f"{k[0]}:{k[1]}": v.hex() for k, v in self.skipped.items()},
        }
        return json.dumps(state).encode()

    @classmethod
    def import_state(cls, data: bytes) -> "DoubleRatchet":
        """Deserialize ratchet state."""
        state = json.loads(data)
        r = cls()
        if state["dh_priv"] and state["dh_pub"]:
            priv = load_x25519_private(bytes.fromhex(state["dh_priv"]))
            pub = load_x25519_public(bytes.fromhex(state["dh_pub"]))
            r.dh_pair = (priv, pub)
        if state["dh_remote"]:
            r.dh_remote = load_x25519_public(bytes.fromhex(state["dh_remote"]))
        r.root_key = bytes.fromhex(state["root_key"])
        r.send_chain_key = bytes.fromhex(state["send_ck"]) if state["send_ck"] else None
        r.recv_chain_key = bytes.fromhex(state["recv_ck"]) if state["recv_ck"] else None
        r.send_n = state["send_n"]
        r.recv_n = state["recv_n"]
        r.prev_send_n = state["prev_send_n"]
        r.skipped = {}
        for k_str, v_hex in state.get("skipped", {}).items():
            parts = k_str.rsplit(":", 1)
            dh_hex = parts[0]
            n = int(parts[1])
            r.skipped[(dh_hex, n)] = bytes.fromhex(v_hex)
        return r


# ---------------------------------------------------------------------------
# Sender Keys (group messaging)
# ---------------------------------------------------------------------------

class SenderKeyState:
    """Sender key chain for group messaging.

    Each sender in a group has their own sender key chain.
    Other group members receive the initial sender_key via pairwise Double Ratchet.
    """

    def __init__(self, sender_key: bytes | None = None):
        if sender_key is None:
            sender_key = os.urandom(32)
        self.sender_key = sender_key
        self.chain_id = hashlib.sha256(sender_key).digest()
        self.chain_key = hkdf_derive(sender_key, salt=b"\x00" * 32, info=b"SenderKeyChain", length=32)
        self.n = 0
        # For receivers: track chain state to allow fast-forward
        self._known_keys: dict[int, bytes] = {}

    def encrypt(self, plaintext: bytes) -> dict:
        """Encrypt with current chain key.

        Returns {chain_id: hex, n: int, ciphertext: bytes, nonce: bytes}.
        """
        self.chain_key, message_key = kdf_ck(self.chain_key)
        nonce = os.urandom(12)
        aesgcm = AESGCM(message_key)
        # AAD includes chain_id and message number
        aad = self.chain_id + struct.pack(">I", self.n)
        ct_with_tag = aesgcm.encrypt(nonce, plaintext, aad)
        result = {
            "chain_id": self.chain_id.hex(),
            "n": self.n,
            "ciphertext": ct_with_tag,
            "nonce": nonce,
        }
        self.n += 1
        return result

    MAX_SENDER_KEY_SKIP = 256

    def decrypt(self, chain_id_hex: str, n: int, ciphertext: bytes, nonce: bytes) -> bytes:
        """Decrypt a group message. Fast-forwards the chain if needed.

        State is snapshotted before modification and restored on failure (M9 fix).
        """
        chain_id = bytes.fromhex(chain_id_hex)
        if chain_id != self.chain_id:
            raise ValueError("Chain ID mismatch")

        if n - self.n > self.MAX_SENDER_KEY_SKIP:
            raise ValueError(f"Sender key skip too large ({n - self.n} > {self.MAX_SENDER_KEY_SKIP})")

        # Snapshot before fast-forward
        snap_chain_key = self.chain_key
        snap_n = self.n
        snap_known = dict(self._known_keys)

        try:
            # Fast-forward the chain to reach message n
            while self.n <= n:
                self.chain_key, mk = kdf_ck(self.chain_key)
                self._known_keys[self.n] = mk
                self.n += 1

            mk = self._known_keys.pop(n, None)
            if mk is None:
                raise ValueError(f"Message key for n={n} not available (already consumed)")

            aad = chain_id + struct.pack(">I", n)
            aesgcm = AESGCM(mk)
            return aesgcm.decrypt(nonce, ciphertext, aad)
        except Exception:
            self.chain_key = snap_chain_key
            self.n = snap_n
            self._known_keys = snap_known
            raise

    def export_key(self) -> bytes:
        """Export sender key for distribution to group members.

        Contains everything needed to initialize a receiving SenderKeyState.
        """
        return json.dumps({
            "sender_key": self.sender_key.hex(),
        }).encode()

    def export_state(self) -> bytes:
        """Serialize full state for persistent storage."""
        return json.dumps({
            "sender_key": self.sender_key.hex(),
            "chain_id": self.chain_id.hex(),
            "chain_key": self.chain_key.hex(),
            "n": self.n,
            "known_keys": {str(k): v.hex() for k, v in self._known_keys.items()},
        }).encode()

    @classmethod
    def import_state(cls, data: bytes) -> "SenderKeyState":
        state = json.loads(data)
        obj = cls.__new__(cls)
        obj.sender_key = bytes.fromhex(state["sender_key"])
        obj.chain_id = bytes.fromhex(state["chain_id"])
        obj.chain_key = bytes.fromhex(state["chain_key"])
        obj.n = state["n"]
        obj._known_keys = {int(k): bytes.fromhex(v) for k, v in state.get("known_keys", {}).items()}
        return obj

    @classmethod
    def from_key(cls, exported_key: bytes) -> "SenderKeyState":
        """Initialize a receiving SenderKeyState from an exported key."""
        data = json.loads(exported_key)
        return cls(sender_key=bytes.fromhex(data["sender_key"]))


# ---------------------------------------------------------------------------
# Contact Key Verification (Safety Numbers / Fingerprints / QR Codes)
# ---------------------------------------------------------------------------

FINGERPRINT_VERSION = 0


def compute_fingerprint(user_id: str, identity_key_bytes: bytes, iterations: int = 5200) -> bytes:
    """Compute a 32-byte fingerprint for a user's identity key.

    Uses iterated SHA-512 (Signal's NumericFingerprint algorithm).
    Seed: version(2B) + identity_key(32B) + user_id(UTF-8).
    Each iteration: SHA-512(previous_hash + identity_key).
    Output: first 32 bytes of final hash.
    """
    version_bytes = FINGERPRINT_VERSION.to_bytes(2, "big")
    data = version_bytes + identity_key_bytes + user_id.encode("utf-8")
    for _ in range(iterations):
        data = hashlib.sha512(data + identity_key_bytes).digest()
    return data[:32]


def format_fingerprint(fp_bytes: bytes) -> str:
    """Format 32-byte fingerprint as 6 groups of 5 zero-padded digits (30 digits).

    Each group: int(bytes[i*5:(i+1)*5], big-endian) % 100000.
    Output: two lines of 3 groups each, space-separated.
    """
    groups = []
    for i in range(6):
        num = int.from_bytes(fp_bytes[i * 5:(i + 1) * 5], "big") % 100000
        groups.append(f"{num:05d}")
    return " ".join(groups[:3]) + "\n" + " ".join(groups[3:])


def compute_safety_number(my_uid: str, my_ik_bytes: bytes,
                          their_uid: str, their_ik_bytes: bytes) -> str:
    """Compute a 60-digit safety number for a pair of users.

    Both users see the same number regardless of who computes it.
    Lower user_id's fingerprint comes first (deterministic ordering).
    Output: 12 groups of 5 digits, formatted as 3 lines of 4 groups.
    """
    fp_mine = compute_fingerprint(my_uid, my_ik_bytes)
    fp_theirs = compute_fingerprint(their_uid, their_ik_bytes)
    if my_uid < their_uid:
        combined = fp_mine + fp_theirs
    else:
        combined = fp_theirs + fp_mine
    # 64 bytes -> 12 groups of 5 digits
    groups = []
    for i in range(12):
        num = int.from_bytes(combined[i * 5:(i + 1) * 5], "big") % 100000
        groups.append(f"{num:05d}")
    lines = [
        " ".join(groups[0:4]),
        " ".join(groups[4:8]),
        " ".join(groups[8:12]),
    ]
    return "\n".join(lines)


def encode_verification_qr(user_id: str, identity_key_bytes: bytes) -> bytes:
    """Encode user identity for QR code verification.

    Format: version(1B=0x01) + uid_len(1B) + uid(UTF-8) + identity_key(32B).
    """
    uid_bytes = user_id.encode("utf-8")
    return b"\x01" + len(uid_bytes).to_bytes(1, "big") + uid_bytes + identity_key_bytes


def decode_verification_qr(data: bytes) -> tuple[str, bytes]:
    """Decode QR code verification payload.

    Returns (user_id, identity_key_bytes).
    Raises ValueError on invalid format.
    """
    if len(data) < 3:
        raise ValueError("QR data too short")
    if data[0] != 0x01:
        raise ValueError(f"Unknown QR version: {data[0]}")
    uid_len = data[1]
    if len(data) < 2 + uid_len + 32:
        raise ValueError("QR data truncated")
    user_id = data[2:2 + uid_len].decode("utf-8")
    identity_key = data[2 + uid_len:2 + uid_len + 32]
    return user_id, identity_key


# ---------------------------------------------------------------------------
# Message Padding (metadata privacy — hide plaintext length)
# ---------------------------------------------------------------------------

_PAD_MAGIC = b"\x01"
_PAD_BUCKETS = [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]


def pad_plaintext(plaintext: bytes) -> bytes:
    """Pad plaintext to nearest bucket size to hide message length.

    Format: 0x01 + plaintext + random_padding + pad_length(4B big-endian)
    Prefix 0x01 distinguishes padded messages from legacy unpadded (which start with '{').
    """
    content = _PAD_MAGIC + plaintext
    # +4 for the length suffix
    min_size = len(content) + 4
    target = next((b for b in _PAD_BUCKETS if b >= min_size), min_size)
    pad_len = target - len(content)
    return content + os.urandom(pad_len - 4) + struct.pack(">I", pad_len)


def unpad_plaintext(data: bytes) -> bytes:
    """Remove padding. Returns raw plaintext for both padded and legacy unpadded messages."""
    if not data or data[0:1] != _PAD_MAGIC:
        return data  # legacy unpadded message (starts with '{' for JSON)
    if len(data) < 5:
        return data  # too short to be validly padded
    pad_len = struct.unpack(">I", data[-4:])[0]
    if pad_len < 4 or pad_len > len(data) - 1:
        return data  # invalid padding metadata, treat as legacy
    return data[1:len(data) - pad_len]
