/* Test for Phase 6 get_public_key response format (dispatcher.c).
 *
 * Verifies:
 *   - secp256k1: get_public_key returns plain hex string (backward compat)
 *   - secp256k1 with {"format":"structured"}: returns structured JSON
 *   - ed25519: returns structured JSON with algorithm "ed25519"
 *   - ML-DSA-65: returns structured JSON with algorithm "ml-dsa-65" and large pubkey
 *   - ML-KEM-768: returns structured JSON with algorithm "ml-kem-768"
 *   - Structured result contains algorithm, public_key, and key_id fields
 *   - key_id is the first 16 hex chars of the public key
 */
/* 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_entry_t *role_table_find_by_nostr_index(role_table_t *table, int index);
role_entry_t *role_table_find_by_path(role_table_t *table, const char *path);
role_entry_t *role_table_get_default(role_table_t *table);
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);
int role_table_register_nostr_index(role_table_t *table, int nostr_index);


/* 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 VERB_SIGN_EVENT     "sign_event"
#define VERB_GET_PUBLIC_KEY "get_public_key"
#define VERB_NIP44_ENCRYPT  "nip44_encrypt"
#define VERB_NIP44_DECRYPT  "nip44_decrypt"
#define VERB_NIP04_ENCRYPT  "nip04_encrypt"
#define VERB_NIP04_DECRYPT  "nip04_decrypt"
#define VERB_MINE_EVENT     "mine_event"
#define VERB_SIGN_DATA      "sign_data"
#define VERB_VERIFY_SIG     "verify_signature"
#define VERB_SSH_SIGN       "ssh_sign"
#define VERB_KEM_ENCAPS     "kem_encapsulate"
#define VERB_KEM_DECAPS     "kem_decapsulate"

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

/* Check whether caller_id is allowed to invoke `verb` with the given
 * algorithm and index (algorithm-based policy). Returns POLICY_ALLOW,
 * POLICY_DENY, POLICY_PROMPT, or POLICY_NO_MATCH. */
int policy_check_algorithm(const policy_table_t *table, const char *caller_id,
                           const char *verb, const char *algorithm, int index,
                           policy_source_t *out_source);

/* Check whether caller_id is allowed to invoke `verb` with the given
 * algorithm and index (algorithm-based policy). Returns POLICY_ALLOW,
 * POLICY_DENY, POLICY_PROMPT, or POLICY_NO_MATCH. */
int policy_check_algorithm(const policy_table_t *table, const char *caller_id,
                           const char *verb, const char *algorithm, int index,
                           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);
int crypto_ml_kem_768_keygen_from_seed(const unsigned char *seed, size_t seed_len,
                                         unsigned char *priv_out, unsigned char *pub_out);
int crypto_ml_kem_768_encaps(const unsigned char *pub, size_t pub_len,
                               unsigned char *ct_out, unsigned char *ss_out);
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>

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);
}

/* Parse the dispatcher response's "result" field. If result is a JSON string
 * (plain hex), returns a newly-allocated copy of the string in *out_str and
 * returns 1. If result is a JSON object serialized as a string, parses it into
 * *out_obj and returns 2. Returns 0 on failure. Caller frees accordingly. */
static int parse_result(const char *response, char **out_str, cJSON **out_obj) {
    cJSON *root = NULL;
    cJSON *result_item = NULL;
    int ret = 0;

    if (out_str != NULL) *out_str = NULL;
    if (out_obj != NULL) *out_obj = NULL;

    if (response == NULL) return 0;
    root = cJSON_Parse(response);
    if (root == NULL) return 0;
    result_item = cJSON_GetObjectItemCaseSensitive(root, "result");
    if (cJSON_IsString(result_item) && result_item->valuestring != NULL) {
        /* Could be a plain hex string OR a serialized JSON object string. */
        cJSON *maybe_obj = cJSON_Parse(result_item->valuestring);
        if (maybe_obj != NULL && cJSON_IsObject(maybe_obj)) {
            if (out_obj != NULL) {
                *out_obj = maybe_obj;
                maybe_obj = NULL;
                ret = 2;
            } else {
                cJSON_Delete(maybe_obj);
                ret = 0;
            }
        } else {
            cJSON_Delete(maybe_obj);
            if (out_str != NULL) {
                *out_str = strdup(result_item->valuestring);
                ret = 1;
            }
        }
    }
    cJSON_Delete(root);
    return ret;
}

