#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdarg.h>
#include "mbedtls/x509_crt.h"
#include "mbedtls/platform_time.h"
#include "mbedtls/gcm.h"
#include "mbedtls/ssl_ciphersuites.h"
#include "mbedtls/ecdh.h"
#include "mbedtls/ctr_drbg.h"
#include "talkingq_wss.h"

typedef union { size_t bytes; long double align; void *ptr; } alloc_header;
static size_t live_bytes, peak_bytes;
void *mbedtls_calloc(size_t n, size_t size) {
    alloc_header *h;
    size_t bytes;
    if (size && n > (SIZE_MAX - sizeof(*h)) / size) return NULL;
    bytes = n * size;
    h = (alloc_header *)calloc(1, sizeof(*h) + bytes);
    if (!h) return NULL;
    h->bytes = bytes; live_bytes += bytes;
    if (live_bytes > peak_bytes) peak_bytes = live_bytes;
    return h + 1;
}
void mbedtls_free(void *p) {
    if (p) { alloc_header *h = (alloc_header *)p - 1; live_bytes -= h->bytes; free(h); }
}
static time_t test_time;
static time_t read_test_time(time_t *p) { if (p) *p = test_time; return test_time; }
mbedtls_time_t (*mbedtls_time)(mbedtls_time_t *) = read_test_time;
void CPUartLogPrintf(const char *format, ...) { (void)format; }
int mbedtls_snprintf(char *buffer, size_t size, const char *format, ...)
{
    int result;
    va_list args;
    va_start(args, format);
    result = vsnprintf(buffer, size, format, args);
    va_end(args);
    return result;
}

/* Deterministic entropy is confined to this host benchmark, never firmware. */
static int test_entropy(void *ctx, unsigned char *out, size_t n)
{
    size_t i; unsigned int *counter = (unsigned int *)ctx;
    for (i = 0; i < n; ++i) out[i] = (unsigned char)(++*counter);
    return 0;
}

static void benchmark_ecdh(mbedtls_ecp_group_id id, const char *name)
{
    mbedtls_ecp_group peer, client;
    mbedtls_ecp_point peer_q, client_q;
    mbedtls_mpi peer_d, client_d, secret, expected;
    mbedtls_ctr_drbg_context rng;
    unsigned int counter = 0;
    clock_t elapsed = 0, started;
    size_t base;
    int i;
    mbedtls_ctr_drbg_init(&rng);
    assert(mbedtls_ctr_drbg_seed(&rng, test_entropy, &counter, NULL, 0) == 0);
    mbedtls_ecp_group_init(&peer); mbedtls_ecp_point_init(&peer_q); mbedtls_mpi_init(&peer_d);
    assert(mbedtls_ecp_group_load(&peer, id) == 0);
    assert(mbedtls_ecdh_gen_public(&peer, &peer_d, &peer_q, mbedtls_ctr_drbg_random, &rng) == 0);
    base = live_bytes; peak_bytes = base;
    for (i = 0; i < 20; ++i) {
        mbedtls_ecp_group_init(&client); mbedtls_ecp_point_init(&client_q);
        mbedtls_mpi_init(&client_d); mbedtls_mpi_init(&secret); mbedtls_mpi_init(&expected);
        assert(mbedtls_ecp_group_load(&client, id) == 0);
        started = clock();
        assert(mbedtls_ecdh_gen_public(&client, &client_d, &client_q, mbedtls_ctr_drbg_random, &rng) == 0);
        assert(mbedtls_ecdh_compute_shared(&client, &secret, &peer_q, &client_d, mbedtls_ctr_drbg_random, &rng) == 0);
        elapsed += clock() - started;
        assert(mbedtls_ecdh_compute_shared(&peer, &expected, &client_q, &peer_d, mbedtls_ctr_drbg_random, &rng) == 0);
        assert(mbedtls_mpi_cmp_mpi(&secret, &expected) == 0);
        mbedtls_mpi_free(&expected); mbedtls_mpi_free(&secret); mbedtls_mpi_free(&client_d);
        mbedtls_ecp_point_free(&client_q); mbedtls_ecp_group_free(&client);
        assert(live_bytes == base);
    }
    printf("Host fresh ECDH %s: %.2f ms, extra peak heap %zu bytes (20 runs, includes peer check heap)\n",
           name, 1000.0 * elapsed / CLOCKS_PER_SEC / 20, peak_bytes - base);
    mbedtls_ecp_group_free(&peer); mbedtls_ecp_point_free(&peer_q); mbedtls_mpi_free(&peer_d);
    mbedtls_ctr_drbg_free(&rng);
}

