/* Test for Phase 5 ML-KEM-768 (FIPS 203) post-quantum key encapsulation.
 *
 * Verifies:
 *   - ML-KEM-768 keygen from seed: 1184-byte pub, 2400-byte priv
 *   - ML-KEM-768 keygen determinism: same seed -> same keypair
 *   - Encaps/decaps roundtrip: shared secrets match
 *   - Encaps produces different ciphertexts each call (non-deterministic)
 *   - Decaps with wrong ciphertext: different shared secret (implicit rejection)
 *   - Integration: derive ML-KEM-768 key from mnemonic, encaps via dispatcher,
 *     decaps via dispatcher, shared secrets match
 *   - Enforcement: encapsulate/decapsulate allowed on pq-kem+ml-kem-768
 */
/* NSIGNER_HEADERLESS_DECLS_BEGIN */
#include <stddef.h>
#include <stdint.h>
#include <sys/types.h>
#include <cJSON.h>

/* from secure_mem.h */
typedef struct {
    void *data;
    size_t size;
    int locked;
} secure_buf_t;

int secure_buf_alloc(secure_buf_t *buf, size_t size);
void secure_buf_free(secure_buf_t *buf);
void secure_memzero(void *ptr, size_t len);

/* from mnemonic.h */
#define MNEMONIC_MAX_LEN 256
typedef struct {
    secure_buf_t buf;
    int loaded;
    int word_count;
} mnemonic_state_t;

void mnemonic_init(mnemonic_state_t *state);
int mnemonic_load(mnemonic_state_t *state, const char *phrase);
void mnemonic_unload(mnemonic_state_t *state);
int mnemonic_is_loaded(const mnemonic_state_t *state);
const char *mnemonic_get_phrase(const mnemonic_state_t *state);
int mnemonic_generate(int word_count, char *out, size_t out_len);

/* from role_table.h */
#define ROLE_NAME_MAX 64
#define ROLE_PATH_MAX 128
#define ROLE_PURPOSE_MAX 32
#define ROLE_CURVE_MAX 16
#define ROLE_PUBKEY_HEX_MAX 66
#define ROLE_TABLE_MAX_ENTRIES 256

typedef enum {
    PURPOSE_NOSTR = 0, PURPOSE_BITCOIN, PURPOSE_SSH, PURPOSE_AGE,
    PURPOSE_FIPS, PURPOSE_PQ_SIG, PURPOSE_PQ_KEM, PURPOSE_UNKNOWN
} role_purpose_t;

typedef enum {
    CURVE_SECP256K1 = 0, CURVE_ED25519, CURVE_X25519,
    CURVE_ML_DSA_65, CURVE_SLH_DSA_128S, CURVE_ML_KEM_768, CURVE_UNKNOWN
} role_curve_t;

typedef enum { SELECTOR_NOSTR_INDEX, SELECTOR_ROLE_PATH } role_selector_type_t;

typedef struct {
    char name[ROLE_NAME_MAX];
    char purpose_str[ROLE_PURPOSE_MAX];
    char curve_str[ROLE_CURVE_MAX];
    role_purpose_t purpose;
    role_curve_t curve;
    role_selector_type_t selector_type;
    int nostr_index;
    char role_path[ROLE_PATH_MAX];
    char pubkey_hex[ROLE_PUBKEY_HEX_MAX];
    int derived;
} role_entry_t;

typedef struct {
    role_entry_t entries[ROLE_TABLE_MAX_ENTRIES];
    int count;
} role_table_t;

void role_table_init(role_table_t *table);
int role_table_add(role_table_t *table, const role_entry_t *entry);
role_entry_t *role_table_find_by_name(role_table_t *table, const char *name);
role_purpose_t role_purpose_from_str(const char *s);
role_curve_t role_curve_from_str(const char *s);
const char *role_purpose_to_str(role_purpose_t p);
const char *role_curve_to_str(role_curve_t c);

/* from selector.h */
#define SELECTOR_OK 0
#define SELECTOR_ERR_AMBIGUOUS -1
#define SELECTOR_ERR_NOT_FOUND -2
#define SELECTOR_ERR_NO_DEFAULT -3

