"""
MeshBay cryptographic primitives.

Validated in Spike 1 and Spike 6 of the POC.
All operations use PyCA cryptography (OpenSSL-backed, hardware-accelerated).
"""

import base64
import os

import blake3
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.argon2 import Argon2id
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

# ── Key serialisation helpers ─────────────────────────────────────────────────

def sk_to_raw(sk: Ed25519PrivateKey | X25519PrivateKey) -> bytes:
    return sk.private_bytes(
        serialization.Encoding.Raw,
        serialization.PrivateFormat.Raw,
        serialization.NoEncryption(),
    )

def pk_to_raw(pk: Ed25519PublicKey | X25519PublicKey) -> bytes:
    return pk.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)

def pk_to_b64(pk: Ed25519PublicKey | X25519PublicKey) -> str:
    return base64.b64encode(pk_to_raw(pk)).decode()

def sk_to_b64(sk: Ed25519PrivateKey | X25519PrivateKey) -> str:
    return base64.b64encode(sk_to_raw(sk)).decode()

# ── GEK (Group Encryption Key) ────────────────────────────────────────────────

def generate_gek() -> bytes:
    """Generate a fresh 256-bit Group Encryption Key."""
    return ChaCha20Poly1305.generate_key()

def file_hash(path_or_bytes) -> bytes:
    """Compute blake3 hash of a file (bytes or path-like)."""
    if isinstance(path_or_bytes, (str, bytes)) and not isinstance(path_or_bytes, bytes):
        import pathlib
        data = pathlib.Path(path_or_bytes).read_bytes()
    elif isinstance(path_or_bytes, bytes):
        data = path_or_bytes
    else:
        data = path_or_bytes.read_bytes()
    return blake3.blake3(data).digest()

# ── GEK wrapping (ECIES-like) ─────────────────────────────────────────────────

GEK_WRAP_INFO_AES = b"meshbay:gek_wrap:v1:aes"

def wrap_gek_aes(gek: bytes, pk_recipient: bytes) -> dict:
    """ECIES wrap using AES-256-GCM — compatible with browser WebCrypto."""
    sk_eph = X25519PrivateKey.generate()
    pk_eph_raw = pk_to_raw(sk_eph.public_key())

    shared = sk_eph.exchange(X25519PublicKey.from_public_bytes(pk_recipient))
    wrap_key = HKDF(
        algorithm=hashes.SHA256(), length=32,
        salt=pk_eph_raw, info=GEK_WRAP_INFO_AES,
    ).derive(shared)

    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    nonce = os.urandom(12)
    wrapped = AESGCM(wrap_key).encrypt(nonce, gek, pk_recipient)

    return {
        "pk_eph_b64":  base64.b64encode(pk_eph_raw).decode(),
        "nonce_b64":   base64.b64encode(nonce).decode(),
        "wrapped_b64": base64.b64encode(wrapped).decode(),
    }

def unwrap_gek_aes(bundle: dict, sk_recipient: bytes, pk_recipient: bytes) -> bytes:
    """Unwrap a GEK bundle created by browser (AES-256-GCM ECIES)."""
    pk_eph_raw = base64.b64decode(bundle["pk_eph_b64"])
    nonce      = base64.b64decode(bundle["nonce_b64"])
    wrapped    = base64.b64decode(bundle["wrapped_b64"])

    shared = X25519PrivateKey.from_private_bytes(sk_recipient).exchange(
        X25519PublicKey.from_public_bytes(pk_eph_raw)
    )
    wrap_key = HKDF(
        algorithm=hashes.SHA256(), length=32,
        salt=pk_eph_raw, info=GEK_WRAP_INFO_AES,
    ).derive(shared)

    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    return AESGCM(wrap_key).decrypt(nonce, wrapped, pk_recipient)


# ── Keystore (local key storage) ──────────────────────────────────────────────

# Argon2id parameters for the node keystore.
#
# Finding M2: these sat at 64 MB long after the hub's password verifier was raised
# to 256 MB, and the docs recorded the bump as done — true for the hub, false here.
# The keystore protects the node's Ed25519 and X25519 private keys, so it is the
# more valuable target of the two.
#
# Parameters are recorded in each keystore envelope, so raising them does not
# invalidate existing files: LEGACY_* is used when an envelope predates the field.
ARGON2_ITERATIONS   = 3
ARGON2_MEMORY_COST  = 262144   # 256 MB
ARGON2_LANES        = 4
ARGON2_KEY_LENGTH   = 32

LEGACY_ARGON2_ITERATIONS  = 3
LEGACY_ARGON2_MEMORY_COST = 65536    # 64 MB — keystores written before M2
LEGACY_ARGON2_LANES       = 4


def derive_keystore_key(
    password: str,
    salt: bytes,
    *,
    iterations: int | None = None,
    memory_cost: int | None = None,
    lanes: int | None = None,
) -> bytes:
    """
    Derive an AES-256 key from a password using Argon2id.

    Parameters default to the current production values; callers pass the values
    recorded in an existing envelope when opening an older keystore.
    """
    return Argon2id(
        salt=salt,
        length=ARGON2_KEY_LENGTH,
        iterations=ARGON2_ITERATIONS if iterations is None else iterations,
        lanes=ARGON2_LANES if lanes is None else lanes,
        memory_cost=ARGON2_MEMORY_COST if memory_cost is None else memory_cost,
    ).derive(password.encode())

def encrypt_keystore(plaintext: bytes, key: bytes) -> tuple[bytes, bytes, bytes]:
    """Encrypt keystore blob with AES-256-GCM. Returns (iv, ciphertext, tag)."""
    iv = os.urandom(12)
    enc = Cipher(algorithms.AES(key), modes.GCM(iv)).encryptor()
    ct  = enc.update(plaintext) + enc.finalize()
    return iv, ct, enc.tag

def decrypt_keystore(iv: bytes, ciphertext: bytes, tag: bytes, key: bytes) -> bytes:
    """Decrypt keystore blob. Raises on authentication failure."""
    dec = Cipher(algorithms.AES(key), modes.GCM(iv, tag)).decryptor()
    return dec.update(ciphertext) + dec.finalize()