"""
MeshBay — Double Ratchet Algorithm implementation.

Based on the Signal Protocol specification:
  https://signal.org/docs/specifications/doubleratchet/

Provides forward secrecy and break-in recovery for group messaging.
Each message is encrypted with a unique key derived from the ratchet state.
Compromise of the current state does not reveal past message keys.

Key components:
  - DH Ratchet:       rotates Diffie-Hellman keys to achieve break-in recovery
  - Symmetric Ratchet: derives unique per-message keys from a chain key
  - KDF functions:     HKDF-based key derivation following the Signal spec

Usage:
    # Initialise from a shared secret (e.g. from X3DH or GEK)
    alice_state = RatchetState.init_sender(shared_secret, bob_public_key)
    bob_state   = RatchetState.init_receiver(shared_secret, bob_private_key)

    # Alice sends
    header, ciphertext = alice_state.encrypt(b"Hello Bob")

    # Bob receives
    plaintext = bob_state.decrypt(header, ciphertext)
    assert plaintext == b"Hello Bob"
"""

import os
import struct
from dataclasses import dataclass, field
from typing import Optional

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 import hashes, serialization

# KDF info strings (stable identifiers)
_KDF_RK_INFO  = b"meshbay:ratchet:root:v1"
_KDF_CK_INFO  = b"meshbay:ratchet:chain:v1"
_KDF_MSG_INFO = b"meshbay:ratchet:msg:v1"

# Maximum out-of-order messages stored per ratchet step
MAX_SKIP = 1000


# ── Key helpers ───────────────────────────────────────────────────────────────

def _dh_generate() -> X25519PrivateKey:
    return X25519PrivateKey.generate()

def _dh_pub_bytes(sk: X25519PrivateKey) -> bytes:
    return sk.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)

def _dh_pub_from_bytes(b: bytes) -> X25519PublicKey:
    return X25519PublicKey.from_public_bytes(b)

def _dh(sk: X25519PrivateKey, pk: X25519PublicKey) -> bytes:
    return sk.exchange(pk)

def _kdf_rk(root_key: bytes, dh_out: bytes) -> tuple[bytes, bytes]:
    """KDF_RK: derive new root key + chain key from root key + DH output."""
    material = HKDF(
        algorithm=hashes.SHA256(), length=64,
        salt=root_key, info=_KDF_RK_INFO,
    ).derive(dh_out)
    return material[:32], material[32:]   # new_rk, ck

def _kdf_ck(chain_key: bytes) -> tuple[bytes, bytes]:
    """KDF_CK: advance chain key, derive message key."""
    material = HKDF(
        algorithm=hashes.SHA256(), length=64,
        salt=chain_key, info=_KDF_CK_INFO,
    ).derive(b"\x01")
    return material[:32], material[32:]   # new_ck, mk

def _kdf_msg(msg_key: bytes) -> tuple[bytes, bytes, bytes]:
    """Expand msg_key into encryption_key, auth_key, iv (96-bit nonce)."""
    material = HKDF(
        algorithm=hashes.SHA256(), length=80,
        salt=b"\x00" * 32, info=_KDF_MSG_INFO,
    ).derive(msg_key)
    return material[:32], material[32:64], material[64:]   # enc_key, auth_key, iv

def _encrypt(mk: bytes, plaintext: bytes, ad: bytes) -> bytes:
    enc_key, _, iv = _kdf_msg(mk)
    return AESGCM(enc_key).encrypt(iv, plaintext, ad)

def _decrypt(mk: bytes, ciphertext: bytes, ad: bytes) -> bytes:
    enc_key, _, iv = _kdf_msg(mk)
    return AESGCM(enc_key).decrypt(iv, ciphertext, ad)


# ── Message header ────────────────────────────────────────────────────────────

@dataclass
class MessageHeader:
    """
    Wire header for a ratchet-encrypted message.
    dh_pub:       sender's current DH public key (32 bytes raw)
    prev_chain_n: number of messages in previous sending chain
    msg_num:      message number in current sending chain
    """
    dh_pub:       bytes
    prev_chain_n: int
    msg_num:      int

    def encode(self) -> bytes:
        return self.dh_pub + struct.pack(">II", self.prev_chain_n, self.msg_num)

    @classmethod
    def decode(cls, data: bytes) -> "MessageHeader":
        dh_pub = data[:32]
        prev_chain_n, msg_num = struct.unpack(">II", data[32:40])
        return cls(dh_pub=dh_pub, prev_chain_n=prev_chain_n, msg_num=msg_num)

    @property
    def encoded_len(self) -> int:
        return 40   # 32 (dh) + 4 + 4


# ── Ratchet state ─────────────────────────────────────────────────────────────

