"""
Message Signing — Ed25519 signatures for room messages.

Every message sent through the room WebSocket is signed so it cannot be
spoofed.  Key pairs are generated on first run and stored at
``~/.brynq/keys/``.

Public keys are exchanged between room members on join so each peer can
verify the other's messages.  If verification fails the message is dropped
with a warning.
"""

from __future__ import annotations

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

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

KEYS_DIR = Path.home() / ".brynq" / "keys"
PRIVATE_KEY_FILE = KEYS_DIR / "ed25519.key"
PUBLIC_KEY_FILE = KEYS_DIR / "ed25519.pub"

# We use the ``cryptography`` library which is already a transitive
# dependency via PyJWT[crypto].
try:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import (
        Ed25519PrivateKey,
        Ed25519PublicKey,
    )
    from cryptography.hazmat.primitives import serialization
    from cryptography.exceptions import InvalidSignature

    _HAS_CRYPTO = True
except ImportError:
    _HAS_CRYPTO = False


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


class MessageSigner:
    """Ed25519 message signing and verification."""

    def __init__(self, keys_dir: Path | str | None = None):
        _ensure_crypto()
        self._keys_dir = Path(keys_dir) if keys_dir else KEYS_DIR
        self._private_key: Ed25519PrivateKey | None = None
        self._public_key: Ed25519PublicKey | None = None
        self._peer_keys: dict[str, Ed25519PublicKey] = {}  # peer_id -> pub key
        self._load_or_generate()

    # ── key management ────────────────────────────────────────────────

    def _load_or_generate(self) -> None:
        """Load existing key pair or generate a new one."""
        priv_path = self._keys_dir / "ed25519.key"
        pub_path = self._keys_dir / "ed25519.pub"

        if priv_path.exists() and pub_path.exists():
            priv_bytes = priv_path.read_bytes()
            self._private_key = serialization.load_pem_private_key(
                priv_bytes, password=None
            )
            self._public_key = self._private_key.public_key()
            logger.info("Loaded existing Ed25519 key pair from %s", self._keys_dir)
        else:
            self._generate_keys()

    def _generate_keys(self) -> None:
        """Generate a fresh Ed25519 key pair and persist it."""
        self._keys_dir.mkdir(parents=True, exist_ok=True)

        self._private_key = Ed25519PrivateKey.generate()
        self._public_key = self._private_key.public_key()

        priv_path = self._keys_dir / "ed25519.key"
        pub_path = self._keys_dir / "ed25519.pub"

        priv_bytes = self._private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        pub_bytes = self._public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        priv_path.write_bytes(priv_bytes)
        pub_path.write_bytes(pub_bytes)

        # Restrict private key permissions on Unix
        try:
            os.chmod(str(priv_path), 0o600)
        except OSError:
            pass  # Windows doesn't support chmod the same way

        logger.info("Generated new Ed25519 key pair at %s", self._keys_dir)

    def get_public_key_pem(self) -> str:
        """Return the local public key in PEM format (for sharing with peers)."""
        return self._public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("utf-8")

    def register_peer_key(self, peer_id: str, public_key_pem: str) -> None:
        """Register a peer's public key for future verification."""
        pub_key = serialization.load_pem_public_key(public_key_pem.encode("utf-8"))
        self._peer_keys[peer_id] = pub_key
        logger.info("Registered public key for peer %s", peer_id)

    def remove_peer_key(self, peer_id: str) -> None:
        """Remove a peer's public key."""
        self._peer_keys.pop(peer_id, None)

    # ── sign / verify ─────────────────────────────────────────────────

    def sign_message(self, message: dict) -> dict:
        """Sign a message dict, adding ``_signature`` and ``_signed_payload`` fields.

        The signature covers a canonical JSON serialization of the message
        (excluding any existing signature fields).
        """
        # Strip existing signature fields before signing
        payload = {k: v for k, v in message.items() if not k.startswith("_sig")}
        canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        sig_bytes = self._private_key.sign(canonical.encode("utf-8"))

        signed = dict(message)
        signed["_signature"] = base64.b64encode(sig_bytes).decode("ascii")
        signed["_signed_payload"] = canonical
        return signed

    def verify_message(self, message: dict, sender_id: str) -> bool:
        """Verify the signature on a received message.

        Returns ``True`` if valid, ``False`` if the signature is missing,
        the sender's key is unknown, or verification fails.
        """
        sig_b64 = message.get("_signature")
        canonical = message.get("_signed_payload")

        if not sig_b64 or not canonical:
            logger.warning("Message from %s has no signature — dropping", sender_id)
            return False

        pub_key = self._peer_keys.get(sender_id)
        if pub_key is None:
            logger.warning("No public key registered for peer %s — dropping", sender_id)
            return False

        try:
            sig_bytes = base64.b64decode(sig_b64)
            pub_key.verify(sig_bytes, canonical.encode("utf-8"))
            return True
        except InvalidSignature:
            logger.warning("Invalid signature from peer %s — message dropped", sender_id)
            return False
        except Exception as exc:
            logger.error("Signature verification error for %s: %s", sender_id, exc)
            return False
