from __future__ import annotations

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

# HPKE support requires pyhpke
try:
    from pyhpke import CipherSuite as HPKESuite, KEMId, KDFId, AEADId
    _HAS_HPKE = True
except ImportError:
    _HAS_HPKE = False

from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives import serialization


_HPKE_INFO_V2 = b"mrp-v2-msg"


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


def _ensure_hpke() -> None:
    if not _HAS_HPKE:
        raise ImportError(
            "pyhpke is required for HPKE encryption. Install with: pip install mrp-sdk[e2e]"
        )


def _get_hpke_suite():
    _ensure_hpke()
    return HPKESuite.new(
        KEMId.DHKEM_X25519_HKDF_SHA256,
        KDFId.HKDF_SHA256,
        AEADId.CHACHA20_POLY1305,
    )


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 using HPKE Auth mode (v2).

    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, enc, ct, ct_content_type
    """
    _ensure_nacl()
    _ensure_hpke()

    # Convert Ed25519 keys to X25519 raw bytes
    recipient_x25519 = ed25519_to_x25519_public(recipient_public_key)
    recipient_x25519_bytes = recipient_x25519.public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw
    )
    sender_x25519 = ed25519_to_x25519_private(sender_seed)
    sender_x25519_bytes = sender_x25519.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption()
    )

    suite = _get_hpke_suite()
    pkr = suite.kem.deserialize_public_key(recipient_x25519_bytes)
    sks = suite.kem.deserialize_private_key(sender_x25519_bytes)

    enc, ctx = suite.create_sender_context(pkr=pkr, info=_HPKE_INFO_V2, sks=sks)
    ct = ctx.seal(plaintext)

    return {
        "v": 2,
        "enc": _b64url_encode(enc),
        "ct": _b64url_encode(ct),
        "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
        encrypted_body: dict with encryption fields

    Returns:
        (plaintext, original_content_type)
    """
    v = encrypted_body.get("v")
    if v != 2:
        raise ValueError(f"Unsupported encryption version: {v}")
    return _decrypt_v2(recipient_seed, sender_public_key, encrypted_body)


def _decrypt_v2(
    recipient_seed: bytes,
    sender_public_key: bytes,
    encrypted_body: dict,
) -> tuple[bytes, str]:
    """Decrypt a v2 encrypted message (HPKE Auth mode)."""
    _ensure_nacl()
    _ensure_hpke()

    enc = _b64url_decode(encrypted_body["enc"])
    ct = _b64url_decode(encrypted_body["ct"])
    content_type = encrypted_body["ct_content_type"]

    recipient_x25519 = ed25519_to_x25519_private(recipient_seed)
    recipient_x25519_bytes = recipient_x25519.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw, serialization.NoEncryption()
    )
    sender_x25519 = ed25519_to_x25519_public(sender_public_key)
    sender_x25519_bytes = sender_x25519.public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw
    )

    suite = _get_hpke_suite()
    skr = suite.kem.deserialize_private_key(recipient_x25519_bytes)
    pks = suite.kem.deserialize_public_key(sender_x25519_bytes)

    ctx = suite.create_recipient_context(enc=enc, skr=skr, info=_HPKE_INFO_V2, pks=pks)
    plaintext = ctx.open(ct)

    return plaintext, content_type