typedef struct {
    int has_role;
    char role_name[ROLE_NAME_MAX];
    int has_nostr_index;
    int nostr_index;
    int has_role_path;
    char role_path[ROLE_PATH_MAX];
} selector_request_t;

void selector_request_init(selector_request_t *req);
int selector_resolve(const selector_request_t *req, role_table_t *table, role_entry_t **out);
const char *selector_strerror(int err);

/* from enforcement.h */
#define ENFORCE_OK 0
#define ENFORCE_ERR_PURPOSE -1
#define ENFORCE_ERR_CURVE -2
#define ENFORCE_ERR_UNKNOWN_VERB -3
#define ENFORCE_ERR_ALGORITHM -4

#define VERB_GET_PUBLIC_KEY       "get_public_key"
#define VERB_SIGN                 "sign"
#define VERB_VERIFY               "verify"
#define VERB_ENCAPSULATE          "encapsulate"
#define VERB_DECAPSULATE          "decapsulate"
#define VERB_DERIVE_SHARED        "derive_shared_secret"
#define VERB_DERIVE           "derive"
#define VERB_NOSTR_GET_PUBLIC_KEY "nostr_get_public_key"
#define VERB_NOSTR_SIGN_EVENT     "nostr_sign_event"
#define VERB_NOSTR_MINE_EVENT     "nostr_mine_event"
#define VERB_NOSTR_NIP44_ENCRYPT  "nostr_nip44_encrypt"
#define VERB_NOSTR_NIP44_DECRYPT  "nostr_nip44_decrypt"
#define VERB_NOSTR_NIP04_ENCRYPT  "nostr_nip04_encrypt"
#define VERB_NOSTR_NIP04_DECRYPT  "nostr_nip04_decrypt"

int enforce_verb_role(const char *verb, const role_entry_t *role);
const char *enforce_strerror(int err);

/* from policy.h */
#define POLICY_MAX_ENTRIES 32
#define POLICY_MAX_VERBS 16
#define POLICY_MAX_ROLES 16
#define POLICY_MAX_PURPOSES 8
#define POLICY_VERB_MAX_LEN 32
#define POLICY_CALLER_MAX_LEN 160
#define POLICY_MAX_ALGS       16
#define POLICY_MAX_ALGS       16

typedef enum { PROMPT_NEVER = 0, PROMPT_FIRST_PER_BOOT, PROMPT_EVERY_REQUEST, PROMPT_DENY } prompt_mode_t;
typedef enum { POLICY_SOURCE_DEFAULT = 0, POLICY_SOURCE_PREAPPROVE, POLICY_SOURCE_SESSION_GRANT } policy_source_t;

typedef struct {
    char caller[POLICY_CALLER_MAX_LEN];
    char verbs[POLICY_MAX_VERBS][POLICY_VERB_MAX_LEN];
    int verb_count;
    char roles[POLICY_MAX_ROLES][ROLE_NAME_MAX];
    int role_count;
    char purposes[POLICY_MAX_PURPOSES][ROLE_PURPOSE_MAX];
    int purpose_count;
    /* Algorithm-based (new) */
    char algorithms[16][32];  /* algorithm names; POLICY_MAX_ALGS */
    int alg_count;
    int index_min;            /* -1 = any */
    int index_max;            /* -1 = any */
    /* Common */
    prompt_mode_t prompt;
    policy_source_t source;
} policy_entry_t;

typedef struct {
    policy_entry_t entries[POLICY_MAX_ENTRIES];
    int count;
} policy_table_t;

#define POLICY_ALLOW 0
#define POLICY_DENY -1
#define POLICY_PROMPT -2
#define POLICY_NO_MATCH -3

void policy_table_init(policy_table_t *table);
void policy_init_default(policy_table_t *table, uid_t owner_uid);
int policy_table_add(policy_table_t *table, const policy_entry_t *entry);
int policy_check(const policy_table_t *table, const char *caller_id, const char *verb, const char *role_name, const char *purpose, policy_source_t *out_source);
prompt_mode_t prompt_mode_from_str(const char *s);
const char *prompt_mode_to_str(prompt_mode_t m);

