/*
 * pq_kem_example.c — connect to a running n_signer over its abstract UNIX
 * socket and demonstrate post-quantum key encapsulation with ML-KEM-768.
 *
 * The example:
 *   1. Sends a get_public_key request for an ML-KEM-768 role ("kem_main").
 *   2. Prints the structured public key (algorithm, public_key, key_id).
 *   3. Sends a kem_encapsulate request with the public key, obtaining a
 *      ciphertext + shared secret.
 *   4. Sends a kem_decapsulate request with the ciphertext, recovering the
 *      shared secret on the signer side.
 *   5. Prints both shared secrets — they should match.
 *
 * Prerequisites:
 *   - n_signer must be running with a role configured for purpose=pq-kem,
 *     curve=ml-kem-768, named "kem_main" (or pass the role name as the 2nd arg).
 *   - A mnemonic must be loaded in the signer.
 *
 * Usage: ./pq_kem_example [socket_name] [role_name]
 *
 * Default socket_name: nsigner
 * Default role_name:   kem_main
 */

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

#include "nostr_common.h"
#include "nsigner_transport.h"
#include "nsigner_client.h"
#include "../cjson/cJSON.h"

static int get_structured_pubkey(nsigner_client_t *client, const char *role,
                                  char **out_pub_hex) {
    cJSON *params = NULL;
    cJSON *opts = NULL;
    cJSON *result = NULL;
    cJSON *parsed = NULL;
    int rc = -1;

    *out_pub_hex = NULL;

    params = cJSON_CreateArray();
    if (params == NULL) return -1;

    opts = cJSON_CreateObject();
    if (opts == NULL) {
        cJSON_Delete(params);
        return -1;
    }
    cJSON_AddStringToObject(opts, "algorithm", "ml-kem-768");
    cJSON_AddNumberToObject(opts, "index", 0);
    cJSON_AddItemToArray(params, opts);
    opts = NULL;

    if (nsigner_client_call(client, "get_public_key", params, &result) != NOSTR_SUCCESS) {
        cJSON_Delete(params);
        return -1;
    }
    params = NULL;

    if (cJSON_IsString(result)) {
        parsed = cJSON_Parse(result->valuestring);
        if (parsed != NULL) {
            cJSON *pk_item = cJSON_GetObjectItemCaseSensitive(parsed, "public_key");
            if (cJSON_IsString(pk_item)) {
                *out_pub_hex = strdup(pk_item->valuestring);
                rc = 0;
            }
        }
    }

    cJSON_Delete(parsed);
    cJSON_Delete(result);
    cJSON_Delete(params);
    return rc;
}

/* kem_encapsulate: returns ciphertext_hex and shared_secret_hex (newly
 * allocated, caller frees). */
static int kem_encapsulate(nsigner_client_t *client, const char *role,
                            const char *pub_hex,
                            char **out_ct_hex, char **out_ss_hex) {
    cJSON *params = NULL;
    cJSON *opts = NULL;
    cJSON *result = NULL;
    int rc = -1;

    *out_ct_hex = NULL;
    *out_ss_hex = NULL;

    params = cJSON_CreateArray();
    if (params == NULL) return -1;

    cJSON_AddItemToArray(params, cJSON_CreateString(pub_hex));

    opts = cJSON_CreateObject();
    if (opts == NULL) {
        cJSON_Delete(params);
        return -1;
    }
    cJSON_AddStringToObject(opts, "algorithm", "ml-kem-768");
    cJSON_AddNumberToObject(opts, "index", 0);
    cJSON_AddItemToArray(params, opts);
    opts = NULL;

    if (nsigner_client_call(client, "encapsulate", params, &result) != NOSTR_SUCCESS) {
        cJSON_Delete(params);
        return -1;
    }
    params = NULL;

    if (cJSON_IsString(result)) {
        cJSON *parsed = cJSON_Parse(result->valuestring);
        if (parsed != NULL) {
            cJSON *ct_item = cJSON_GetObjectItemCaseSensitive(parsed, "ciphertext");
            cJSON *ss_item = cJSON_GetObjectItemCaseSensitive(parsed, "shared_secret");
            if (cJSON_IsString(ct_item) && cJSON_IsString(ss_item)) {
                *out_ct_hex = strdup(ct_item->valuestring);
                *out_ss_hex = strdup(ss_item->valuestring);
                if (*out_ct_hex != NULL && *out_ss_hex != NULL) {
                    rc = 0;
                }
            }
            cJSON_Delete(parsed);
        }
    }

    cJSON_Delete(result);
    cJSON_Delete(params);
    return rc;
}

