"""Standalone cryptographic functions for the ARBI SDK.

Only depends on PyNaCl, hashlib, base64, os, time. No server-side imports.
"""

from __future__ import annotations

import base64
import hashlib
import hmac as hmac_mod
import os
import time

from nacl.public import Box, PrivateKey, PublicKey, SealedBox
from nacl.pwhash import argon2id
from nacl.secret import SecretBox
from nacl.signing import SigningKey, VerifyKey

# ---------------------------------------------------------------------------
# Key derivation
# ---------------------------------------------------------------------------


def generate_keypair(
    password: str,
    email: str,
    deployment_domain: str,
) -> tuple[str, bytes, str, bytes]:
    """Generate Ed25519 + X25519 keypairs from password.

    Deterministic derivation matching the frontend behaviour:
    salt = SHA-256(email|audience)[:16], Argon2id(password, salt) -> Ed25519 seed.

    Args:
        password: User password.
        email: User email (lowercased for salt).
        deployment_domain: e.g. ``"app.arbi.city"`` — used to form the audience URL.

    Returns:
        ``(signing_pub_b64, signing_priv_bytes, encryption_pub_b64, encryption_priv_bytes)``
    """
    audience = f"https://{deployment_domain}"
    salt_string = f"{email.lower()}|{audience}"
    salt = hashlib.sha256(salt_string.encode()).digest()[:16]

    seed = argon2id.kdf(32, password.encode(), salt, opslimit=2, memlimit=67108864)

    signing_key = SigningKey(seed)
    signing_pub_b64 = base64.b64encode(bytes(signing_key.verify_key)).decode()
    signing_priv_bytes = bytes(signing_key)  # 64 bytes (seed + pub)

    x25519_priv = signing_key.to_curve25519_private_key()
    encryption_pub_b64 = base64.b64encode(bytes(x25519_priv.public_key)).decode()
    encryption_priv_bytes = bytes(x25519_priv)

    return signing_pub_b64, signing_priv_bytes, encryption_pub_b64, encryption_priv_bytes


# ---------------------------------------------------------------------------
# Login signature
# ---------------------------------------------------------------------------


def create_login_signature(
    email: str,
    signing_private_key: bytes,
) -> tuple[str, int]:
    """Create an Ed25519 login signature.

    Args:
        email: User email.
        signing_private_key: 64-byte Ed25519 private key.

    Returns:
        ``(signature_b64, timestamp)``
    """
    timestamp = int(time.time())
    message = f"{email}|{timestamp}".encode()

    signing_key = SigningKey(signing_private_key[:32])
    signed = signing_key.sign(message)
    signature_b64 = base64.b64encode(signed.signature).decode()

    return signature_b64, timestamp


# ---------------------------------------------------------------------------
# Workspace key helpers
# ---------------------------------------------------------------------------


def generate_workspace_key() -> bytes:
    """Generate a random 32-byte workspace key for AES-256 encryption."""
    return os.urandom(32)


def encrypt_workspace_key_for_header(
    workspace_key: bytes,
    session_public_key_b64: str,
) -> str:
    """Encrypt workspace key using SealedBox with the session public key.

    Args:
        workspace_key: 32-byte workspace symmetric key.
        session_public_key_b64: Server's X25519 session public key (base64).

    Returns:
        Base64-encoded SealedBox ciphertext.
    """
    pub = PublicKey(base64.b64decode(session_public_key_b64))
    encrypted = SealedBox(pub).encrypt(workspace_key)
    return base64.b64encode(encrypted).decode()


def wrap_workspace_key_for_sharing(
    workspace_key: bytes,
    recipient_pub_b64: str,
) -> str:
    """Wrap workspace key with recipient's X25519 public key (SealedBox).

    Args:
        workspace_key: 32-byte workspace key.
        recipient_pub_b64: Recipient's X25519 public key (base64).

    Returns:
        Hex-encoded wrapped key.
    """
    pub = PublicKey(base64.b64decode(recipient_pub_b64))
    encrypted = SealedBox(pub).encrypt(workspace_key)
    return encrypted.hex()


# ---------------------------------------------------------------------------
# Key wrapping / unwrapping (standalone reimplementation of encryption.py)
# ---------------------------------------------------------------------------


def unwrap_key_with_x25519(wrapped_key_b64: str, private_key_bytes: bytes) -> bytes:
    """Unwrap a SealedBox-encrypted key.

    Args:
        wrapped_key_b64: Base64-encoded wrapped key.
        private_key_bytes: 32-byte X25519 private key.

    Returns:
        Unwrapped key bytes (32 bytes).
    """
    wrapped = base64.b64decode(wrapped_key_b64)
    priv = PrivateKey(private_key_bytes)
    return SealedBox(priv).decrypt(wrapped)