/* from pq_crypto.h */
typedef enum {
    CRYPTO_ALG_SECP256K1 = 0, CRYPTO_ALG_ED25519, CRYPTO_ALG_X25519,
    CRYPTO_ALG_ML_DSA_65, CRYPTO_ALG_SLH_DSA_128S, CRYPTO_ALG_ML_KEM_768,
    CRYPTO_ALG_UNKNOWN
} crypto_alg_t;

typedef struct {
    size_t priv_key_len;
    size_t pub_key_len;
    size_t sig_len;
    size_t ciphertext_len;
    size_t shared_secret_len;
} crypto_alg_sizes_t;

const crypto_alg_sizes_t *crypto_alg_get_sizes(crypto_alg_t alg);
crypto_alg_t crypto_alg_from_role(role_curve_t curve, role_purpose_t purpose);
const char *crypto_alg_to_str(crypto_alg_t alg);
crypto_alg_t crypto_alg_from_str(const char *s);

int crypto_ed25519_keygen_from_seed(const unsigned char *seed, size_t seed_len, unsigned char *priv_out, unsigned char *pub_out);
int crypto_ed25519_sign(const unsigned char *priv, size_t priv_len, const unsigned char *msg, size_t msg_len, unsigned char *sig_out, size_t *sig_out_len);
int crypto_ed25519_verify(const unsigned char *pub, size_t pub_len, const unsigned char *msg, size_t msg_len, const unsigned char *sig, size_t sig_len);
int crypto_x25519_keygen_from_seed(const unsigned char *seed, size_t seed_len, unsigned char *priv_out, unsigned char *pub_out);
int crypto_x25519_ecdh(const unsigned char *our_priv, size_t priv_len, const unsigned char *peer_pub, size_t pub_len, unsigned char *shared_out, size_t *shared_out_len);
int crypto_derive_seed_from_mnemonic(const char *mnemonic, const char *path, unsigned char *seed_out, size_t seed_out_len);
int crypto_ml_dsa_65_keygen_from_seed(const unsigned char *seed, size_t seed_len, unsigned char *priv_out, unsigned char *pub_out);
int crypto_ml_dsa_65_sign(const unsigned char *priv, size_t priv_len, const unsigned char *msg, size_t msg_len, unsigned char *sig_out, size_t *sig_out_len);
int crypto_ml_dsa_65_verify(const unsigned char *pub, size_t pub_len, const unsigned char *msg, size_t msg_len, const unsigned char *sig, size_t sig_len);
int crypto_slh_dsa_128s_keygen_from_seed(const unsigned char *seed, size_t seed_len, unsigned char *priv_out, unsigned char *pub_out);
int crypto_slh_dsa_128s_sign(const unsigned char *priv, size_t priv_len, const unsigned char *msg, size_t msg_len, unsigned char *sig_out, size_t *sig_out_len);
int crypto_slh_dsa_128s_verify(const unsigned char *pub, size_t pub_len, const unsigned char *msg, size_t msg_len, const unsigned char *sig, size_t sig_len);

/* ML-KEM-768: generate keypair from a 32-byte seed (deterministic). */
int crypto_ml_kem_768_keygen_from_seed(const unsigned char *seed, size_t seed_len, unsigned char *priv_out, unsigned char *pub_out);
/* ML-KEM-768: encapsulate. */
int crypto_ml_kem_768_encaps(const unsigned char *pub, size_t pub_len, unsigned char *ct_out, unsigned char *ss_out);
/* ML-KEM-768: decapsulate. */
int crypto_ml_kem_768_decaps(const unsigned char *priv, size_t priv_len, const unsigned char *ct, size_t ct_len, unsigned char *ss_out);

void pq_drbg_init(const unsigned char *seed, size_t seed_len);
int  pq_drbg_randombytes(unsigned char *buf, size_t len);
void pq_drbg_zeroize(void);

/* from crypto.h */
typedef struct {
    secure_buf_t private_key;
    secure_buf_t public_key;
    char pubkey_hex[8192];
    char npub[128];
    crypto_alg_t alg;
    int valid;
} derived_key_t;

typedef struct {
    derived_key_t keys[ROLE_TABLE_MAX_ENTRIES];
    int count;
} key_store_t;

