"""
communication/crypto.py - Cryptographic Primitives
===================================================
Ed25519 identity, X25519 key exchange, AES-256-GCM encryption/decryption,
message signing/verification, and key serialization.

Uses the ``cryptography`` library.  Falls back to a pure-Python stub when it
is not available (signing only, encryption disabled).
"""
from __future__ import annotations

import hashlib
import hmac
import os
import base64
import json
from typing import Optional, Tuple

# ---------------------------------------------------------------------------
# Try importing the real crypto library
# ---------------------------------------------------------------------------
_HAS_CRYPTOGRAPHY = False
try:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import (
        Ed25519PrivateKey,
        Ed25519PublicKey,
    )
    from cryptography.hazmat.primitives.asymmetric.x25519 import (
        X25519PrivateKey,
        X25519PublicKey,
    )
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
    from cryptography.hazmat.primitives import serialization
    from cryptography.exceptions import InvalidSignature

    _HAS_CRYPTOGRAPHY = True
except ImportError:
    pass

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
NONCE_SIZE = 12          # 96-bit nonce for AES-256-GCM
SHARED_SECRET_SIZE = 32  # 256-bit shared secret


# =========================================================================
# Key Generation
# =========================================================================

def generate_ed25519_keypair() -> Tuple[bytes, bytes]:
    """Generate an Ed25519 key pair.

    Returns:
        (private_bytes, public_bytes) — both 32 bytes raw.
    """
    if _HAS_CRYPTOGRAPHY:
        private_key = Ed25519PrivateKey.generate()
        priv_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PrivateFormat.Raw,
            encryption_algorithm=serialization.NoEncryption(),
        )
        pub_bytes = private_key.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
        return priv_bytes, pub_bytes
    else:
        # Fallback: generate random bytes (signing will not be available)
        priv_bytes = os.urandom(32)
        pub_bytes = hashlib.sha256(priv_bytes).digest()
        return priv_bytes, pub_bytes


def generate_x25519_keypair() -> Tuple[bytes, bytes]:
    """Generate an X25519 key pair for Diffie-Hellman.

    Returns:
        (private_bytes, public_bytes) — both 32 bytes raw.
    """
    if _HAS_CRYPTOGRAPHY:
        private_key = X25519PrivateKey.generate()
        priv_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PrivateFormat.Raw,
            encryption_algorithm=serialization.NoEncryption(),
        )
        pub_bytes = private_key.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
        return priv_bytes, pub_bytes
    else:
        priv_bytes = os.urandom(32)
        pub_bytes = hashlib.sha256(priv_bytes).digest()
        return priv_bytes, pub_bytes


# =========================================================================
# Ed25519 <-> X25519 Conversion
# =========================================================================

def ed25519_to_x25519_private(ed_priv: bytes) -> bytes:
    """Convert an Ed25519 private key to X25519 private key bytes.

    According to RFC 7748, we hash the 32-byte Ed25519 seed with SHA-512
    and clamp the first 32 bytes.
    """
    if _HAS_CRYPTOGRAPHY:
        # The raw Ed25519 private key from cryptography is already the seed.
        # Use the standard clamping procedure.
        h = hashlib.sha512(ed_priv).digest()
        clamped = bytearray(h[:32])
        clamped[0] &= 248
        clamped[31] &= 127
        clamped[31] |= 64
        return bytes(clamped)
    else:
        return hashlib.sha256(ed_priv).digest()


def ed25519_to_x25519_public(ed_pub: bytes) -> bytes:
    """Convert an Ed25519 public key to X25519 public key.

    In practice with the cryptography library, Ed25519 and X25519 public
    keys share the same representation for curve25519.  We return a simple
    hash fallback when the library is not available.
    """
    if _HAS_CRYPTOGRAPHY:
        # Both use the same Montgomery u-coordinate internally.
        return ed_pub
    else:
        return hashlib.sha256(ed_pub).digest()


# =========================================================================
# X25519 Diffie-Hellman Key Exchange
# =========================================================================

def x25519_key_exchange(
    my_private: bytes,
    their_public: bytes,
) -> bytes:
    """Perform X25519 Diffie-Hellman key exchange.

    Returns:
        32-byte shared secret.
    """
    if _HAS_CRYPTOGRAPHY:
        my_priv_key = X25519PrivateKey.from_private_bytes(my_private)
        their_pub_key = X25519PublicKey.from_public_bytes(their_public)
        shared = my_priv_key.exchange(their_pub_key)
        return shared
    else:
        # Fallback: HMAC-based pseudo-DH (NOT secure, for testing only)
        return hmac.new(my_private, their_public, hashlib.sha256).digest()


# =========================================================================
# HKDF — Derive Encryption Key from Shared Secret
# =========================================================================

def derive_shared_key(shared_secret: bytes, context: bytes = b"") -> bytes:
    """Derive a 256-bit AES key from a raw shared secret using HKDF-SHA256.

    Args:
        shared_secret: Raw DH shared secret (32 bytes).
        context: Optional context string for domain separation.

    Returns:
        32-byte key suitable for AES-256-GCM.
    """
    if _HAS_CRYPTOGRAPHY:
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b"myagent-comm-v1" + context,
        )
        return hkdf.derive(shared_secret)
    else:
        # Fallback: simple hash-based derivation
        derived = hashlib.sha256(shared_secret + b"myagent-comm-v1" + context).digest()
        return derived


