# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Familiar Secure Transport Protocol

Signal-grade end-to-end encryption for mesh communication.

Implements:
- X3DH (Extended Triple Diffie-Hellman) for initial key exchange
- Double Ratchet for forward secrecy and post-compromise recovery
- HKDF for key derivation
- AES-256-GCM or ChaCha20-Poly1305 for message encryption

Security Properties:
- Forward secrecy: Past messages safe if long-term key compromised
- Post-compromise security: Future messages safe after breach heals
- Replay protection: Nonce + message counter
- Deniable authentication: Cannot prove who sent a message to third party

Cryptographic Primitives (from PyNaCl/libsodium):
- X25519: Elliptic curve Diffie-Hellman
- Ed25519: Signatures
- XSalsa20-Poly1305 / AES-256-GCM: Authenticated encryption
- BLAKE2b / SHA-512: Hashing
- HKDF: Key derivation

References:
- Signal Protocol: https://signal.org/docs/
- X3DH: https://signal.org/docs/specifications/x3dh/
- Double Ratchet: https://signal.org/docs/specifications/doubleratchet/

Author: Familiar Project
License: MIT
"""

from __future__ import annotations

import hashlib
import hmac
import json
import logging
import os
import secrets
import struct
import time
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


# ============================================================
# CRYPTOGRAPHIC PRIMITIVES
# ============================================================

try:
    import nacl.encoding  # noqa: F401
    from nacl.bindings import (
        crypto_scalarmult,
        crypto_sign_ed25519_pk_to_curve25519,
        crypto_sign_ed25519_sk_to_curve25519,
    )
    from nacl.exceptions import CryptoError  # noqa: F401
    from nacl.public import (  # noqa: F401 - Box/PublicKey available for callers
        Box,
        PrivateKey,
        PublicKey,
    )
    from nacl.secret import SecretBox
    from nacl.signing import SigningKey, VerifyKey
    from nacl.utils import random as nacl_random

    NACL_AVAILABLE = True
except ImportError:
    NACL_AVAILABLE = False
    logger.debug("PyNaCl not installed. Secure transport unavailable.")


def _check_nacl():
    """Ensure PyNaCl is available."""
    if not NACL_AVAILABLE:
        raise ImportError("PyNaCl required for secure transport.")


# ============================================================
# CONSTANTS
# ============================================================

# Protocol version for future compatibility
PROTOCOL_VERSION = 1

# Key sizes (bytes)
KEY_SIZE = 32  # X25519/Ed25519 key size
NONCE_SIZE = 24  # XSalsa20-Poly1305 nonce
MAC_SIZE = 16  # Poly1305 MAC
SIGNATURE_SIZE = 64  # Ed25519 signature

# HKDF constants
HKDF_HASH = hashlib.sha512

# HKDF info strings
HKDF_INFO_ROOT = b"FamiliarRootKey"
HKDF_INFO_CHAIN = b"FamiliarChainKey"
HKDF_INFO_MESSAGE = b"FamiliarMessageKey"

# Ratchet constants
MAX_SKIP = 1000  # Maximum skipped message keys to store
MAX_SKIP_AGE = 86400  # Maximum age of skipped keys (24 hours)


# Message types
class MessageType(Enum):
    X3DH_INIT = 1  # Initial X3DH key exchange
    X3DH_RESPONSE = 2  # X3DH response with signed prekey
    RATCHET = 3  # Normal ratcheted message
    PREKEY_REQUEST = 4  # Request fresh prekey bundle
    ACK = 5  # Acknowledgment


# ============================================================
# HKDF - HMAC-based Key Derivation Function (RFC 5869)
# ============================================================


def hkdf_extract(salt: bytes, input_key_material: bytes) -> bytes:
    """
    HKDF-Extract: Extract a pseudorandom key from input material.

    Args:
        salt: Random value (can be zero-length)
        input_key_material: Source key material

    Returns:
        Pseudorandom key (PRK)
    """
    if not salt:
        salt = b"\x00" * HKDF_HASH().digest_size
    return hmac.new(salt, input_key_material, HKDF_HASH).digest()


def hkdf_expand(prk: bytes, info: bytes, length: int) -> bytes:
    """
    HKDF-Expand: Expand a pseudorandom key to desired length.

    Args:
        prk: Pseudorandom key from HKDF-Extract
        info: Context/application-specific info
        length: Desired output length in bytes

    Returns:
        Output key material (OKM)
    """
    hash_len = HKDF_HASH().digest_size
    n = (length + hash_len - 1) // hash_len

    okm = b""
    t = b""

    for i in range(1, n + 1):
        t = hmac.new(prk, t + info + bytes([i]), HKDF_HASH).digest()
        okm += t

    return okm[:length]


def hkdf(input_key_material: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    """
    Full HKDF: Extract-then-Expand.

    Args:
        input_key_material: Source key material
        salt: Random value
        info: Context info
        length: Desired output length

    Returns:
        Derived key material
    """
    prk = hkdf_extract(salt, input_key_material)
    return hkdf_expand(prk, info, length)


# ============================================================
# KEY TYPES
# ============================================================


@dataclass
class KeyPair:
    """X25519 key pair for Diffie-Hellman."""

    private: bytes  # 32 bytes
    public: bytes  # 32 bytes

    @classmethod
    def generate(cls) -> "KeyPair":
        """Generate a new X25519 key pair."""
        _check_nacl()
        private_key = PrivateKey.generate()
        return cls(private=bytes(private_key), public=bytes(private_key.public_key))

    @classmethod
    def from_private(cls, private: bytes) -> "KeyPair":
        """Reconstruct key pair from private key."""
        _check_nacl()
        private_key = PrivateKey(private)
        return cls(private=bytes(private_key), public=bytes(private_key.public_key))

    def dh(self, their_public: bytes) -> bytes:
        """Perform X25519 Diffie-Hellman."""
        _check_nacl()
        return crypto_scalarmult(self.private, their_public)


@dataclass
class SigningKeyPair:
    """Ed25519 key pair for signatures."""

    private: bytes  # 32 bytes (seed)
    public: bytes  # 32 bytes

    @classmethod
    def generate(cls) -> "SigningKeyPair":
        """Generate a new Ed25519 signing key pair."""
        _check_nacl()
        signing_key = SigningKey.generate()
        return cls(private=bytes(signing_key), public=bytes(signing_key.verify_key))

    @classmethod
    def from_private(cls, private: bytes) -> "SigningKeyPair":
        """Reconstruct from private key."""
        _check_nacl()
        signing_key = SigningKey(private)
        return cls(private=bytes(signing_key), public=bytes(signing_key.verify_key))

    def sign(self, message: bytes) -> bytes:
        """Sign a message, return signature."""
        _check_nacl()
        signing_key = SigningKey(self.private)
        signed = signing_key.sign(message)
        return signed.signature

    def to_x25519_keypair(self) -> KeyPair:
        """
        Convert Ed25519 signing key to X25519 key pair for DH.

        This is a standard cryptographic conversion used by Signal, WireGuard, etc.
        """
        _check_nacl()
        signing_key = SigningKey(self.private)
        # Get full 64-byte secret key (seed + public)
        sk_bytes = bytes(signing_key._signing_key)
        # Convert to X25519
        x25519_private = crypto_sign_ed25519_sk_to_curve25519(sk_bytes)
        x25519_public = crypto_sign_ed25519_pk_to_curve25519(self.public)
        return KeyPair(private=x25519_private, public=x25519_public)

    def to_x25519_public(self) -> bytes:
        """Convert Ed25519 public key to X25519 public key."""
        _check_nacl()
        return crypto_sign_ed25519_pk_to_curve25519(self.public)

    @staticmethod
    def verify(public: bytes, message: bytes, signature: bytes) -> bool:
        """Verify a signature."""
        _check_nacl()
        try:
            verify_key = VerifyKey(public)
            verify_key.verify(message, signature)
            return True
        except Exception:
            return False


# ============================================================
# X3DH - Extended Triple Diffie-Hellman
# ============================================================


@dataclass
class PreKeyBundle:
    """
    Public key bundle for X3DH key exchange.

    Published by each node, used by others to initiate secure sessions.
    """

    identity_key: bytes  # Long-term identity public key (Ed25519)
    signed_prekey: bytes  # Medium-term prekey (X25519)
    signed_prekey_signature: bytes  # Signature of signed_prekey
    signed_prekey_id: int  # ID for key rotation
    one_time_prekey: Optional[bytes] = None  # Optional one-time key
    one_time_prekey_id: Optional[int] = None

    def to_bytes(self) -> bytes:
        """Serialize bundle for transmission."""
        data = {
            "ik": self.identity_key.hex(),
            "spk": self.signed_prekey.hex(),
            "spk_sig": self.signed_prekey_signature.hex(),
            "spk_id": self.signed_prekey_id,
        }
        if self.one_time_prekey:
            data["opk"] = self.one_time_prekey.hex()
            data["opk_id"] = self.one_time_prekey_id
        return json.dumps(data).encode()

    @classmethod
    def from_bytes(cls, data: bytes) -> "PreKeyBundle":
        """Deserialize bundle."""
        d = json.loads(data.decode())
        return cls(
            identity_key=bytes.fromhex(d["ik"]),
            signed_prekey=bytes.fromhex(d["spk"]),
            signed_prekey_signature=bytes.fromhex(d["spk_sig"]),
            signed_prekey_id=d["spk_id"],
            one_time_prekey=bytes.fromhex(d["opk"]) if "opk" in d else None,
            one_time_prekey_id=d.get("opk_id"),
        )

    def verify(self) -> bool:
        """Verify the signed prekey signature."""
        # Convert Ed25519 identity key to verify the X25519 prekey
        # The signature is over the prekey bytes
        return SigningKeyPair.verify(
            self.identity_key, self.signed_prekey, self.signed_prekey_signature
        )


@dataclass
class X3DHKeys:
    """Local keys for X3DH protocol."""

    identity_key: SigningKeyPair  # Long-term identity
    signed_prekey: KeyPair  # Medium-term (rotate weekly)
    signed_prekey_id: int
    signed_prekey_signature: bytes
    one_time_prekeys: Dict[int, KeyPair] = field(default_factory=dict)
    next_prekey_id: int = 1

    @classmethod
    def generate(cls, num_one_time_keys: int = 10) -> "X3DHKeys":
        """Generate a fresh set of X3DH keys."""
        identity = SigningKeyPair.generate()
        signed_prekey = KeyPair.generate()

        # Sign the prekey with identity key
        signature = identity.sign(signed_prekey.public)

        # Generate one-time prekeys
        one_time = {}
        for i in range(num_one_time_keys):
            one_time[i + 1] = KeyPair.generate()

        return cls(
            identity_key=identity,
            signed_prekey=signed_prekey,
            signed_prekey_id=1,
            signed_prekey_signature=signature,
            one_time_prekeys=one_time,
            next_prekey_id=num_one_time_keys + 1,
        )

    def get_bundle(self, include_one_time: bool = True) -> PreKeyBundle:
        """Get public key bundle for distribution."""
        opk = None
        opk_id = None

        if include_one_time and self.one_time_prekeys:
            # Get first available one-time key
            opk_id = next(iter(self.one_time_prekeys))
            opk = self.one_time_prekeys[opk_id].public

        return PreKeyBundle(
            identity_key=self.identity_key.public,
            signed_prekey=self.signed_prekey.public,
            signed_prekey_signature=self.signed_prekey_signature,
            signed_prekey_id=self.signed_prekey_id,
            one_time_prekey=opk,
            one_time_prekey_id=opk_id,
        )

    def consume_one_time_key(self, key_id: int) -> Optional[KeyPair]:
        """Consume and return a one-time prekey."""
        return self.one_time_prekeys.pop(key_id, None)

    def generate_one_time_keys(self, count: int) -> List[Tuple[int, bytes]]:
        """Generate additional one-time prekeys."""
        new_keys = []
        for _ in range(count):
            key_id = self.next_prekey_id
            self.next_prekey_id += 1
            keypair = KeyPair.generate()
            self.one_time_prekeys[key_id] = keypair
            new_keys.append((key_id, keypair.public))
        return new_keys

    def rotate_signed_prekey(self):
        """Rotate the signed prekey (call periodically)."""
        self.signed_prekey = KeyPair.generate()
        self.signed_prekey_id += 1
        self.signed_prekey_signature = self.identity_key.sign(self.signed_prekey.public)

    def to_dict(self) -> dict:
        """Serialize for storage."""
        return {
            "identity_private": self.identity_key.private.hex(),
            "identity_public": self.identity_key.public.hex(),
            "signed_prekey_private": self.signed_prekey.private.hex(),
            "signed_prekey_public": self.signed_prekey.public.hex(),
            "signed_prekey_id": self.signed_prekey_id,
            "signed_prekey_signature": self.signed_prekey_signature.hex(),
            "one_time_prekeys": {
                str(k): {"private": v.private.hex(), "public": v.public.hex()}
                for k, v in self.one_time_prekeys.items()
            },
            "next_prekey_id": self.next_prekey_id,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "X3DHKeys":
        """Deserialize from storage."""
        return cls(
            identity_key=SigningKeyPair(
                private=bytes.fromhex(data["identity_private"]),
                public=bytes.fromhex(data["identity_public"]),
            ),
            signed_prekey=KeyPair(
                private=bytes.fromhex(data["signed_prekey_private"]),
                public=bytes.fromhex(data["signed_prekey_public"]),
            ),
            signed_prekey_id=data["signed_prekey_id"],
            signed_prekey_signature=bytes.fromhex(data["signed_prekey_signature"]),
            one_time_prekeys={
                int(k): KeyPair(
                    private=bytes.fromhex(v["private"]),
                    public=bytes.fromhex(v["public"]),
                )
                for k, v in data["one_time_prekeys"].items()
            },
            next_prekey_id=data["next_prekey_id"],
        )


@dataclass
class X3DHInitMessage:
    """
    Initial message in X3DH key exchange.

    Sent by initiator to establish shared secret.
    """

    identity_key: bytes  # Initiator's identity public key
    ephemeral_key: bytes  # Initiator's ephemeral public key
    signed_prekey_id: int  # Which signed prekey was used
    one_time_prekey_id: Optional[int]  # Which OPK was used (if any)
    ciphertext: bytes  # Initial encrypted message

    def to_bytes(self) -> bytes:
        """Serialize for transmission."""
        data = {
            "ik": self.identity_key.hex(),
            "ek": self.ephemeral_key.hex(),
            "spk_id": self.signed_prekey_id,
            "ct": self.ciphertext.hex(),
        }
        if self.one_time_prekey_id is not None:
            data["opk_id"] = self.one_time_prekey_id
        return json.dumps(data).encode()

    @classmethod
    def from_bytes(cls, data: bytes) -> "X3DHInitMessage":
        """Deserialize."""
        d = json.loads(data.decode())
        return cls(
            identity_key=bytes.fromhex(d["ik"]),
            ephemeral_key=bytes.fromhex(d["ek"]),
            signed_prekey_id=d["spk_id"],
            one_time_prekey_id=d.get("opk_id"),
            ciphertext=bytes.fromhex(d["ct"]),
        )


class X3DH:
    """
    X3DH key agreement protocol implementation.

    Provides initial key exchange with forward secrecy.
    """

    @staticmethod
    def initiate(
        our_identity: SigningKeyPair, their_bundle: PreKeyBundle, initial_message: bytes = b""
    ) -> Tuple[bytes, X3DHInitMessage]:
        """
        Initiate X3DH key exchange.

        Args:
            our_identity: Our long-term identity key
            their_bundle: Recipient's published key bundle
            initial_message: Optional first message to encrypt

        Returns:
            (shared_secret, init_message) tuple
        """
        _check_nacl()

        # Verify their signed prekey
        if not their_bundle.verify():
            raise ValueError("Invalid signed prekey signature")

        # Generate ephemeral key
        ephemeral = KeyPair.generate()

        # Convert our Ed25519 identity to X25519 for DH
        our_identity_dh = our_identity.to_x25519_keypair()

        # Convert their Ed25519 identity to X25519
        their_identity_x25519 = crypto_sign_ed25519_pk_to_curve25519(their_bundle.identity_key)

        # Perform DH calculations
        # DH1 = DH(IKa, SPKb) - our identity with their signed prekey
        dh1 = our_identity_dh.dh(their_bundle.signed_prekey)

        # DH2 = DH(EKa, IKb) - our ephemeral with their identity
        dh2 = ephemeral.dh(their_identity_x25519)

        # DH3 = DH(EKa, SPKb) - our ephemeral with their signed prekey
        dh3 = ephemeral.dh(their_bundle.signed_prekey)

        # DH4 = DH(EKa, OPKb) - our ephemeral with their one-time key (if available)
        dh4 = b""
        if their_bundle.one_time_prekey:
            dh4 = ephemeral.dh(their_bundle.one_time_prekey)

        # Combine DH outputs
        dh_concat = dh1 + dh2 + dh3 + dh4

        # Derive shared secret using HKDF
        shared_secret = hkdf(
            input_key_material=dh_concat,
            salt=b"\x00" * 32,  # Protocol salt
            info=HKDF_INFO_ROOT,
            length=32,
        )

        # Encrypt initial message with derived key
        if initial_message:
            ciphertext = X3DH._encrypt(shared_secret, initial_message)
        else:
            ciphertext = b""

        # Create init message
        init_msg = X3DHInitMessage(
            identity_key=our_identity.public,
            ephemeral_key=ephemeral.public,
            signed_prekey_id=their_bundle.signed_prekey_id,
            one_time_prekey_id=their_bundle.one_time_prekey_id,
            ciphertext=ciphertext,
        )

        return shared_secret, init_msg

    @staticmethod
    def respond(our_keys: X3DHKeys, init_message: X3DHInitMessage) -> Tuple[bytes, bytes]:
        """
        Respond to X3DH key exchange.

        Args:
            our_keys: Our X3DH key set
            init_message: Received init message

        Returns:
            (shared_secret, decrypted_message) tuple
        """
        _check_nacl()

        # Verify signed prekey ID matches
        if init_message.signed_prekey_id != our_keys.signed_prekey_id:
            raise ValueError("Signed prekey ID mismatch - key may have rotated")

        # Convert our Ed25519 identity to X25519 for DH
        our_identity_dh = our_keys.identity_key.to_x25519_keypair()

        # Convert their Ed25519 identity to X25519
        their_identity_x25519 = crypto_sign_ed25519_pk_to_curve25519(init_message.identity_key)

        # Perform DH calculations (mirror of initiate)
        # DH1 = DH(SPKb, IKa) - Note: same result as DH(IKa, SPKb) due to commutativity
        dh1 = our_keys.signed_prekey.dh(their_identity_x25519)

        # DH2 = DH(IKb, EKa)
        dh2 = our_identity_dh.dh(init_message.ephemeral_key)

        # DH3 = DH(SPKb, EKa)
        dh3 = our_keys.signed_prekey.dh(init_message.ephemeral_key)

        # DH4 = DH(OPKb, EKa) if one-time key was used
        dh4 = b""
        if init_message.one_time_prekey_id is not None:
            opk = our_keys.consume_one_time_key(init_message.one_time_prekey_id)
            if opk:
                dh4 = opk.dh(init_message.ephemeral_key)
            else:
                logger.warning(f"One-time prekey {init_message.one_time_prekey_id} not found")

        # Combine DH outputs
        dh_concat = dh1 + dh2 + dh3 + dh4

        # Derive shared secret
        shared_secret = hkdf(
            input_key_material=dh_concat, salt=b"\x00" * 32, info=HKDF_INFO_ROOT, length=32
        )

        # Decrypt initial message
        plaintext = b""
        if init_message.ciphertext:
            plaintext = X3DH._decrypt(shared_secret, init_message.ciphertext)

        return shared_secret, plaintext

    @staticmethod
    def _encrypt(key: bytes, plaintext: bytes) -> bytes:
        """Encrypt with key using XSalsa20-Poly1305."""
        _check_nacl()
        box = SecretBox(key)
        return bytes(box.encrypt(plaintext))

    @staticmethod
    def _decrypt(key: bytes, ciphertext: bytes) -> bytes:
        """Decrypt with key using XSalsa20-Poly1305."""
        _check_nacl()
        box = SecretBox(key)
        return bytes(box.decrypt(ciphertext))


# ============================================================
# DOUBLE RATCHET
# ============================================================


@dataclass
class SkippedKey:
    """A message key that was skipped (for out-of-order messages)."""

    key: bytes
    timestamp: float


@dataclass
class RatchetState:
    """
    Double Ratchet state for a session.

    Maintains symmetric ratchet chains and DH ratchet state.
    """

    # DH Ratchet
    dh_keypair: Optional[KeyPair] = None  # Our current ratchet key pair
    dh_remote_public: Optional[bytes] = None  # Their current ratchet public key
    root_key: Optional[bytes] = None  # Current root key

    # Sending chain
    chain_key_send: Optional[bytes] = None  # Current sending chain key
    send_count: int = 0  # Messages sent in current chain
    prev_send_count: int = 0  # Messages sent in previous sending chain

    # Receiving chain
    chain_key_recv: Optional[bytes] = None  # Current receiving chain key
    recv_count: int = 0  # Messages received in current chain

    # Previous receiving chains (for out-of-order)
    previous_chains: Dict[bytes, Tuple[bytes, int]] = field(default_factory=dict)

    # Skipped message keys
    skipped_keys: Dict[Tuple[bytes, int], SkippedKey] = field(default_factory=dict)

    # Session metadata
    session_id: str = ""
    created_at: float = 0.0
    last_activity: float = 0.0

    def __post_init__(self):
        if not self.session_id:
            self.session_id = secrets.token_hex(16)
        if not self.created_at:
            self.created_at = time.time()
        self.last_activity = time.time()

    def to_dict(self) -> dict:
        """Serialize state for storage."""
        return {
            "dh_keypair": {
                "private": self.dh_keypair.private.hex(),
                "public": self.dh_keypair.public.hex(),
            }
            if self.dh_keypair
            else None,
            "dh_remote_public": self.dh_remote_public.hex() if self.dh_remote_public else None,
            "root_key": self.root_key.hex() if self.root_key else None,
            "chain_key_send": self.chain_key_send.hex() if self.chain_key_send else None,
            "send_count": self.send_count,
            "prev_send_count": self.prev_send_count,
            "chain_key_recv": self.chain_key_recv.hex() if self.chain_key_recv else None,
            "recv_count": self.recv_count,
            "previous_chains": {
                k.hex(): (v[0].hex(), v[1]) for k, v in self.previous_chains.items()
            },
            "skipped_keys": {
                f"{k[0].hex()}:{k[1]}": {"key": v.key.hex(), "ts": v.timestamp}
                for k, v in self.skipped_keys.items()
            },
            "session_id": self.session_id,
            "created_at": self.created_at,
            "last_activity": self.last_activity,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "RatchetState":
        """Deserialize from storage."""
        dh_keypair = None
        if data.get("dh_keypair"):
            dh_keypair = KeyPair(
                private=bytes.fromhex(data["dh_keypair"]["private"]),
                public=bytes.fromhex(data["dh_keypair"]["public"]),
            )

        previous_chains = {}
        for k, v in data.get("previous_chains", {}).items():
            previous_chains[bytes.fromhex(k)] = (bytes.fromhex(v[0]), v[1])

        skipped_keys = {}
        for k, v in data.get("skipped_keys", {}).items():
            pk_hex, count = k.rsplit(":", 1)
            key_tuple = (bytes.fromhex(pk_hex), int(count))
            skipped_keys[key_tuple] = SkippedKey(key=bytes.fromhex(v["key"]), timestamp=v["ts"])

        return cls(
            dh_keypair=dh_keypair,
            dh_remote_public=bytes.fromhex(data["dh_remote_public"])
            if data.get("dh_remote_public")
            else None,
            root_key=bytes.fromhex(data["root_key"]) if data.get("root_key") else None,
            chain_key_send=bytes.fromhex(data["chain_key_send"])
            if data.get("chain_key_send")
            else None,
            send_count=data.get("send_count", 0),
            prev_send_count=data.get("prev_send_count", 0),
            chain_key_recv=bytes.fromhex(data["chain_key_recv"])
            if data.get("chain_key_recv")
            else None,
            recv_count=data.get("recv_count", 0),
            previous_chains=previous_chains,
            skipped_keys=skipped_keys,
            session_id=data.get("session_id", ""),
            created_at=data.get("created_at", 0.0),
            last_activity=data.get("last_activity", 0.0),
        )


@dataclass
class RatchetMessage:
    """
    Encrypted message with Double Ratchet header.
    """

    dh_public: bytes  # Sender's current ratchet public key
    prev_chain_len: int  # Number of messages in previous sending chain
    message_num: int  # Message number in current sending chain
    ciphertext: bytes  # Encrypted payload

    def to_bytes(self) -> bytes:
        """Serialize for transmission."""
        # Format: [dh_public(32)][prev_chain_len(4)][message_num(4)][ciphertext]
        header = struct.pack(">32sII", self.dh_public, self.prev_chain_len, self.message_num)
        return header + self.ciphertext

    @classmethod
    def from_bytes(cls, data: bytes) -> "RatchetMessage":
        """Deserialize."""
        dh_public = data[:32]
        prev_chain_len, message_num = struct.unpack(">II", data[32:40])
        ciphertext = data[40:]
        return cls(
            dh_public=dh_public,
            prev_chain_len=prev_chain_len,
            message_num=message_num,
            ciphertext=ciphertext,
        )


class DoubleRatchet:
    """
    Double Ratchet Algorithm implementation.

    Provides forward secrecy and post-compromise security for ongoing messaging.
    """

    def __init__(self, state: Optional[RatchetState] = None):
        """Initialize with optional existing state."""
        self.state = state or RatchetState()

    @classmethod
    def init_sender(cls, shared_secret: bytes, their_ratchet_key: bytes) -> "DoubleRatchet":
        """
        Initialize as sender (initiator of X3DH).

        Args:
            shared_secret: From X3DH
            their_ratchet_key: Recipient's initial ratchet public key (their signed prekey)
        """
        ratchet = cls()

        # Generate our first ratchet key pair
        ratchet.state.dh_keypair = KeyPair.generate()
        ratchet.state.dh_remote_public = their_ratchet_key

        # Perform initial DH
        dh_output = ratchet.state.dh_keypair.dh(their_ratchet_key)

        # Derive root key and sending chain key
        derived = hkdf(
            input_key_material=shared_secret + dh_output,
            salt=b"\x00" * 32,
            info=HKDF_INFO_ROOT,
            length=64,
        )
        ratchet.state.root_key = derived[:32]
        ratchet.state.chain_key_send = derived[32:]

        return ratchet

    @classmethod
    def init_receiver(cls, shared_secret: bytes, our_ratchet_keypair: KeyPair) -> "DoubleRatchet":
        """
        Initialize as receiver (responder to X3DH).

        Args:
            shared_secret: From X3DH
            our_ratchet_keypair: Our initial ratchet key pair (our signed prekey)
        """
        ratchet = cls()

        ratchet.state.dh_keypair = our_ratchet_keypair
        ratchet.state.root_key = shared_secret

        return ratchet

    def encrypt(self, plaintext: bytes) -> RatchetMessage:
        """
        Encrypt a message with the current ratchet state.

        Args:
            plaintext: Message to encrypt

        Returns:
            Encrypted ratchet message
        """
        _check_nacl()

        self.state.last_activity = time.time()

        # Derive message key from chain key
        message_key = self._kdf_chain_key(self.state.chain_key_send, b"msg")

        # Advance chain key
        self.state.chain_key_send = self._kdf_chain_key(self.state.chain_key_send, b"chain")

        # Encrypt
        nonce = nacl_random(NONCE_SIZE)
        box = SecretBox(message_key)
        ciphertext = bytes(box.encrypt(plaintext, nonce))

        # Create message
        msg = RatchetMessage(
            dh_public=self.state.dh_keypair.public,
            prev_chain_len=self.state.prev_send_count,
            message_num=self.state.send_count,
            ciphertext=ciphertext,
        )

        self.state.send_count += 1

        return msg

    def decrypt(self, message: RatchetMessage) -> bytes:
        """
        Decrypt a received ratchet message.

        Args:
            message: Received encrypted message

        Returns:
            Decrypted plaintext
        """
        _check_nacl()

        self.state.last_activity = time.time()

        # Check if we need to perform DH ratchet step
        if message.dh_public != self.state.dh_remote_public:
            self._dh_ratchet(message.dh_public)

        # Check for skipped messages
        skip_key = (message.dh_public, message.message_num)
        if skip_key in self.state.skipped_keys:
            skipped = self.state.skipped_keys.pop(skip_key)
            return self._decrypt_with_key(skipped.key, message.ciphertext)

        # Skip any messages we haven't received yet
        while self.state.recv_count < message.message_num:
            if len(self.state.skipped_keys) >= MAX_SKIP:
                raise ValueError("Too many skipped messages")

            msg_key = self._kdf_chain_key(self.state.chain_key_recv, b"msg")
            self.state.skipped_keys[(message.dh_public, self.state.recv_count)] = SkippedKey(
                key=msg_key, timestamp=time.time()
            )
            self.state.chain_key_recv = self._kdf_chain_key(self.state.chain_key_recv, b"chain")
            self.state.recv_count += 1

        # Derive message key
        message_key = self._kdf_chain_key(self.state.chain_key_recv, b"msg")

        # Advance chain key
        self.state.chain_key_recv = self._kdf_chain_key(self.state.chain_key_recv, b"chain")
        self.state.recv_count += 1

        # Decrypt
        return self._decrypt_with_key(message_key, message.ciphertext)

    def _dh_ratchet(self, their_new_public: bytes):
        """Perform DH ratchet step when we receive a new public key."""
        # Save current receiving chain for out-of-order messages
        if self.state.dh_remote_public and self.state.chain_key_recv:
            self.state.previous_chains[self.state.dh_remote_public] = (
                self.state.chain_key_recv,
                self.state.recv_count,
            )

        self.state.dh_remote_public = their_new_public

        # Derive new receiving chain
        dh_recv = self.state.dh_keypair.dh(their_new_public)
        derived = hkdf(
            input_key_material=self.state.root_key + dh_recv,
            salt=b"\x00" * 32,
            info=HKDF_INFO_ROOT,
            length=64,
        )
        self.state.root_key = derived[:32]
        self.state.chain_key_recv = derived[32:]
        self.state.recv_count = 0

        # Generate new ratchet key pair
        self.state.dh_keypair = KeyPair.generate()

        # Derive new sending chain
        dh_send = self.state.dh_keypair.dh(their_new_public)
        derived = hkdf(
            input_key_material=self.state.root_key + dh_send,
            salt=b"\x00" * 32,
            info=HKDF_INFO_ROOT,
            length=64,
        )
        self.state.root_key = derived[:32]
        self.state.chain_key_send = derived[32:]
        self.state.prev_send_count = self.state.send_count
        self.state.send_count = 0

    def _kdf_chain_key(self, chain_key: bytes, purpose: bytes) -> bytes:
        """Derive key from chain key."""
        return hkdf(
            input_key_material=chain_key,
            salt=b"\x00" * 32,
            info=HKDF_INFO_CHAIN + purpose,
            length=32,
        )

    def _decrypt_with_key(self, key: bytes, ciphertext: bytes) -> bytes:
        """Decrypt with a specific key."""
        box = SecretBox(key)
        return bytes(box.decrypt(ciphertext))

    def cleanup_skipped_keys(self):
        """Remove expired skipped keys."""
        now = time.time()
        expired = [
            k for k, v in self.state.skipped_keys.items() if now - v.timestamp > MAX_SKIP_AGE
        ]
        for k in expired:
            del self.state.skipped_keys[k]


# ============================================================
# SECURE SESSION
# ============================================================


class SecureSession:
    """
    Complete secure session combining X3DH and Double Ratchet.

    Manages the full lifecycle of a secure communication channel.
    """

    def __init__(
        self, session_id: str, local_keys: X3DHKeys, remote_identity: Optional[bytes] = None
    ):
        self.session_id = session_id
        self.local_keys = local_keys
        self.remote_identity = remote_identity
        self.ratchet: Optional[DoubleRatchet] = None
        self.established = False
        self.created_at = time.time()
        self.last_activity = time.time()

    def initiate(self, their_bundle: PreKeyBundle, initial_message: bytes = b"") -> bytes:
        """
        Initiate a secure session with a remote node.

        Args:
            their_bundle: Remote node's prekey bundle
            initial_message: Optional first message

        Returns:
            X3DH init message bytes to send
        """
        # Perform X3DH
        shared_secret, init_msg = X3DH.initiate(
            self.local_keys.identity_key, their_bundle, initial_message
        )

        # Initialize Double Ratchet as sender
        self.ratchet = DoubleRatchet.init_sender(shared_secret, their_bundle.signed_prekey)

        self.remote_identity = their_bundle.identity_key
        self.established = True
        self.last_activity = time.time()

        return init_msg.to_bytes()

    def respond(self, init_message_bytes: bytes) -> bytes:
        """
        Respond to an incoming X3DH initiation.

        Args:
            init_message_bytes: Received X3DH init message

        Returns:
            Decrypted initial message (if any)
        """
        init_msg = X3DHInitMessage.from_bytes(init_message_bytes)

        # Perform X3DH response
        shared_secret, initial_plaintext = X3DH.respond(self.local_keys, init_msg)

        # Initialize Double Ratchet as receiver
        self.ratchet = DoubleRatchet.init_receiver(shared_secret, self.local_keys.signed_prekey)

        self.remote_identity = init_msg.identity_key
        self.established = True
        self.last_activity = time.time()

        return initial_plaintext

    def encrypt(self, plaintext: bytes) -> bytes:
        """Encrypt a message."""
        if not self.established:
            raise RuntimeError("Session not established")

        msg = self.ratchet.encrypt(plaintext)
        self.last_activity = time.time()
        return msg.to_bytes()

    def decrypt(self, ciphertext_bytes: bytes) -> bytes:
        """Decrypt a message."""
        if not self.established:
            raise RuntimeError("Session not established")

        msg = RatchetMessage.from_bytes(ciphertext_bytes)
        plaintext = self.ratchet.decrypt(msg)
        self.last_activity = time.time()
        return plaintext

    def to_dict(self) -> dict:
        """Serialize session for storage."""
        return {
            "session_id": self.session_id,
            "local_keys": self.local_keys.to_dict(),
            "remote_identity": self.remote_identity.hex() if self.remote_identity else None,
            "ratchet": self.ratchet.state.to_dict() if self.ratchet else None,
            "established": self.established,
            "created_at": self.created_at,
            "last_activity": self.last_activity,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "SecureSession":
        """Deserialize from storage."""
        session = cls(
            session_id=data["session_id"],
            local_keys=X3DHKeys.from_dict(data["local_keys"]),
            remote_identity=bytes.fromhex(data["remote_identity"])
            if data.get("remote_identity")
            else None,
        )
        if data.get("ratchet"):
            session.ratchet = DoubleRatchet(RatchetState.from_dict(data["ratchet"]))
        session.established = data.get("established", False)
        session.created_at = data.get("created_at", time.time())
        session.last_activity = data.get("last_activity", time.time())
        return session


# ============================================================
# SESSION MANAGER
# ============================================================


class SecureSessionManager:
    """
    Manages multiple secure sessions with different peers.

    Handles key storage, session persistence, and lifecycle.
    """

    def __init__(self, storage_path: Optional[Path] = None):
        """
        Initialize session manager.

        Args:
            storage_path: Path to store keys and sessions (default: ~/.familiar/secure/)
        """
        _check_nacl()

        if storage_path is None:
            storage_path = Path.home() / ".familiar" / "secure"
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True, mode=0o700)

        # Load or generate our keys
        self.keys_file = self.storage_path / "keys.json"
        self.local_keys = self._load_or_generate_keys()

        # Active sessions
        self.sessions: Dict[str, SecureSession] = {}

        # Load persisted sessions
        self._load_sessions()

    def _get_key_encryption(self):
        """Get encryption storage for keys, if available."""
        try:
            enc_key = os.environ.get("FAMILIAR_ENCRYPTION_KEY")
            if enc_key:
                from .encryption import EncryptedJSONStorage

                storage = EncryptedJSONStorage()
                if storage.is_available:
                    return storage
        except (ImportError, Exception):
            pass
        return None

    def _load_or_generate_keys(self) -> X3DHKeys:
        """Load existing keys or generate new ones."""
        enc = self._get_key_encryption()
        enc_file = self.keys_file.with_suffix(".json.enc")

        # Try encrypted file first
        if enc and enc_file.exists():
            try:
                encrypted_data = enc_file.read_text()
                decrypted = enc.decrypt(encrypted_data)
                data = json.loads(decrypted)
                return X3DHKeys.from_dict(data)
            except Exception as e:
                logger.error(f"Failed to load encrypted keys: {e}")

        # Fall back to plaintext file
        if self.keys_file.exists():
            try:
                with open(self.keys_file, "r") as f:
                    data = json.load(f)
                keys = X3DHKeys.from_dict(data)
                # Migrate to encrypted storage if available
                if enc:
                    self._save_keys(keys)
                    self.keys_file.unlink()
                    logger.info("Migrated keys to encrypted storage")
                return keys
            except Exception as e:
                logger.error(f"Failed to load keys: {e}")

        # Generate new keys
        keys = X3DHKeys.generate()
        self._save_keys(keys)
        return keys

    def _save_keys(self, keys: X3DHKeys):
        """Save keys to storage, encrypted if possible."""
        enc = self._get_key_encryption()
        key_json = json.dumps(keys.to_dict())

        if enc:
            enc_file = self.keys_file.with_suffix(".json.enc")
            encrypted = enc.encrypt(key_json)
            enc_file.write_text(encrypted)
            os.chmod(enc_file, 0o600)
        else:
            with open(self.keys_file, "w") as f:
                f.write(key_json)
            os.chmod(self.keys_file, 0o600)

    def _load_sessions(self):
        """Load persisted sessions."""
        sessions_file = self.storage_path / "sessions.json"
        if sessions_file.exists():
            try:
                with open(sessions_file, "r") as f:
                    data = json.load(f)
                for session_data in data.get("sessions", []):
                    session = SecureSession.from_dict(session_data)
                    self.sessions[session.session_id] = session
            except Exception as e:
                logger.error(f"Failed to load sessions: {e}")

    def _save_sessions(self):
        """Persist sessions to storage."""
        sessions_file = self.storage_path / "sessions.json"
        data = {"sessions": [s.to_dict() for s in self.sessions.values()]}
        with open(sessions_file, "w") as f:
            json.dump(data, f)
        os.chmod(sessions_file, 0o600)

    def get_prekey_bundle(self) -> PreKeyBundle:
        """Get our prekey bundle for distribution."""
        return self.local_keys.get_bundle()

    def create_session(
        self, peer_id: str, their_bundle: PreKeyBundle
    ) -> Tuple[SecureSession, bytes]:
        """
        Create a new session with a peer (as initiator).

        Args:
            peer_id: Identifier for the peer
            their_bundle: Peer's prekey bundle

        Returns:
            (session, init_message_bytes) tuple
        """
        session = SecureSession(
            session_id=peer_id,
            local_keys=self.local_keys,
        )

        init_bytes = session.initiate(their_bundle)

        self.sessions[peer_id] = session
        self._save_sessions()

        return session, init_bytes

    def accept_session(
        self, peer_id: str, init_message_bytes: bytes
    ) -> Tuple[SecureSession, bytes]:
        """
        Accept an incoming session (as responder).

        Args:
            peer_id: Identifier for the peer
            init_message_bytes: Received X3DH init message

        Returns:
            (session, decrypted_initial_message) tuple
        """
        session = SecureSession(
            session_id=peer_id,
            local_keys=self.local_keys,
        )

        initial_plaintext = session.respond(init_message_bytes)

        self.sessions[peer_id] = session
        self._save_sessions()

        return session, initial_plaintext

    def get_session(self, peer_id: str) -> Optional[SecureSession]:
        """Get an existing session."""
        return self.sessions.get(peer_id)

    def encrypt(self, peer_id: str, plaintext: bytes) -> bytes:
        """Encrypt a message for a peer."""
        session = self.sessions.get(peer_id)
        if not session:
            raise ValueError(f"No session for peer: {peer_id}")

        ciphertext = session.encrypt(plaintext)
        self._save_sessions()  # Save updated ratchet state
        return ciphertext

    def decrypt(self, peer_id: str, ciphertext: bytes) -> bytes:
        """Decrypt a message from a peer."""
        session = self.sessions.get(peer_id)
        if not session:
            raise ValueError(f"No session for peer: {peer_id}")

        plaintext = session.decrypt(ciphertext)
        self._save_sessions()  # Save updated ratchet state
        return plaintext

    def rotate_prekey(self):
        """Rotate our signed prekey (call periodically)."""
        self.local_keys.rotate_signed_prekey()
        self._save_keys(self.local_keys)

    def generate_one_time_keys(self, count: int = 10):
        """Generate additional one-time prekeys."""
        self.local_keys.generate_one_time_keys(count)
        self._save_keys(self.local_keys)


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    # Constants
    "PROTOCOL_VERSION",
    "MessageType",
    # Key types
    "KeyPair",
    "SigningKeyPair",
    "PreKeyBundle",
    "X3DHKeys",
    "X3DHInitMessage",
    # Protocols
    "X3DH",
    "DoubleRatchet",
    "RatchetState",
    "RatchetMessage",
    # Session management
    "SecureSession",
    "SecureSessionManager",
    # Utilities
    "hkdf",
    "hkdf_extract",
    "hkdf_expand",
]