int crypto_derive_all(key_store_t *store, role_table_t *table, const mnemonic_state_t *mnemonic);
int crypto_derive_one(key_store_t *store, role_table_t *table, const mnemonic_state_t *mnemonic, int role_index);
const unsigned char *crypto_get_private_key(const key_store_t *store, int role_index);
const char *crypto_get_pubkey_hex(const key_store_t *store, int role_index);
char *crypto_sign_event(const key_store_t *store, int role_index, const char *event_json);
void crypto_wipe(key_store_t *store);


/* from alg_api.h */


/* Check whether a verb is valid for an algorithm (algorithm-based enforcement).
 * Returns ENFORCE_OK, ENFORCE_ERR_ALGORITHM, or ENFORCE_ERR_UNKNOWN_VERB.
 * Does NOT check purpose — purpose is irrelevant for the new verbs. */
int enforce_verb_algorithm(const char *verb, crypto_alg_t alg);

/* secp256k1 Schnorr (BIP-340) sign arbitrary bytes.
 * priv is 32-byte scalar, pub is 32-byte x-only pubkey, sig_out is 64 bytes.
 * Hashes the message with SHA-256 before signing (like Nostr event signing).
 * Returns 0 on success, -1 on error. */
int crypto_secp256k1_schnorr_sign(const unsigned char *priv, size_t priv_len,
                                    const unsigned char *msg, size_t msg_len,
                                    unsigned char *sig_out, size_t *sig_out_len);

/* secp256k1 Schnorr (BIP-340) verify.
 * pub is 32-byte x-only pubkey, sig is 64 bytes.
 * Returns 0 on valid, 1 on invalid, -1 on error. */
int crypto_secp256k1_schnorr_verify(const unsigned char *pub, size_t pub_len,
                                      const unsigned char *msg, size_t msg_len,
                                      const unsigned char *sig, size_t sig_len);

/* secp256k1 ECDSA sign arbitrary bytes.
 * priv is 32-byte scalar, sig_out must be at least 64 bytes (compact DER r||s).
 * Hashes the message with SHA-256 before signing.
 * Returns 0 on success, -1 on error. */
int crypto_secp256k1_ecdsa_sign(const unsigned char *priv, size_t priv_len,
                                  const unsigned char *msg, size_t msg_len,
                                  unsigned char *sig_out, size_t *sig_out_len);

/* secp256k1 ECDSA verify.
 * pub is 32-byte x-only pubkey (converted internally to compressed form).
 * sig is 64-byte compact (r||s). Returns 0 on valid, 1 on invalid, -1 on error. */
int crypto_secp256k1_ecdsa_verify(const unsigned char *pub, size_t pub_len,
                                    const unsigned char *msg, size_t msg_len,
                                    const unsigned char *sig, size_t sig_len);

/* ---- Algorithm key cache ----
 * On-demand key derivation by algorithm+index, separate from the role-based
 * key_store. Holds up to ALG_KEY_CACHE_MAX derived keys in secure memory.
 * When full, the oldest entry is evicted (FIFO). */

#define ALG_KEY_CACHE_MAX 32

typedef struct {
    crypto_alg_t alg;
    int index;
    secure_buf_t private_key;
    secure_buf_t public_key;
    char pubkey_hex[8192];
    char key_id[17];
    int valid;
} alg_key_entry_t;

typedef struct {
    alg_key_entry_t entries[ALG_KEY_CACHE_MAX];
    int count;
} algorithm_key_cache_t;

/* Initialize an empty cache. */
void alg_key_cache_init(algorithm_key_cache_t *cache);

/* Zeroize and free all entries. Idempotent. */
void alg_key_cache_wipe(algorithm_key_cache_t *cache);

/* Look up a cached entry by (alg, index). Returns NULL if not present. */
const alg_key_entry_t *alg_key_cache_get(algorithm_key_cache_t *cache, crypto_alg_t alg, int index);

/* Derive a key on-demand by (alg, index) and store it in the cache.
 * Uses the standard derivation path for the algorithm.
 * Returns 0 on success, -1 on error. */
int alg_key_cache_derive(algorithm_key_cache_t *cache, const mnemonic_state_t *mnemonic, crypto_alg_t alg, int index);


/* from dispatcher.h */
typedef struct {
    role_table_t *role_table;
    mnemonic_state_t *mnemonic;
    key_store_t *key_store;
} dispatcher_ctx_t;