int main(int argc, char **argv)
{
    mbedtls_x509_crt root, chain;
    mbedtls_gcm_context gcm;
    uint32_t flags;
    int ret, i;
    char path[1024];
    char pem[16384];
    FILE *file;
    size_t length;
    clock_t verify_started;
    mbedtls_ecp_group group;
    unsigned char key[16] = {0}, iv[12] = {0}, plain[16] = {0}, cipher[16], tag[16];
    const unsigned char expected_cipher[16] = {0x03,0x88,0xda,0xce,0x60,0xb6,0xa3,0x92,0xf3,0x28,0xc2,0xb9,0x71,0xb2,0xfe,0x78};
    const unsigned char expected_tag[16] = {0xab,0x6e,0x47,0xd4,0x2c,0xec,0x13,0xbd,0xf5,0x3a,0x67,0xb2,0x12,0x57,0xbd,0xdf};
    assert(argc == 2);
    test_time = time(NULL);
    mbedtls_x509_crt_init(&root);
    mbedtls_x509_crt_init(&chain);
    ret = mbedtls_x509_crt_parse(&root, (const unsigned char *)talkingq_root_ca, talkingq_root_ca_size);
    printf("SDK root CA parse: %d (previous device failure: -9774)\n", ret);
    assert(ret == 0);
    for (i = 0; i < 3; ++i) {
        snprintf(path, sizeof(path), "%s/chain-%d.pem", argv[1], i);
        file = fopen(path, "rb");
        assert(file != NULL);
        length = fread(pem, 1, sizeof(pem)-1, file);
        assert(feof(file) && !ferror(file));
        fclose(file);
        pem[length] = 0;
        ret = mbedtls_x509_crt_parse(&chain, (const unsigned char *)pem, length + 1);
        assert(ret == 0);
    }
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    printf("SDK server chain + hostname verify: %d, flags=0x%08x\n", ret, (unsigned int)flags);
    assert(ret == 0 && flags == 0);
    verify_started = clock();
    peak_bytes = live_bytes;
    for (i = 0; i < 20; ++i) {
        ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
        assert(ret == 0 && flags == 0);
    }
    printf("Host certificate chain verify average: %.2f ms (20 runs)\n",
           1000.0 * (clock() - verify_started) / CLOCKS_PER_SEC / 20);
    printf("Host certificate verify peak heap: %zu bytes, ECP window: %d\n", peak_bytes, MBEDTLS_ECP_WINDOW_SIZE);
    mbedtls_ecp_group_init(&group);
    assert(mbedtls_ecp_group_load(&group, MBEDTLS_ECP_DP_SECP384R1) == 0);
#ifdef MBEDTLS_ECP_NIST_OPTIM
    assert(group.modp != NULL);
#else
    assert(group.modp == NULL);
#endif
    mbedtls_ecp_group_free(&group);
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "wrong.invalid", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_CN_MISMATCH));
    ret = mbedtls_x509_crt_verify(&chain, NULL, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_NOT_TRUSTED));
    test_time = 1900000000; /* 2030: this server certificate is expired. */
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_EXPIRED));
    test_time = 1700000000; /* 2023: this server certificate is not yet valid. */
    ret = mbedtls_x509_crt_verify(&chain, &root, NULL, "banban.api.talkingq.com", &flags, NULL, NULL, 0);
    assert(ret != 0 && (flags & MBEDTLS_X509_BADCERT_FUTURE));
    assert(mbedtls_ssl_ciphersuite_from_id(MBEDTLS_TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384) != NULL);
    mbedtls_gcm_init(&gcm);
    assert(mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, key, 128) == 0);
    assert(mbedtls_gcm_crypt_and_tag(&gcm, MBEDTLS_GCM_ENCRYPT, 16, iv, 12, NULL, 0,
                                   plain, cipher, 16, tag) == 0);
    assert(!memcmp(cipher, expected_cipher, 16) && !memcmp(tag, expected_tag, 16));
    mbedtls_gcm_free(&gcm);
    mbedtls_x509_crt_free(&chain);
    mbedtls_x509_crt_free(&root);
    assert(live_bytes == 0);
    benchmark_ecdh(MBEDTLS_ECP_DP_CURVE25519, "X25519");
    benchmark_ecdh(MBEDTLS_ECP_DP_SECP256R1, "P-256");
    benchmark_ecdh(MBEDTLS_ECP_DP_SECP384R1, "P-384");
    benchmark_ecdh(MBEDTLS_ECP_DP_SECP521R1, "P-521");
    assert(live_bytes == 0);
    puts("PASS: SDK ECDSA certificate chain, hostname, trust, dates, TLS suite and AES-GCM vector");
    return 0;
}