@dataclass
class RatchetState:
    """
    Full Double Ratchet state for one participant.

    Fields match Signal spec (Appendix C):
      DHs:     our sending DH key pair
      DHr:     remote's DH public key (None until first message received)
      RK:      32-byte root key
      CKs:     sending chain key (None until first send)
      CKr:     receiving chain key (None until first receive)
      Ns:      number of messages sent on current sending chain
      Nr:      number of messages received on current receiving chain
      PN:      number of messages sent in previous sending chain
      MKSKIP:  skipped message keys {(dh_pub_bytes, msg_num) → mk}
    """
    DHs:    X25519PrivateKey
    DHr:    Optional[X25519PublicKey]
    RK:     bytes
    CKs:    Optional[bytes]
    CKr:    Optional[bytes]
    Ns:     int = 0
    Nr:     int = 0
    PN:     int = 0
    MKSKIP: dict = field(default_factory=dict)

    # ── Initialisation ────────────────────────────────────────────────────────

    @classmethod
    def init_sender(
        cls, shared_secret: bytes, remote_public_key: bytes
    ) -> "RatchetState":
        """
        Initialise state as the message sender.
        shared_secret: pre-shared key from X3DH or GEK
        remote_public_key: recipient's initial X25519 public key (raw 32 bytes)
        """
        dhs = _dh_generate()
        dhr = _dh_pub_from_bytes(remote_public_key)
        rk, cks = _kdf_rk(shared_secret, _dh(dhs, dhr))
        return cls(DHs=dhs, DHr=dhr, RK=rk, CKs=cks, CKr=None)

    @classmethod
    def init_receiver(
        cls, shared_secret: bytes, own_private_key: X25519PrivateKey
    ) -> "RatchetState":
        """
        Initialise state as the message receiver.
        shared_secret: same pre-shared key used by sender
        own_private_key: the X25519 private key whose public key was given to sender
        """
        return cls(
            DHs=own_private_key,
            DHr=None,
            RK=shared_secret,
            CKs=None,
            CKr=None,
        )

    # ── Encryption ────────────────────────────────────────────────────────────

    def encrypt(self, plaintext: bytes, associated_data: bytes = b"") -> tuple[MessageHeader, bytes]:
        """Encrypt a message. Returns (header, ciphertext)."""
        assert self.CKs is not None, "Not initialised as sender"
        self.CKs, mk = _kdf_ck(self.CKs)
        header = MessageHeader(
            dh_pub=_dh_pub_bytes(self.DHs),
            prev_chain_n=self.PN,
            msg_num=self.Ns,
        )
        self.Ns += 1
        ct = _encrypt(mk, plaintext, associated_data + header.encode())
        return header, ct

    # ── Decryption ────────────────────────────────────────────────────────────

    def decrypt(self, header: MessageHeader, ciphertext: bytes,
                associated_data: bytes = b"") -> bytes:
        """Decrypt a message. Handles out-of-order delivery via MKSKIP."""
        ad = associated_data + header.encode()

        # Check skipped keys first
        skip_key = (header.dh_pub, header.msg_num)
        if skip_key in self.MKSKIP:
            mk = self.MKSKIP.pop(skip_key)
            return _decrypt(mk, ciphertext, ad)

        # DH ratchet step if new DH key received
        if self.DHr is None or header.dh_pub != _dh_pub_bytes_from_pk(self.DHr):
            self._skip_message_keys(header.prev_chain_n)
            self._dh_ratchet(header)

        # Advance receiving chain
        self._skip_message_keys(header.msg_num)
        assert self.CKr is not None
        self.CKr, mk = _kdf_ck(self.CKr)
        self.Nr += 1
        return _decrypt(mk, ciphertext, ad)

    def _skip_message_keys(self, until: int) -> None:
        """Store skipped message keys for out-of-order delivery."""
        if self.Nr + MAX_SKIP < until:
            raise ValueError(f"Too many skipped messages: {until - self.Nr}")
        if self.CKr is None:
            return
        while self.Nr < until:
            self.CKr, mk = _kdf_ck(self.CKr)
            self.MKSKIP[(bytes(_dh_pub_bytes_from_pk(self.DHr)), self.Nr)] = mk
            self.Nr += 1

    def _dh_ratchet(self, header: MessageHeader) -> None:
        """Perform a DH ratchet step on receiving a new remote DH key."""
        self.PN = self.Ns
        self.Ns = 0
        self.Nr = 0
        self.DHr = _dh_pub_from_bytes(header.dh_pub)
        self.RK, self.CKr = _kdf_rk(self.RK, _dh(self.DHs, self.DHr))
        self.DHs = _dh_generate()
        self.RK, self.CKs = _kdf_rk(self.RK, _dh(self.DHs, self.DHr))


def _dh_pub_bytes_from_pk(pk: X25519PublicKey | None) -> bytes:
    if pk is None:
        return b"\x00" * 32
    return pk.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)


# ── Group chat helpers ────────────────────────────────────────────────────────

@dataclass
class ChatMessage:
    """A ratchet-encrypted group chat message."""
    sender_id:   str
    header_enc:  bytes      # encoded MessageHeader
    ciphertext:  bytes
    timestamp:   int        # unix timestamp
    msg_id:      str        # uuid4

    def to_dict(self) -> dict:
        import base64
        return {
            "sender_id":  self.sender_id,
            "header_b64": base64.b64encode(self.header_enc).decode(),
            "ct_b64":     base64.b64encode(self.ciphertext).decode(),
            "timestamp":  self.timestamp,
            "msg_id":     self.msg_id,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "ChatMessage":
        import base64
        return cls(
            sender_id=d["sender_id"],
            header_enc=base64.b64decode(d["header_b64"]),
            ciphertext=base64.b64decode(d["ct_b64"]),
            timestamp=d["timestamp"],
            msg_id=d["msg_id"],
        )


def encrypt_chat_message(
    state: RatchetState,
    sender_id: str,
    plaintext: str | bytes,
) -> ChatMessage:
    """Encrypt a chat message using the Double Ratchet state."""
    import time, uuid
    if isinstance(plaintext, str):
        plaintext = plaintext.encode()
    header, ct = state.encrypt(plaintext)
    return ChatMessage(
        sender_id=sender_id,
        header_enc=header.encode(),
        ciphertext=ct,
        timestamp=int(time.time()),
        msg_id=str(uuid.uuid4()),
    )


def decrypt_chat_message(
    state: RatchetState,
    msg: ChatMessage,
) -> bytes:
    """Decrypt a chat message using the Double Ratchet state."""
    header = MessageHeader.decode(msg.header_enc)
    return state.decrypt(header, msg.ciphertext)