"""
End-to-End Encryption — X25519 key exchange + AES-256-GCM for room messages.

Room messages are encrypted so the cloud relay sees only opaque blobs.
On room join, peers exchange X25519 public keys via the cloud signaling
channel.  A shared secret is derived via Diffie-Hellman and used to
encrypt/decrypt all subsequent messages with AES-256-GCM.

Requires the ``cryptography`` library (already a transitive dep via
PyJWT[crypto]).
"""

from __future__ import annotations

import base64
import json
import logging
import os
from pathlib import Path
from typing import Optional

logger = logging.getLogger("brynq.security.encryption")

try:
    from cryptography.hazmat.primitives.asymmetric.x25519 import (
        X25519PrivateKey,
        X25519PublicKey,
    )
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF

    _HAS_CRYPTO = True
except ImportError:
    _HAS_CRYPTO = False


def _ensure_crypto() -> None:
    if not _HAS_CRYPTO:
        raise RuntimeError(
            "cryptography library is required for E2E encryption. "
            "Install it with: pip install cryptography"
        )


class E2EEncryption:
    """X25519 + AES-256-GCM end-to-end encryption for room messages."""

    def __init__(self):
        _ensure_crypto()
        self._private_key: X25519PrivateKey = X25519PrivateKey.generate()
        self._public_key: X25519PublicKey = self._private_key.public_key()
        self._shared_secrets: dict[str, bytes] = {}  # peer_id -> 32-byte key

    # ── key exchange ──────────────────────────────────────────────────

    def get_public_key_bytes(self) -> bytes:
        """Return the raw X25519 public key (32 bytes) for sharing."""
        return self._public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )

    def get_public_key_b64(self) -> str:
        """Return the public key as a base64 string (for JSON transport)."""
        return base64.b64encode(self.get_public_key_bytes()).decode("ascii")

    def derive_shared_secret(self, peer_id: str, peer_public_key_b64: str) -> None:
        """Derive a shared AES key from the peer's X25519 public key.

        Uses HKDF-SHA256 to derive a 32-byte symmetric key from the raw
        Diffie-Hellman shared secret.
        """
        peer_pub_bytes = base64.b64decode(peer_public_key_b64)
        peer_pub_key = X25519PublicKey.from_public_bytes(peer_pub_bytes)
        raw_shared = self._private_key.exchange(peer_pub_key)

        # Derive a proper AES-256 key via HKDF
        derived = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b"brynq-room-e2e-v1",
        ).derive(raw_shared)

        self._shared_secrets[peer_id] = derived
        logger.info("Derived shared secret with peer %s", peer_id)

    def has_shared_secret(self, peer_id: str) -> bool:
        return peer_id in self._shared_secrets

    def remove_peer(self, peer_id: str) -> None:
        self._shared_secrets.pop(peer_id, None)

    # ── encrypt / decrypt ─────────────────────────────────────────────

    def encrypt_message(self, message: dict, peer_id: str) -> dict:
        """Encrypt a message dict for *peer_id*.

        Returns a wrapper dict with ``type=encrypted``, the ciphertext
        (base64), the nonce (base64), and the original message type so
        the receiver knows how to route it after decryption.
        """
        key = self._shared_secrets.get(peer_id)
        if key is None:
            raise ValueError(f"No shared secret with peer {peer_id}")

        plaintext = json.dumps(message, separators=(",", ":")).encode("utf-8")
        nonce = os.urandom(12)  # 96-bit nonce for AES-GCM
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)

        return {
            "type": "encrypted",
            "original_type": message.get("type", "unknown"),
            "ciphertext": base64.b64encode(ciphertext).decode("ascii"),
            "nonce": base64.b64encode(nonce).decode("ascii"),
            "sender_id": message.get("sender_id", ""),
        }

    def decrypt_message(self, envelope: dict, peer_id: str) -> dict | None:
        """Decrypt an encrypted envelope from *peer_id*.

        Returns the original message dict, or ``None`` if decryption fails.
        """
        key = self._shared_secrets.get(peer_id)
        if key is None:
            logger.warning("No shared secret for peer %s — cannot decrypt", peer_id)
            return None

        try:
            ciphertext = base64.b64decode(envelope["ciphertext"])
            nonce = base64.b64decode(envelope["nonce"])
            aesgcm = AESGCM(key)
            plaintext = aesgcm.decrypt(nonce, ciphertext, None)
            return json.loads(plaintext.decode("utf-8"))
        except Exception as exc:
            logger.warning("Decryption failed for peer %s: %s", peer_id, exc)
            return None

    # ── broadcast helpers ─────────────────────────────────────────────

    def encrypt_for_all(self, message: dict) -> dict[str, dict]:
        """Encrypt a message for every peer with a shared secret.

        Returns ``{peer_id: encrypted_envelope}``.
        """
        result = {}
        for peer_id in self._shared_secrets:
            result[peer_id] = self.encrypt_message(message, peer_id)
        return result
