"""
MeshBay cryptographic primitives.

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

import os
import base64

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 ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.argon2 import Argon2id
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes, serialization
import blake3

# ── 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 chunk_key(gek: bytes, file_hash: bytes, chunk_index: int) -> bytes:
    """Derive a per-chunk encryption key from the GEK (deterministic)."""
    return HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=None,
        info=b"file:" + file_hash + b":chunk:" + chunk_index.to_bytes(4, "big"),
    ).derive(gek)

def encrypt_chunk(key: bytes, plaintext: bytes) -> tuple[bytes, bytes]:
    """Encrypt plaintext with ChaCha20-Poly1305. Returns (nonce, ciphertext)."""
    nonce = os.urandom(12)
    ct = ChaCha20Poly1305(key).encrypt(nonce, plaintext, None)
    return nonce, ct

def decrypt_chunk(key: bytes, nonce: bytes, ciphertext: bytes) -> bytes:
    """Decrypt ciphertext. Raises InvalidTag on authentication failure."""
    return ChaCha20Poly1305(key).decrypt(nonce, ciphertext, None)

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 = b"meshbay:gek_wrap:v1"

def wrap_gek(gek: bytes, pk_recipient: bytes) -> dict:
    """
    Wrap a GEK for a recipient using ephemeral X25519 + HKDF + ChaCha20-Poly1305.

    Protocol:
      1. Generate ephemeral (sk_eph, pk_eph)
      2. shared = X25519(sk_eph, pk_recipient)
      3. wrap_key = HKDF(shared, salt=pk_eph, info=GEK_WRAP_INFO)
      4. wrapped = ChaCha20-Poly1305(wrap_key).encrypt(nonce, gek, aad=pk_recipient)

    The hub stores {pk_eph, nonce, wrapped} — opaque, cannot decrypt.
    """
    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,
    ).derive(shared)

    nonce = os.urandom(12)
    wrapped = ChaCha20Poly1305(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(bundle: dict, sk_recipient: bytes, pk_recipient: bytes) -> bytes:
    """
    Unwrap a GEK bundle using the recipient's X25519 private key.
    Raises InvalidTag if the key is wrong or the bundle was tampered.
    """
    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,
    ).derive(shared)

    return ChaCha20Poly1305(wrap_key).decrypt(nonce, wrapped, pk_recipient)


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()

# ── Chunk signing ─────────────────────────────────────────────────────────────

def sign_chunk(sk_node: Ed25519PrivateKey, chunk_index: int,
               nonce: bytes, ct_hash: bytes) -> bytes:
    """Sign chunk metadata. Payload: chunk_index || nonce || ct_hash."""
    payload = chunk_index.to_bytes(4, "big") + nonce + ct_hash
    return sk_node.sign(payload)

def verify_chunk_signature(pk_node: Ed25519PublicKey, chunk_index: int,
                           nonce: bytes, ct_hash: bytes, signature: bytes) -> None:
    """Verify chunk signature. Raises InvalidSignature on failure."""
    payload = chunk_index.to_bytes(4, "big") + nonce + ct_hash
    pk_node.verify(signature, payload)