static role_entry_t make_entry(const char *name, const char *purpose_str,
                                const char *curve_str, int idx) {
    role_entry_t e;
    memset(&e, 0, sizeof(e));
    strncpy(e.name, name, sizeof(e.name) - 1);
    strncpy(e.purpose_str, purpose_str, sizeof(e.purpose_str) - 1);
    strncpy(e.curve_str, curve_str, 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";

    if (nostr_init() != 0) {
        fprintf(stderr, "nostr_init failed\n");
        return 1;
    }

    /* ---- secp256k1: plain hex (backward compat) ---- */
    {
        role_table_t table;
        role_entry_t nostr_role;
        static mnemonic_state_t mnemonic_state;
            dispatcher_ctx_t dispatcher;
        char *resp = NULL;
        char *result_str = NULL;
        cJSON *result_obj = NULL;
        int pr;

        role_table_init(&table);
        nostr_role = make_entry("main", "nostr", "secp256k1", 0);
        role_table_add(&table, &nostr_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);

        resp = dispatcher_handle_request(&dispatcher,
                                          "{\"id\":\"1\",\"method\":\"get_public_key\",\"params\":[{\"role\":\"main\"}]}");
        check_condition("secp256k1 get_public_key returns a response", resp != NULL);

        pr = parse_result(resp, &result_str, &result_obj);
        check_condition("secp256k1 get_public_key result is a plain string (not object)",
                        pr == 1 && result_str != NULL);
        check_condition("secp256k1 plain hex result is 64 hex chars",
                        result_str != NULL && strlen(result_str) == 64);
        check_condition("secp256k1 plain hex result has no 'algorithm' field",
                        !response_has(resp, "algorithm"));

        free(result_str);
        cJSON_Delete(result_obj);
        free(resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- secp256k1 with format:structured ---- */
    {
        role_table_t table;
        role_entry_t nostr_role;
        static mnemonic_state_t mnemonic_state;
        dispatcher_ctx_t dispatcher;
        char *resp = NULL;
        char *result_str = NULL;
        cJSON *result_obj = NULL;
        cJSON *alg_item = NULL;
        cJSON *pk_item = NULL;
        cJSON *kid_item = NULL;
        int pr;

        role_table_init(&table);
        nostr_role = make_entry("main", "nostr", "secp256k1", 0);
        role_table_add(&table, &nostr_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);

        resp = dispatcher_handle_request(&dispatcher,
                                          "{\"id\":\"2\",\"method\":\"get_public_key\",\"params\":[{\"role\":\"main\",\"format\":\"structured\"}]}");
        check_condition("secp256k1 structured get_public_key returns a response", resp != NULL);

        pr = parse_result(resp, &result_str, &result_obj);
        check_condition("secp256k1 structured result is a JSON object",
                        pr == 2 && result_obj != NULL);

        if (result_obj != NULL) {
            alg_item = cJSON_GetObjectItemCaseSensitive(result_obj, "algorithm");
            pk_item = cJSON_GetObjectItemCaseSensitive(result_obj, "public_key");
            kid_item = cJSON_GetObjectItemCaseSensitive(result_obj, "key_id");
            check_condition("secp256k1 structured has algorithm=secp256k1",
                            cJSON_IsString(alg_item) && strcmp(alg_item->valuestring, "secp256k1") == 0);
            check_condition("secp256k1 structured has public_key (64 hex chars)",
                            cJSON_IsString(pk_item) && strlen(pk_item->valuestring) == 64);
            check_condition("secp256k1 structured has key_id (16 hex chars)",
                            cJSON_IsString(kid_item) && strlen(kid_item->valuestring) == 16);
            check_condition("secp256k1 key_id == first 16 chars of public_key",
                            cJSON_IsString(pk_item) && cJSON_IsString(kid_item) &&
                            strncmp(pk_item->valuestring, kid_item->valuestring, 16) == 0);
        } else {
            check_condition("secp256k1 structured has algorithm=secp256k1", 0);
            check_condition("secp256k1 structured has public_key (64 hex chars)", 0);
            check_condition("secp256k1 structured has key_id (16 hex chars)", 0);
            check_condition("secp256k1 key_id == first 16 chars of public_key", 0);
        }

        free(result_str);
        cJSON_Delete(result_obj);
        free(resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- ed25519: structured JSON ---- */
    {
        role_table_t table;
        role_entry_t ssh_role;
        static mnemonic_state_t mnemonic_state;
        dispatcher_ctx_t dispatcher;
        char *resp = NULL;
        char *result_str = NULL;
        cJSON *result_obj = NULL;
        cJSON *alg_item = NULL;
        cJSON *pk_item = NULL;
        cJSON *kid_item = NULL;
        int pr;

        role_table_init(&table);
        ssh_role = make_entry("ssh_main", "ssh", "ed25519", 0);
        role_table_add(&table, &ssh_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);

        resp = dispatcher_handle_request(&dispatcher,
                                          "{\"id\":\"3\",\"method\":\"get_public_key\",\"params\":[{\"role\":\"ssh_main\"}]}");
        check_condition("ed25519 get_public_key returns a response", resp != NULL);

        pr = parse_result(resp, &result_str, &result_obj);
        check_condition("ed25519 get_public_key result is a JSON object",
                        pr == 2 && result_obj != NULL);

        if (result_obj != NULL) {
            alg_item = cJSON_GetObjectItemCaseSensitive(result_obj, "algorithm");
            pk_item = cJSON_GetObjectItemCaseSensitive(result_obj, "public_key");
            kid_item = cJSON_GetObjectItemCaseSensitive(result_obj, "key_id");
            check_condition("ed25519 structured has algorithm=ed25519",
                            cJSON_IsString(alg_item) && strcmp(alg_item->valuestring, "ed25519") == 0);
            check_condition("ed25519 structured has public_key (64 hex chars)",
                            cJSON_IsString(pk_item) && strlen(pk_item->valuestring) == 64);
            check_condition("ed25519 structured has key_id (16 hex chars)",
                            cJSON_IsString(kid_item) && strlen(kid_item->valuestring) == 16);
            check_condition("ed25519 key_id == first 16 chars of public_key",
                            cJSON_IsString(pk_item) && cJSON_IsString(kid_item) &&
                            strncmp(pk_item->valuestring, kid_item->valuestring, 16) == 0);
        } else {
            check_condition("ed25519 structured has algorithm=ed25519", 0);
            check_condition("ed25519 structured has public_key (64 hex chars)", 0);
            check_condition("ed25519 structured has key_id (16 hex chars)", 0);
            check_condition("ed25519 key_id == first 16 chars of public_key", 0);
        }

        free(result_str);
        cJSON_Delete(result_obj);
        free(resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- ML-DSA-65: structured JSON with large pubkey ---- */
    {
        role_table_t table;
        role_entry_t pq_role;
        static mnemonic_state_t mnemonic_state;
        dispatcher_ctx_t dispatcher;
        char *resp = NULL;
        char *result_str = NULL;
        cJSON *result_obj = NULL;
        cJSON *alg_item = NULL;
        cJSON *pk_item = NULL;
        cJSON *kid_item = NULL;
        const crypto_alg_sizes_t *sz;
        int pr;

        role_table_init(&table);
        pq_role = make_entry("pq_sig", "pq-sig", "ml-dsa-65", 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);

        sz = crypto_alg_get_sizes(CRYPTO_ALG_ML_DSA_65);
        check_condition("ML-DSA-65 sizes available", sz != NULL);

        resp = dispatcher_handle_request(&dispatcher,
                                          "{\"id\":\"4\",\"method\":\"get_public_key\",\"params\":[{\"role\":\"pq_sig\"}]}");
        check_condition("ML-DSA-65 get_public_key returns a response", resp != NULL);

        pr = parse_result(resp, &result_str, &result_obj);
        check_condition("ML-DSA-65 get_public_key result is a JSON object",
                        pr == 2 && result_obj != NULL);

        if (result_obj != NULL && sz != NULL) {
            alg_item = cJSON_GetObjectItemCaseSensitive(result_obj, "algorithm");
            pk_item = cJSON_GetObjectItemCaseSensitive(result_obj, "public_key");
            kid_item = cJSON_GetObjectItemCaseSensitive(result_obj, "key_id");
            check_condition("ML-DSA-65 structured has algorithm=ml-dsa-65",
                            cJSON_IsString(alg_item) && strcmp(alg_item->valuestring, "ml-dsa-65") == 0);
            check_condition("ML-DSA-65 structured has large public_key (3904 hex chars)",
                            cJSON_IsString(pk_item) && strlen(pk_item->valuestring) == sz->pub_key_len * 2);
            check_condition("ML-DSA-65 structured has key_id (16 hex chars)",
                            cJSON_IsString(kid_item) && strlen(kid_item->valuestring) == 16);
            check_condition("ML-DSA-65 key_id == first 16 chars of public_key",
                            cJSON_IsString(pk_item) && cJSON_IsString(kid_item) &&
                            strncmp(pk_item->valuestring, kid_item->valuestring, 16) == 0);
        } else {
            check_condition("ML-DSA-65 structured has algorithm=ml-dsa-65", 0);
            check_condition("ML-DSA-65 structured has large public_key (3904 hex chars)", 0);
            check_condition("ML-DSA-65 structured has key_id (16 hex chars)", 0);
            check_condition("ML-DSA-65 key_id == first 16 chars of public_key", 0);
        }

        free(result_str);
        cJSON_Delete(result_obj);
        free(resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

    /* ---- ML-KEM-768: structured JSON ---- */
    {
        role_table_t table;
        role_entry_t kem_role;
        static mnemonic_state_t mnemonic_state;
        dispatcher_ctx_t dispatcher;
        char *resp = NULL;
        char *result_str = NULL;
        cJSON *result_obj = NULL;
        cJSON *alg_item = NULL;
        cJSON *pk_item = NULL;
        cJSON *kid_item = NULL;
        const crypto_alg_sizes_t *sz;
        int pr;

        role_table_init(&table);
        kem_role = make_entry("kem_main", "pq-kem", "ml-kem-768", 0);
        role_table_add(&table, &kem_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);

        sz = crypto_alg_get_sizes(CRYPTO_ALG_ML_KEM_768);
        check_condition("ML-KEM-768 sizes available", sz != NULL);

        resp = dispatcher_handle_request(&dispatcher,
                                          "{\"id\":\"5\",\"method\":\"get_public_key\",\"params\":[{\"role\":\"kem_main\"}]}");
        check_condition("ML-KEM-768 get_public_key returns a response", resp != NULL);

        pr = parse_result(resp, &result_str, &result_obj);
        check_condition("ML-KEM-768 get_public_key result is a JSON object",
                        pr == 2 && result_obj != NULL);

        if (result_obj != NULL && sz != NULL) {
            alg_item = cJSON_GetObjectItemCaseSensitive(result_obj, "algorithm");
            pk_item = cJSON_GetObjectItemCaseSensitive(result_obj, "public_key");
            kid_item = cJSON_GetObjectItemCaseSensitive(result_obj, "key_id");
            check_condition("ML-KEM-768 structured has algorithm=ml-kem-768",
                            cJSON_IsString(alg_item) && strcmp(alg_item->valuestring, "ml-kem-768") == 0);
            check_condition("ML-KEM-768 structured has public_key (2368 hex chars)",
                            cJSON_IsString(pk_item) && strlen(pk_item->valuestring) == sz->pub_key_len * 2);
            check_condition("ML-KEM-768 structured has key_id (16 hex chars)",
                            cJSON_IsString(kid_item) && strlen(kid_item->valuestring) == 16);
            check_condition("ML-KEM-768 key_id == first 16 chars of public_key",
                            cJSON_IsString(pk_item) && cJSON_IsString(kid_item) &&
                            strncmp(pk_item->valuestring, kid_item->valuestring, 16) == 0);
        } else {
            check_condition("ML-KEM-768 structured has algorithm=ml-kem-768", 0);
            check_condition("ML-KEM-768 structured has public_key (2368 hex chars)", 0);
            check_condition("ML-KEM-768 structured has key_id (16 hex chars)", 0);
            check_condition("ML-KEM-768 key_id == first 16 chars of public_key", 0);
        }

        free(result_str);
        cJSON_Delete(result_obj);
        free(resp);
        crypto_wipe(&g_key_store);
        mnemonic_unload(&mnemonic_state);
    }

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