#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/provider.h>

// --- HARDCODED ML-KEM-768 Public Key (Replace with a real key) ---
// A real ML-KEM-768 public key is 608 bytes long.
// This is a placeholder for demonstration.
unsigned char server_public_key_bytes[] = {
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    // ... 592 more bytes of a real ML-KEM-768 public key would go here ...
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    // Note: This placeholder is NOT a valid ML-KEM-768 key and will likely fail.
    // Replace with a 608-byte key generated via 'openssl genpkey' for real use.
};
size_t public_key_len = 608; // ML-KEM-768 Public Key size

// Error handling macro
#define HANDLE_ERROR(msg) \
    { \
        fprintf(stderr, "%s failed.\n", msg); \
        ERR_print_errors_fp(stderr); \
        goto cleanup; \
    }

int main(void) {
    EVP_PKEY_CTX *kctx = NULL;
    EVP_PKEY *peer_pub_key = NULL;
    OSSL_PARAM params[2];
    unsigned char *shared_secret = NULL;
    unsigned char *ciphertext = NULL;
    size_t shared_secret_len = 0;
    size_t ciphertext_len = 0;
    int ret = 0;

    printf("Starting ML-KEM Client Key Encapsulation (ML-KEM-768).\n");

    // 1. Load the OpenSSL default provider (ML-KEM is in the default provider since 3.5)
    if (!OSSL_PROVIDER_load(NULL, "default")) {
        HANDLE_ERROR("Failed to load default provider");
    }

    // 2. Import the raw public key bytes into an EVP_PKEY structure
    printf("2. Importing ML-KEM-768 Public Key...\n");

    // a. Create parameter array to describe the key components
    params[0] = OSSL_PARAM_construct_octet_string("pub", server_public_key_bytes, public_key_len);
    params[1] = OSSL_PARAM_construct_end();

    // b. Import the key. ML-KEM-768 is named "ML-KEM-768"
    peer_pub_key = EVP_PKEY_new_from_params(NULL, params, "ML-KEM-768");
    if (!peer_pub_key) {
        // NOTE: This will fail if the hardcoded bytes are not a valid key.
        HANDLE_ERROR("EVP_PKEY_new_from_params failed (Public Key Import)");
    }

    // 3. Initialize the Key Encapsulation Context
    kctx = EVP_PKEY_CTX_new_from_pkey(NULL, peer_pub_key, NULL);
    if (!kctx) {
        HANDLE_ERROR("EVP_PKEY_CTX_new_from_pkey failed");
    }

    if (EVP_PKEY_encapsulate_init(kctx) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate_init failed");
    }

    // 4. Determine buffer sizes
    printf("4. Determining buffer sizes...\n");
    if (EVP_PKEY_encapsulate(kctx, NULL, &shared_secret_len, NULL, &ciphertext_len) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate (size determination) failed");
    }
    
    printf("   - Shared Secret Length: %zu bytes\n", shared_secret_len);
    printf("   - Ciphertext Length: %zu bytes\n", ciphertext_len);

    // 5. Allocate buffers
    shared_secret = OPENSSL_malloc(shared_secret_len);
    ciphertext = OPENSSL_malloc(ciphertext_len);

    if (!shared_secret || !ciphertext) {
        HANDLE_ERROR("Memory allocation failed");
    }

    // 6. Perform the Encapsulation (Client's Key Exchange Step)
    printf("6. Performing Key Encapsulation...\n");
    if (EVP_PKEY_encapsulate(kctx, shared_secret, &shared_secret_len, ciphertext, &ciphertext_len) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate failed");
    }

    printf("✅ Encapsulation successful!\n");
    printf("   The client has generated a **Shared Secret** and the **Ciphertext**.\n");
    
    printf("   - Shared Secret (First 16 bytes): ");
    for (size_t i = 0; i < (shared_secret_len > 16 ? 16 : shared_secret_len); i++) {
        printf("%02x", shared_secret[i]);
    }
    printf("...\n");

    printf("   - Ciphertext (First 16 bytes, to be sent to Server): ");
    for (size_t i = 0; i < (ciphertext_len > 16 ? 16 : ciphertext_len); i++) {
        printf("%02x", ciphertext[i]);
    }
    printf("...\n");

    ret = 0; // Success
    
cleanup:
    // Free all allocated resources
    EVP_PKEY_CTX_free(kctx);
    EVP_PKEY_free(peer_pub_key);
    OPENSSL_free(shared_secret);
    OPENSSL_free(ciphertext);
    OSSL_PROVIDER_unload(OSSL_PROVIDER_load(NULL, "default")); // Unload provider

    if (ret != 0) {
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}