# =========================================================================
# AES-256-GCM Encrypt / Decrypt
# =========================================================================

def encrypt(key: bytes, plaintext: bytes) -> bytes:
    """Encrypt *plaintext* with AES-256-GCM.

    Args:
        key: 32-byte encryption key.
        plaintext: Arbitrary data to encrypt.

    Returns:
        ``nonce (12) || ciphertext+tag`` — concatenation of nonce and sealed data.
    """
    if _HAS_CRYPTOGRAPHY:
        nonce = os.urandom(NONCE_SIZE)
        aesgcm = AESGCM(key)
        ct = aesgcm.encrypt(nonce, plaintext, None)
        return nonce + ct
    else:
        # Fallback: XOR-based "encryption" (NOT secure)
        nonce = os.urandom(NONCE_SIZE)
        stream = b""
        for i, b in enumerate(plaintext):
            stream += bytes([b ^ key[i % len(key)]])
        return nonce + stream


def decrypt(key: bytes, sealed: bytes) -> bytes:
    """Decrypt data produced by :func:`encrypt`.

    Args:
        key: 32-byte encryption key.
        sealed: ``nonce (12) || ciphertext+tag``.

    Returns:
        Original plaintext.

    Raises:
        ValueError: If decryption fails (tampered data).
    """
    if _HAS_CRYPTOGRAPHY:
        nonce = sealed[:NONCE_SIZE]
        ct = sealed[NONCE_SIZE:]
        aesgcm = AESGCM(key)
        try:
            return aesgcm.decrypt(nonce, ct, None)
        except Exception:
            raise ValueError("Decryption failed: data may be tampered")
    else:
        nonce = sealed[:NONCE_SIZE]
        ct = sealed[NONCE_SIZE:]
        stream = b""
        for i, b in enumerate(ct):
            stream += bytes([b ^ key[i % len(key)]])
        return stream


# =========================================================================
# Ed25519 Sign / Verify
# =========================================================================

def sign(private_key: bytes, message: bytes) -> bytes:
    """Sign *message* with an Ed25519 *private_key*.

    Returns:
        64-byte signature.
    """
    if _HAS_CRYPTOGRAPHY:
        priv = Ed25519PrivateKey.from_private_bytes(private_key)
        return priv.sign(message)
    else:
        # Fallback: HMAC-SHA512 (NOT a real Ed25519 signature)
        return hmac.new(private_key, message, hashlib.sha512).digest()


def verify(public_key: bytes, message: bytes, signature: bytes) -> bool:
    """Verify an Ed25519 *signature* for *message* with *public_key*."""
    if _HAS_CRYPTOGRAPHY:
        try:
            pub = Ed25519PublicKey.from_public_bytes(public_key)
            pub.verify(signature, message)
            return True
        except InvalidSignature:
            return False
        except Exception:
            return False
    else:
        expected = hmac.new(public_key, message, hashlib.sha512).digest()
        return hmac.compare_digest(expected, signature)


# =========================================================================
# Key Serialization (Hex)
# =========================================================================

def serialize_public_key(pub: bytes) -> str:
    """Serialize a 32-byte public key to a hex string."""
    return pub.hex()


def deserialize_public_key(hex_str: str) -> bytes:
    """Deserialize a hex public key to 32 bytes."""
    return bytes.fromhex(hex_str)


def serialize_private_key(priv: bytes) -> str:
    """Serialize a 32-byte private key to a hex string."""
    return priv.hex()


def deserialize_private_key(hex_str: str) -> bytes:
    """Deserialize a hex private key to 32 bytes."""
    return bytes.fromhex(hex_str)


# =========================================================================
# Agent ID Generation
# =========================================================================

def generate_agent_id(public_key: bytes) -> str:
    """Generate a short, human-readable Agent ID from a 32-byte Ed25519 public key.

    Uses the first 8 bytes (16 hex chars) of SHA-256(public_key).

    Returns:
        A 16-char lowercase hex string, e.g. ``"a1b2c3d4e5f6a7b8"``.
    """
    digest = hashlib.sha256(public_key).hexdigest()
    return digest[:16]


# =========================================================================
# Convenience: encrypt a text message payload
# =========================================================================

def encrypt_message(
    shared_key: bytes,
    message_data: dict,
) -> str:
    """Encrypt a JSON-serializable dict to a base64 sealed string.

    This is the high-level helper used by the channel layer.
    """
    plaintext = json.dumps(message_data, ensure_ascii=False).encode("utf-8")
    sealed = encrypt(shared_key, plaintext)
    return base64.b64encode(sealed).decode("ascii")


def decrypt_message(
    shared_key: bytes,
    sealed_b64: str,
) -> dict:
    """Decrypt a base64 sealed string back to a dict."""
    sealed = base64.b64decode(sealed_b64)
    plaintext = decrypt(shared_key, sealed)
    return json.loads(plaintext.decode("utf-8"))