static char *kem_decapsulate(nsigner_client_t *client, const char *role,
                              const char *ct_hex) {
    cJSON *params = NULL;
    cJSON *opts = NULL;
    cJSON *result = NULL;
    char *ss_hex = NULL;

    params = cJSON_CreateArray();
    if (params == NULL) return NULL;

    cJSON_AddItemToArray(params, cJSON_CreateString(ct_hex));

    opts = cJSON_CreateObject();
    if (opts == NULL) {
        cJSON_Delete(params);
        return NULL;
    }
    cJSON_AddStringToObject(opts, "algorithm", "ml-kem-768");
    cJSON_AddNumberToObject(opts, "index", 0);
    cJSON_AddItemToArray(params, opts);
    opts = NULL;

    if (nsigner_client_call(client, "decapsulate", params, &result) != NOSTR_SUCCESS) {
        cJSON_Delete(params);
        return NULL;
    }
    params = NULL;

    if (cJSON_IsString(result)) {
        cJSON *parsed = cJSON_Parse(result->valuestring);
        if (parsed != NULL) {
            cJSON *ss_item = cJSON_GetObjectItemCaseSensitive(parsed, "shared_secret");
            if (cJSON_IsString(ss_item)) {
                ss_hex = strdup(ss_item->valuestring);
            }
            cJSON_Delete(parsed);
        }
    }

    cJSON_Delete(result);
    cJSON_Delete(params);
    return ss_hex;
}

int main(int argc, char **argv) {
    const char *socket_name = "nsigner";
    const char *role = "kem_main";
    nsigner_transport_t *transport = NULL;
    nsigner_client_t *client = NULL;
    char *pub_hex = NULL;
    char *ct_hex = NULL;
    char *encap_ss_hex = NULL;
    char *decap_ss_hex = NULL;
    int rc = 1;

    if (argc > 1 && argv[1] != NULL && argv[1][0] != '\0') {
        socket_name = argv[1];
    }
    if (argc > 2 && argv[2] != NULL && argv[2][0] != '\0') {
        role = argv[2];
    }

    if (nostr_init() != NOSTR_SUCCESS) {
        fprintf(stderr, "failed to initialize crypto subsystem\n");
        return 1;
    }

    transport = nsigner_transport_open_unix(socket_name, 10000);
    if (transport == NULL) {
        fprintf(stderr, "connect failed: cannot open unix transport @%s\n", socket_name);
        goto cleanup;
    }

    client = nsigner_client_new(transport);
    if (client == NULL) {
        fprintf(stderr, "connect failed: cannot create nsigner client\n");
        transport->close(transport);
        goto cleanup;
    }
    transport = NULL;

    printf("=== PQ KEM Example (ML-KEM-768) ===\n");
    printf("socket: %s\n", socket_name);
    printf("role:   %s\n", role);
    printf("\n");

    /* 1. Get the ML-KEM-768 public key. */
    if (get_structured_pubkey(client, role, &pub_hex) != 0 || pub_hex == NULL) {
        fprintf(stderr, "get_public_key failed: %s\n",
                nsigner_client_last_error(client));
        goto cleanup;
    }
    printf("Public Key:\n");
    printf("  pub_len: %zu hex chars (%zu bytes)\n",
           strlen(pub_hex), strlen(pub_hex) / 2);
    printf("  pub_head: %.64s...\n", pub_hex);
    printf("\n");

    /* 2. Encapsulate with the public key. */
    printf("Encapsulating with public key...\n");
    if (kem_encapsulate(client, role, pub_hex, &ct_hex, &encap_ss_hex) != 0) {
        fprintf(stderr, "kem_encapsulate failed: %s\n",
                nsigner_client_last_error(client));
        goto cleanup;
    }
    printf("Ciphertext:\n");
    printf("  ct_len: %zu hex chars (%zu bytes)\n",
           strlen(ct_hex), strlen(ct_hex) / 2);
    printf("  ct_head: %.64s...\n", ct_hex);
    printf("Encapsulated shared secret:\n");
    printf("  ss: %s\n", encap_ss_hex);
    printf("\n");

    /* 3. Decapsulate with the ciphertext (uses the role's private key). */
    printf("Decapsulating ciphertext on signer side...\n");
    decap_ss_hex = kem_decapsulate(client, role, ct_hex);
    if (decap_ss_hex == NULL) {
        fprintf(stderr, "kem_decapsulate failed: %s\n",
                nsigner_client_last_error(client));
        goto cleanup;
    }
    printf("Decapsulated shared secret:\n");
    printf("  ss: %s\n", decap_ss_hex);
    printf("\n");

    /* 4. Verify the shared secrets match. */
    if (strcmp(encap_ss_hex, decap_ss_hex) == 0) {
        printf("SUCCESS: shared secrets match!\n");
        rc = 0;
    } else {
        printf("FAILURE: shared secrets do NOT match!\n");
    }

cleanup:
    free(pub_hex);
    free(ct_hex);
    free(encap_ss_hex);
    free(decap_ss_hex);
    nsigner_client_free(client);
    nostr_cleanup();
    return rc;
}