void dispatcher_init(dispatcher_ctx_t *ctx, role_table_t *table, mnemonic_state_t *mnemonic, key_store_t *key_store, algorithm_key_cache_t *alg_key_cache);
char *dispatcher_handle_request(dispatcher_ctx_t *ctx, const char *json_request);

/* NSIGNER_HEADERLESS_DECLS_END */

static key_store_t g_key_store;

static algorithm_key_cache_t g_alg_key_cache;

#include <nostr_core/nostr_common.h>
#include <nostr_core/utils.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* ML-KEM-768 key sizes (FIPS 203) */
#define ML_KEM_768_PUBKEY_BYTES  1184
#define ML_KEM_768_PRIVKEY_BYTES 2400
#define ML_KEM_768_CT_BYTES      1088
#define ML_KEM_768_SS_BYTES      32

static int g_passes = 0;
static int g_total = 0;

static void check_condition(const char *name, int condition) {
    g_total++;
    if (condition) {
        printf("PASS: %s\n", name);
        g_passes++;
    } else {
        printf("FAIL: %s\n", name);
    }
}

static int response_has(const char *response, const char *needle) {
    return (response != NULL && needle != NULL && strstr(response, needle) != NULL);
}

static role_entry_t make_pq_kem_ml_kem_768_entry(const char *name, int idx) {
    role_entry_t e;
    memset(&e, 0, sizeof(e));
    strncpy(e.name, name, sizeof(e.name) - 1);
    strncpy(e.purpose_str, "pq-kem", sizeof(e.purpose_str) - 1);
    strncpy(e.curve_str, "ml-kem-768", sizeof(e.curve_str) - 1);
    e.purpose = role_purpose_from_str(e.purpose_str);
    e.curve = role_curve_from_str(e.curve_str);
    e.selector_type = SELECTOR_NOSTR_INDEX;
    e.nostr_index = idx;
    e.derived = 0;
    return e;
}

static role_entry_t make_nostr_secp_entry(const char *name, int idx) {
    role_entry_t e;
    memset(&e, 0, sizeof(e));
    strncpy(e.name, name, sizeof(e.name) - 1);
    strncpy(e.purpose_str, "nostr", sizeof(e.purpose_str) - 1);
    strncpy(e.curve_str, "secp256k1", sizeof(e.curve_str) - 1);
    e.purpose = role_purpose_from_str(e.purpose_str);
    e.curve = role_curve_from_str(e.curve_str);
    e.selector_type = SELECTOR_NOSTR_INDEX;
    e.nostr_index = idx;
    e.derived = 0;
    return e;
}

