"""Mesh handshake/session crypto primitives."""
from __future__ import annotations

import base64
import json
import os

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

from mesh.identity import decode_public_key


def b64e(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def b64d(data: str) -> bytes:
    return base64.b64decode(data.encode("ascii"))


def random_salt_b64() -> str:
    return b64e(os.urandom(32))


def generate_ephemeral_keypair() -> tuple[X25519PrivateKey, str]:
    private_key = X25519PrivateKey.generate()
    public_key_b64 = b64e(
        private_key.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
    )
    return private_key, public_key_b64


def derive_session_key(
    *,
    my_static_private: X25519PrivateKey,
    my_ephemeral_private: X25519PrivateKey,
    their_static_public_b64: str,
    their_ephemeral_public_b64: str,
    salt_b64: str,
) -> bytes:
    their_static = decode_public_key(their_static_public_b64)
    their_ephemeral = decode_public_key(their_ephemeral_public_b64)
    static_shared = my_static_private.exchange(their_static)
    ephemeral_shared = my_ephemeral_private.exchange(their_ephemeral)
    ikm = ephemeral_shared + static_shared
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=b64d(salt_b64),
        info=b"mugi-mesh-v1-session",
    )
    return hkdf.derive(ikm)


def nonce_from_counter(counter: int) -> bytes:
    if counter < 0:
        raise ValueError("counter must be non-negative")
    return counter.to_bytes(12, byteorder="big", signed=False)


def encrypt_payload(session_key: bytes, payload: dict, counter: int) -> dict:
    aes = AESGCM(session_key)
    nonce = nonce_from_counter(counter)
    plaintext = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    ciphertext = aes.encrypt(nonce, plaintext, None)
    return {
        "ciphertext": b64e(ciphertext),
        "nonce": b64e(nonce),
        "counter": int(counter),
    }


def decrypt_payload(session_key: bytes, *, ciphertext_b64: str, nonce_b64: str) -> dict:
    aes = AESGCM(session_key)
    plaintext = aes.decrypt(b64d(nonce_b64), b64d(ciphertext_b64), None)
    data = json.loads(plaintext.decode("utf-8"))
    return data if isinstance(data, dict) else {"payload": data}
