from __future__ import annotations

import os

from mrp.auth import _b64url_decode, _b64url_encode

# These imports are optional - E2E encryption requires PyNaCl
try:
    import nacl.bindings
    _HAS_NACL = True
except ImportError:
    _HAS_NACL = False

from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes


_HKDF_INFO = b"m2m-v1-msg"
_XCHACHA_NONCE_SIZE = 24
_XCHACHA_KEY_SIZE = 32


def _ensure_nacl() -> None:
    if not _HAS_NACL:
        raise ImportError(
            "PyNaCl is required for E2E encryption. Install with: pip install mrp-sdk[e2e]"
        )


def ed25519_to_x25519_private(ed_seed: bytes) -> X25519PrivateKey:
    """Convert Ed25519 seed (32 bytes) to X25519 private key using nacl."""
    _ensure_nacl()
    # Use libsodium's conversion
    x_priv_bytes = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(ed_seed + b"\x00" * 32)
    return X25519PrivateKey.from_private_bytes(x_priv_bytes[:32])


def ed25519_to_x25519_public(ed_pub: bytes) -> X25519PublicKey:
    """Convert Ed25519 public key (32 bytes) to X25519 public key using nacl."""
    _ensure_nacl()
    x_pub_bytes = nacl.bindings.crypto_sign_ed25519_pk_to_curve25519(ed_pub)
    return X25519PublicKey.from_public_bytes(x_pub_bytes)


def encrypt_message(
    sender_seed: bytes,
    recipient_public_key: bytes,
    plaintext: bytes,
    content_type: str = "application/json",
) -> dict:
    """Encrypt a message for a recipient.

    Args:
        sender_seed: 32-byte Ed25519 seed of the sender
        recipient_public_key: 32-byte Ed25519 public key of the recipient
        plaintext: message bytes to encrypt
        content_type: original content type

    Returns:
        dict with fields: v, eph, nonce, ct, ct_content_type
    """
    _ensure_nacl()

    # Convert recipient's Ed25519 public key to X25519
    recipient_x25519 = ed25519_to_x25519_public(recipient_public_key)

    # Generate ephemeral X25519 keypair
    ephemeral_private = X25519PrivateKey.generate()
    ephemeral_public = ephemeral_private.public_key()
    ephemeral_public_bytes = ephemeral_public.public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw
    )

    # X25519 key agreement
    shared_secret = ephemeral_private.exchange(recipient_x25519)

    # HKDF-SHA256
    derived_key = HKDF(
        algorithm=hashes.SHA256(),
        length=_XCHACHA_KEY_SIZE,
        salt=ephemeral_public_bytes,
        info=_HKDF_INFO,
    ).derive(shared_secret)

    # XChaCha20-Poly1305 encrypt
    nonce = os.urandom(_XCHACHA_NONCE_SIZE)
    ciphertext = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_encrypt(
        plaintext, None, nonce, derived_key
    )

    return {
        "v": 1,
        "eph": _b64url_encode(ephemeral_public_bytes),
        "nonce": _b64url_encode(nonce),
        "ct": _b64url_encode(ciphertext),
        "ct_content_type": content_type,
    }


def decrypt_message(
    recipient_seed: bytes,
    sender_public_key: bytes,
    encrypted_body: dict,
) -> tuple[bytes, str]:
    """Decrypt an encrypted message.

    Args:
        recipient_seed: 32-byte Ed25519 seed of the recipient
        sender_public_key: 32-byte Ed25519 public key of the sender (unused in current scheme, reserved)
        encrypted_body: dict with v, eph, nonce, ct, ct_content_type fields

    Returns:
        (plaintext, original_content_type)
    """
    _ensure_nacl()

    if encrypted_body.get("v") != 1:
        raise ValueError(f"Unsupported encryption version: {encrypted_body.get('v')}")

    # Decode fields
    ephemeral_public_bytes = _b64url_decode(encrypted_body["eph"])
    nonce = _b64url_decode(encrypted_body["nonce"])
    ciphertext = _b64url_decode(encrypted_body["ct"])
    content_type = encrypted_body["ct_content_type"]

    # Convert recipient's Ed25519 key to X25519
    recipient_x25519 = ed25519_to_x25519_private(recipient_seed)

    # Reconstruct ephemeral public key
    ephemeral_public = X25519PublicKey.from_public_bytes(ephemeral_public_bytes)

    # X25519 key agreement
    shared_secret = recipient_x25519.exchange(ephemeral_public)

    # HKDF-SHA256
    derived_key = HKDF(
        algorithm=hashes.SHA256(),
        length=_XCHACHA_KEY_SIZE,
        salt=ephemeral_public_bytes,
        info=_HKDF_INFO,
    ).derive(shared_secret)

    # XChaCha20-Poly1305 decrypt
    plaintext = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_decrypt(
        ciphertext, None, nonce, derived_key
    )

    return plaintext, content_type
