"""
pqc_kem.py — hybrid X25519 + ML-KEM-768 key encapsulation in Python.
Mirrors the wire format of the Next.js kit. MIT.

Install:
    pip install pqcrypto pynacl cryptography
"""
from __future__ import annotations
from pqcrypto.kem import ml_kem_768
from nacl.bindings import crypto_scalarmult
from nacl.public import PrivateKey
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF


def _digest(b: bytes) -> bytes:
    h = hashes.Hash(hashes.SHA384())
    h.update(b)
    return h.finalize()


def derive_session(ss_x: bytes, ss_pq: bytes, transcript: bytes) -> bytes:
    return HKDF(
        algorithm=hashes.SHA384(),
        length=32,
        salt=transcript,
        info=b"qbit404-hybrid",
    ).derive(ss_x + ss_pq)


def server_keygen() -> dict:
    pq_pub, pq_priv = ml_kem_768.generate_keypair()
    x_priv = PrivateKey.generate()
    return {
        "pq_pub": pq_pub,
        "pq_priv": pq_priv,
        "x_pub": bytes(x_priv.public_key),
        "x_priv": bytes(x_priv),
    }


def client_encapsulate(server_pq_pub: bytes, server_x_pub: bytes) -> dict:
    pq_ct, ss_pq = ml_kem_768.encrypt(server_pq_pub)
    eph_priv = PrivateKey.generate()
    eph_pub = bytes(eph_priv.public_key)
    ss_x = crypto_scalarmult(bytes(eph_priv), server_x_pub)
    transcript = _digest(pq_ct + eph_pub)
    session = derive_session(ss_x, ss_pq, transcript)
    return {"pq_ct": pq_ct, "x_client_pub": eph_pub, "session": session}


def server_decapsulate(pq_ct: bytes, client_x_pub: bytes,
                       pq_priv: bytes, x_priv: bytes) -> bytes:
    ss_pq = ml_kem_768.decrypt(pq_priv, pq_ct)
    ss_x = crypto_scalarmult(x_priv, client_x_pub)
    transcript = _digest(pq_ct + client_x_pub)
    return derive_session(ss_x, ss_pq, transcript)


if __name__ == "__main__":
    s = server_keygen()
    c = client_encapsulate(s["pq_pub"], s["x_pub"])
    session_s = server_decapsulate(c["pq_ct"], c["x_client_pub"],
                                   s["pq_priv"], s["x_priv"])
    print("client session :", c["session"].hex())
    print("server session :", session_s.hex())
    print("match          :", c["session"] == session_s)