def wrap_key_with_x25519(key_bytes: bytes, public_key_bytes: bytes) -> str:
    """Wrap a key using SealedBox (X25519 + XSalsa20-Poly1305).

    Args:
        key_bytes: Key to wrap (32 bytes).
        public_key_bytes: X25519 public key bytes.

    Returns:
        Base64-encoded wrapped key.
    """
    pub = PublicKey(public_key_bytes)
    encrypted = SealedBox(pub).encrypt(key_bytes)
    return base64.b64encode(encrypted).decode()


# ---------------------------------------------------------------------------
# Ed25519 <-> X25519 conversion & signature verification
# ---------------------------------------------------------------------------


def ed25519_pk_to_x25519_pk(ed25519_public_key: bytes) -> bytes:
    """Convert Ed25519 public key to X25519 public key.

    Args:
        ed25519_public_key: 32-byte Ed25519 public key.

    Returns:
        32-byte X25519 public key.
    """
    return VerifyKey(ed25519_public_key).to_curve25519_public_key().encode()


def derive_encryption_key_b64(signing_key_b64: str) -> str:
    """Derive X25519 encryption key from Ed25519 signing key (base64 in/out)."""
    signing_bytes = base64.b64decode(signing_key_b64)
    enc_bytes = ed25519_pk_to_x25519_pk(signing_bytes)
    return base64.b64encode(enc_bytes).decode()


def verify_ed25519_signature(
    message: bytes,
    signature: bytes,
    public_key: bytes,
) -> bool:
    """Verify an Ed25519 signature.

    Returns:
        ``True`` if valid, ``False`` otherwise.
    """
    try:
        VerifyKey(public_key).verify(message, signature)
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# SecretBox encrypt / decrypt
# ---------------------------------------------------------------------------


def encrypt_string(plaintext: str, key: bytes) -> str:
    """Encrypt a string with NaCl SecretBox. Returns base64-encoded ciphertext."""
    cipher = SecretBox(key)
    encrypted = cipher.encrypt(plaintext.encode())
    return base64.b64encode(encrypted).decode()


def decrypt_string(encoded: str, key: bytes) -> bytes:
    """Decrypt base64-encoded SecretBox ciphertext. Returns raw bytes."""
    cipher = SecretBox(key)
    encrypted_bytes = base64.b64decode(encoded)
    return cipher.decrypt(encrypted_bytes)


# ---------------------------------------------------------------------------
# E2E messaging (X25519 DH shared secret)
# ---------------------------------------------------------------------------


def encrypt_message(
    plaintext: str,
    recipient_pub_b64: str,
    my_private_key: bytes,
) -> str:
    """Encrypt message using X25519 DH shared secret (crypto_box_easy).

    Args:
        plaintext: Message to encrypt.
        recipient_pub_b64: Recipient's X25519 public key (base64).
        my_private_key: Sender's X25519 private key (32 bytes).

    Returns:
        Base64-encoded encrypted content (nonce + ciphertext).
    """
    my_priv = PrivateKey(my_private_key)
    their_pub = PublicKey(base64.b64decode(recipient_pub_b64))
    box = Box(my_priv, their_pub)
    encrypted = box.encrypt(plaintext.encode())
    return base64.b64encode(encrypted).decode()


def decrypt_message(
    encrypted_b64: str,
    sender_pub_b64: str,
    my_private_key: bytes,
) -> str:
    """Decrypt message using X25519 DH shared secret (crypto_box_open_easy).

    Args:
        encrypted_b64: Base64-encoded encrypted content.
        sender_pub_b64: Sender's X25519 public key (base64).
        my_private_key: Recipient's X25519 private key (32 bytes).

    Returns:
        Decrypted plaintext string.
    """
    my_priv = PrivateKey(my_private_key)
    their_pub = PublicKey(base64.b64decode(sender_pub_b64))
    box = Box(my_priv, their_pub)
    encrypted = base64.b64decode(encrypted_b64)
    return box.decrypt(encrypted).decode()


# ---------------------------------------------------------------------------
# Content hashing
# ---------------------------------------------------------------------------


def create_content_hash(content: bytes, workspace_key: bytes) -> str:
    """Create workspace-scoped content hash for deduplication.

    Uses HMAC-SHA256 keyed with workspace_key.

    Returns:
        16-character hex string.
    """
    return hmac_mod.new(workspace_key, content, hashlib.sha256).hexdigest()[:16]
