"""
AgentDrop Encryption - X25519 key exchange + HKDF + AES-256-GCM

Zero-knowledge encryption: files are encrypted client-side before upload.
The server never sees plaintext.
"""

import os
import base64
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.serialization import (
    Encoding, PublicFormat, PrivateFormat, NoEncryption,
)
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.hashes import SHA256
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


ALGORITHM = "X25519-HKDF-AES-256-GCM"
CHANNEL_ALGORITHM = "X25519-CHANNEL-HKDF-AES-256-GCM"
KEY_BYTES = 32  # AES-256
NONCE_BYTES = 12  # GCM standard
SALT_BYTES = 32  # Per-transfer salt for channel encryption
HKDF_INFO = b"agentdrop-file-transfer-v1"
HKDF_CHANNEL_INFO = b"agentdrop-transfer-v1"


class KeyPair:
    """X25519 key pair for encryption."""

    def __init__(self, private_key: X25519PrivateKey):
        self._private = private_key
        self._public = private_key.public_key()

    @classmethod
    def generate(cls) -> "KeyPair":
        return cls(X25519PrivateKey.generate())

    @classmethod
    def from_private_bytes(cls, b64: str) -> "KeyPair":
        raw = base64.b64decode(b64)
        return cls(X25519PrivateKey.from_private_bytes(raw))

    @property
    def public_key_b64(self) -> str:
        raw = self._public.public_bytes(Encoding.Raw, PublicFormat.Raw)
        return base64.b64encode(raw).decode()

    @property
    def private_key_b64(self) -> str:
        raw = self._private.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
        return base64.b64encode(raw).decode()

    def derive_shared_key(self, recipient_public_b64: str) -> bytes:
        """X25519 key exchange + HKDF to derive AES-256 key."""
        recipient_raw = base64.b64decode(recipient_public_b64)
        recipient_pub = X25519PublicKey.from_public_bytes(recipient_raw)
        shared_secret = self._private.exchange(recipient_pub)

        # HKDF to derive a proper AES key from the shared secret
        derived = HKDF(
            algorithm=SHA256(),
            length=KEY_BYTES,
            salt=None,
            info=HKDF_INFO,
        ).derive(shared_secret)

        return derived


def encrypt_file(data: bytes, shared_key: bytes) -> tuple[bytes, bytes]:
    """
    Encrypt file data with AES-256-GCM.

    Returns (ciphertext, nonce).
    The nonce must be stored per-file for decryption.
    """
    nonce = os.urandom(NONCE_BYTES)
    aesgcm = AESGCM(shared_key)
    ciphertext = aesgcm.encrypt(nonce, data, None)
    return ciphertext, nonce


def decrypt_file(ciphertext: bytes, shared_key: bytes, nonce: bytes) -> bytes:
    """Decrypt file data with AES-256-GCM."""
    aesgcm = AESGCM(shared_key)
    return aesgcm.decrypt(nonce, ciphertext, None)


def encrypt_symmetric_key(symmetric_key: bytes, sender_keys: KeyPair, recipient_public_b64: str) -> str:
    """
    Encrypt the symmetric AES key for the recipient using X25519 + AES-256-GCM.

    This is the 'encrypted_key' field sent with the transfer.
    The recipient derives the same shared key from their private key + sender's public key.
    """
    shared_key = sender_keys.derive_shared_key(recipient_public_b64)
    nonce = os.urandom(NONCE_BYTES)
    aesgcm = AESGCM(shared_key)
    encrypted = aesgcm.encrypt(nonce, symmetric_key, None)
    # Pack nonce + ciphertext together
    return base64.b64encode(nonce + encrypted).decode()


def decrypt_symmetric_key(encrypted_key_b64: str, recipient_keys: KeyPair, sender_public_b64: str) -> bytes:
    """
    Decrypt the symmetric AES key using X25519 key exchange.

    The recipient uses their private key + sender's public key to derive the shared key.
    """
    raw = base64.b64decode(encrypted_key_b64)
    nonce = raw[:NONCE_BYTES]
    ciphertext = raw[NONCE_BYTES:]
    shared_key = recipient_keys.derive_shared_key(sender_public_b64)
    aesgcm = AESGCM(shared_key)
    return aesgcm.decrypt(nonce, ciphertext, None)


# ─── Pairwise Channel Encryption ────────────────────────────


def derive_shared_secret(my_private_key_b64: str, their_public_key_b64: str) -> bytes:
    """
    Compute the raw X25519 Diffie-Hellman shared secret.

    This is the pairwise secret between two agents that persists across
    transfers on the same channel. It MUST NOT be used directly as an
    encryption key - pass it through derive_transfer_key() first.

    Args:
        my_private_key_b64: Base64-encoded X25519 private key.
        their_public_key_b64: Base64-encoded X25519 public key of the peer.

    Returns:
        32-byte raw shared secret.
    """
    my_private = X25519PrivateKey.from_private_bytes(base64.b64decode(my_private_key_b64))
    their_public = X25519PublicKey.from_public_bytes(base64.b64decode(their_public_key_b64))
    return my_private.exchange(their_public)


def derive_transfer_key(shared_secret: bytes, salt: bytes) -> bytes:
    """
    Derive a per-transfer AES-256 key from a pairwise shared secret.

    Uses HKDF-SHA256 with a unique random salt per transfer so that
    compromising one transfer key does not reveal the shared secret
    or any other transfer key on the same channel.

    Args:
        shared_secret: 32-byte X25519 DH output from derive_shared_secret().
        salt: 32-byte random salt unique to this transfer.

    Returns:
        32-byte AES-256 key for encrypt_file / decrypt_file.
    """
    return HKDF(
        algorithm=SHA256(),
        length=KEY_BYTES,
        salt=salt,
        info=HKDF_CHANNEL_INFO,
    ).derive(shared_secret)