int main(void) {
    const char *mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about";
    unsigned char seed[32];
    unsigned char *priv1, *pub1, *priv2, *pub2;
    unsigned char *ct1, *ct2, *ss1, *ss2, *ss3;
    int rc;

    priv1 = (unsigned char *)malloc(ML_KEM_768_PRIVKEY_BYTES);
    pub1 = (unsigned char *)malloc(ML_KEM_768_PUBKEY_BYTES);
    priv2 = (unsigned char *)malloc(ML_KEM_768_PRIVKEY_BYTES);
    pub2 = (unsigned char *)malloc(ML_KEM_768_PUBKEY_BYTES);
    ct1 = (unsigned char *)malloc(ML_KEM_768_CT_BYTES);
    ct2 = (unsigned char *)malloc(ML_KEM_768_CT_BYTES);
    ss1 = (unsigned char *)malloc(ML_KEM_768_SS_BYTES);
    ss2 = (unsigned char *)malloc(ML_KEM_768_SS_BYTES);
    ss3 = (unsigned char *)malloc(ML_KEM_768_SS_BYTES);

    if (!priv1 || !pub1 || !priv2 || !pub2 || !ct1 || !ct2 || !ss1 || !ss2 || !ss3) {
        printf("FAIL: memory allocation\n");
        return 1;
    }

    /* ---- ML-KEM-768 keygen from seed ---- */
    memset(seed, 0x42, 32);
    rc = crypto_ml_kem_768_keygen_from_seed(seed, 32, priv1, pub1);
    check_condition("ML-KEM-768 keygen from seed succeeds", rc == 0);

    /* Keygen determinism: same seed -> same keypair */
    memset(seed, 0x42, 32);
    rc = crypto_ml_kem_768_keygen_from_seed(seed, 32, priv2, pub2);
    check_condition("ML-KEM-768 keygen determinism (same seed -> same key)", rc == 0 &&
                    memcmp(priv1, priv2, ML_KEM_768_PRIVKEY_BYTES) == 0 &&
                    memcmp(pub1, pub2, ML_KEM_768_PUBKEY_BYTES) == 0);

    /* Different seed -> different keypair */
    memset(seed, 0x99, 32);
    rc = crypto_ml_kem_768_keygen_from_seed(seed, 32, priv2, pub2);
    check_condition("ML-KEM-768 keygen different seed -> different key", rc == 0 &&
                    memcmp(pub1, pub2, ML_KEM_768_PUBKEY_BYTES) != 0);

    /* ---- Encaps/decaps roundtrip ---- */
    memset(seed, 0x42, 32);
    crypto_ml_kem_768_keygen_from_seed(seed, 32, priv1, pub1);

    rc = crypto_ml_kem_768_encaps(pub1, ML_KEM_768_PUBKEY_BYTES, ct1, ss1);
    check_condition("ML-KEM-768 encaps succeeds", rc == 0);

    rc = crypto_ml_kem_768_decaps(priv1, ML_KEM_768_PRIVKEY_BYTES, ct1, ML_KEM_768_CT_BYTES, ss2);
    check_condition("ML-KEM-768 decaps succeeds", rc == 0);

    check_condition("ML-KEM-768 encaps/decaps shared secrets match",
                    memcmp(ss1, ss2, ML_KEM_768_SS_BYTES) == 0);

    /* ---- Encaps produces different ciphertexts each call ---- */
    rc = crypto_ml_kem_768_encaps(pub1, ML_KEM_768_PUBKEY_BYTES, ct2, ss3);
    check_condition("ML-KEM-768 second encaps succeeds", rc == 0);
    check_condition("ML-KEM-768 encaps produces different ciphertexts",
                    memcmp(ct1, ct2, ML_KEM_768_CT_BYTES) != 0);

    /* Decaps with second ciphertext also matches its shared secret */
    {
        unsigned char ss4[32];
        crypto_ml_kem_768_decaps(priv1, ML_KEM_768_PRIVKEY_BYTES, ct2, ML_KEM_768_CT_BYTES, ss4);
        check_condition("ML-KEM-768 second encaps/decaps shared secrets match",
                        memcmp(ss3, ss4, ML_KEM_768_SS_BYTES) == 0);
    }

    /* ---- Decaps with wrong ciphertext: implicit rejection ---- */
    {
        ct1[0] ^= 0xFF;
        rc = crypto_ml_kem_768_decaps(priv1, ML_KEM_768_PRIVKEY_BYTES, ct1, ML_KEM_768_CT_BYTES, ss3);
        ct1[0] ^= 0xFF; /* restore */
        check_condition("ML-KEM-768 decaps wrong ct succeeds (no crash)", rc == 0);
        /* The shared secret should be different from the original (implicit rejection) */
        check_condition("ML-KEM-768 wrong ct gives different shared secret",
                        memcmp(ss1, ss3, ML_KEM_768_SS_BYTES) != 0);
    }

    /* ---- Integration: derive ML-KEM-768 key via crypto_derive_all ---- */
    {
        role_table_t table;
        role_entry_t pq_role;
        static mnemonic_state_t mnemonic_state;
            int derived;

        role_table_init(&table);
        pq_role = make_pq_kem_ml_kem_768_entry("pq_kem", 0);
        role_table_add(&table, &pq_role);

        mnemonic_init(&mnemonic_state);
        mnemonic_load(&mnemonic_state, mnemonic);

        memset(&g_key_store, 0, sizeof(g_key_store));
        alg_key_cache_init(&g_alg_key_cache);
        derived = crypto_derive_all(&g_key_store, &table, &mnemonic_state);
        check_condition("crypto_derive_all derives ML-KEM-768 key", derived == 1);

        {
            const char *pub_hex = crypto_get_pubkey_hex(&g_key_store, 0);
            /* 1184 bytes hex-encoded = 2368 hex chars + null */
            check_condition("ML-KEM-768 derived pubkey hex is 2368 chars",
                            pub_hex != NULL && strlen(pub_hex) == 2368);
        }

        {
            const unsigned char *priv = crypto_get_private_key(&g_key_store, 0);
            check_condition("ML-KEM-768 derived private key not NULL", priv != NULL);
        }

        /* Determinism: re-derive and check same key */
        {
            static key_store_t key_store2;
            memset(&key_store2, 0, sizeof(key_store2));
            crypto_derive_all(&key_store2, &table, &mnemonic_state);
            {
                const char *pub_hex1 = crypto_get_pubkey_hex(&g_key_store, 0);
                const char *pub_hex2 = crypto_get_pubkey_hex(&key_store2, 0);
                check_condition("ML-KEM-768 derivation determinism",
                                pub_hex1 != NULL && pub_hex2 != NULL &&
                                strcmp(pub_hex1, pub_hex2) == 0);
            }
            crypto_wipe(&key_store2);
        }

        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- Integration: encapsulate / decapsulate via dispatcher ---- */
    {
        role_table_t table;
        role_entry_t pq_role;
        static mnemonic_state_t mnemonic_state;
        dispatcher_ctx_t dispatcher;
        char *encaps_resp, *decaps_resp;
        const char *pub_hex;
        char encaps_req[6000]; /* pub_hex is 2368 chars + JSON overhead */
        char *decaps_req = (char *)malloc(6000);
        cJSON *encaps_json, *result_item, *ct_item, *ss_item;
        const char *ct_str = NULL, *ss_encaps_str = NULL;
        const char *ss_decaps_str = NULL;

        role_table_init(&table);
        pq_role = make_pq_kem_ml_kem_768_entry("pq_kem", 0);
        role_table_add(&table, &pq_role);

        mnemonic_init(&mnemonic_state);
        mnemonic_load(&mnemonic_state, mnemonic);

        memset(&g_key_store, 0, sizeof(g_key_store));
        alg_key_cache_init(&g_alg_key_cache);
        crypto_derive_all(&g_key_store, &table, &mnemonic_state);
        dispatcher_init(&dispatcher, &table, &mnemonic_state, &g_key_store, &g_alg_key_cache);

        /* Get the public key hex for encaps */
        pub_hex = crypto_get_pubkey_hex(&g_key_store, 0);
        check_condition("dispatcher integration: pubkey available", pub_hex != NULL);

        if (pub_hex != NULL && decaps_req != NULL) {
            /* encapsulate request */
            snprintf(encaps_req, sizeof(encaps_req),
                     "{\"id\":\"e1\",\"method\":\"encapsulate\",\"params\":[\"%s\",{\"algorithm\":\"ml-kem-768\",\"index\":0}]}",
                     pub_hex);
            encaps_resp = dispatcher_handle_request(&dispatcher, encaps_req);
            check_condition("encapsulate via dispatcher returns result",
                            encaps_resp != NULL && response_has(encaps_resp, "\"result\""));
            check_condition("encapsulate result contains ciphertext",
                            encaps_resp != NULL && response_has(encaps_resp, "ciphertext"));
            check_condition("encapsulate result contains shared_secret",
                            encaps_resp != NULL && response_has(encaps_resp, "shared_secret"));
            check_condition("encapsulate result contains algorithm ml-kem-768",
                            encaps_resp != NULL && response_has(encaps_resp, "ml-kem-768"));

            /* Extract ciphertext and shared_secret from result */
            if (encaps_resp != NULL) {
                encaps_json = cJSON_Parse(encaps_resp);
                if (encaps_json != NULL) {
                    result_item = cJSON_GetObjectItemCaseSensitive(encaps_json, "result");
                    if (cJSON_IsString(result_item)) {
                        cJSON *result_obj = cJSON_Parse(result_item->valuestring);
                        if (result_obj != NULL) {
                            ct_item = cJSON_GetObjectItemCaseSensitive(result_obj, "ciphertext");
                            ss_item = cJSON_GetObjectItemCaseSensitive(result_obj, "shared_secret");
                            if (cJSON_IsString(ct_item)) ct_str = strdup(ct_item->valuestring);
                            if (cJSON_IsString(ss_item)) ss_encaps_str = strdup(ss_item->valuestring);
                            cJSON_Delete(result_obj);
                        }
                    }
                    cJSON_Delete(encaps_json);
                }
            }
            check_condition("encapsulate: extracted ciphertext hex", ct_str != NULL);
            check_condition("encapsulate: extracted shared_secret hex", ss_encaps_str != NULL);

            /* decapsulate request */
            if (ct_str != NULL) {
                snprintf(decaps_req, 6000,
                         "{\"id\":\"d1\",\"method\":\"decapsulate\",\"params\":[\"%s\",{\"algorithm\":\"ml-kem-768\",\"index\":0}]}",
                         ct_str);
                decaps_resp = dispatcher_handle_request(&dispatcher, decaps_req);
                check_condition("decapsulate via dispatcher returns result",
                                decaps_resp != NULL && response_has(decaps_resp, "\"result\""));
                check_condition("decapsulate result contains shared_secret",
                                decaps_resp != NULL && response_has(decaps_resp, "shared_secret"));

                /* Extract shared_secret from decaps result */
                if (decaps_resp != NULL) {
                    cJSON *decaps_json = cJSON_Parse(decaps_resp);
                    if (decaps_json != NULL) {
                        cJSON *d_result = cJSON_GetObjectItemCaseSensitive(decaps_json, "result");
                        if (cJSON_IsString(d_result)) {
                            cJSON *d_obj = cJSON_Parse(d_result->valuestring);
                            if (d_obj != NULL) {
                                cJSON *d_ss = cJSON_GetObjectItemCaseSensitive(d_obj, "shared_secret");
                                if (cJSON_IsString(d_ss)) ss_decaps_str = strdup(d_ss->valuestring);
                                cJSON_Delete(d_obj);
                            }
                        }
                        cJSON_Delete(decaps_json);
                    }
                }

                /* Check shared secrets match */
                if (ss_encaps_str != NULL && ss_decaps_str != NULL) {
                    check_condition("dispatcher encaps/decaps shared secrets match",
                                    strcmp(ss_encaps_str, ss_decaps_str) == 0);
                } else {
                    check_condition("dispatcher encaps/decaps shared secrets match", 0);
                }

                free((void *)ss_decaps_str);
                free(decaps_resp);
            } else {
                check_condition("decapsulate via dispatcher returns result", 0);
                check_condition("decapsulate result contains shared_secret", 0);
                check_condition("dispatcher encaps/decaps shared secrets match", 0);
            }

            free((void *)ct_str);
            free((void *)ss_encaps_str);
        }

        free(decaps_req);
        free(encaps_resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- Enforcement: encapsulate/decapsulate on pq-kem+ml-kem-768 ---- */
    {
        role_entry_t pq_kem = make_pq_kem_ml_kem_768_entry("pqk", 0);
        role_entry_t nostr_secp = make_nostr_secp_entry("nostr", 0);

        /* encapsulate/decapsulate are algorithm-based now. */
        check_condition("enforce encapsulate + ml-kem-768 -> OK",
                        enforce_verb_algorithm(VERB_ENCAPSULATE, CRYPTO_ALG_ML_KEM_768) == ENFORCE_OK);

        check_condition("enforce encapsulate + secp256k1 -> ALGORITHM err",
                        enforce_verb_algorithm(VERB_ENCAPSULATE, CRYPTO_ALG_SECP256K1) == ENFORCE_ERR_ALGORITHM);

        check_condition("enforce decapsulate + ml-kem-768 -> OK",
                        enforce_verb_algorithm(VERB_DECAPSULATE, CRYPTO_ALG_ML_KEM_768) == ENFORCE_OK);

        check_condition("enforce decapsulate + ed25519 -> ALGORITHM err",
                        enforce_verb_algorithm(VERB_DECAPSULATE, CRYPTO_ALG_ED25519) == ENFORCE_ERR_ALGORITHM);
    }

    /* ---- Cleanup ---- */
    free(priv1); free(pub1); free(priv2); free(pub2);
    free(ct1); free(ct2); free(ss1); free(ss2); free(ss3);

    printf("\n%d/%d tests passed\n", g_passes, g_total);
    return (g_passes == g_total) ? 0 : 1;
}