/*
 * test_dm.c — тест криптографии прямого чата (dm_crypto)
 *
 * Проверяет:
 *   - детерминированность и симметричность conv_id (одинаков на обеих сторонах);
 *   - симметричность content_key (ECDH: обе стороны выводят одинаковый ключ);
 *   - round-trip шифрования AES-256-CCM и детекцию неверного ключа.
 */

#include "../src/dm/dm_crypto.h"
#include "../src/transport_layer/secure_channel.h"
#include "../lib/debug_config.h"

#include <string.h>
#include <stdio.h>
#include "../lib/mem.h"

int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_ERROR);

    int failures = 0;

    /* ── 1. conv_id: симметричность и уникальность ── */
    uint64_t a = 0x1122334455667788ULL;
    uint64_t b = 0x8877665544332211ULL;
    uint64_t id_ab = dm_derive_conv_id(a, b);
    uint64_t id_ba = dm_derive_conv_id(b, a);
    if (id_ab == 0 || id_ab != id_ba) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "conv_id not symmetric: ab=%llu ba=%llu",
                    (unsigned long long)id_ab, (unsigned long long)id_ba);
        failures++;
    } else {
        DEBUG_INFO(DEBUG_CATEGORY_DM, "conv_id symmetric: %llu == %llu",
                   (unsigned long long)id_ab, (unsigned long long)id_ba);
    }

    uint64_t c = 0xDEADBEEF00000001ULL;
    if (dm_derive_conv_id(a, c) == id_ab) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "conv_id collision for different pairs");
        failures++;
    }

    /* ── 2. content_key: симметричность ECDH ── */
    struct SC_MYKEYS k1, k2;
    if (sc_generate_keypair(&k1) != SC_OK || sc_generate_keypair(&k2) != SC_OK) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "keypair generation failed");
        return 1;
    }
    uint8_t key1[DM_CONTENT_KEY_SIZE], key2[DM_CONTENT_KEY_SIZE];
    if (dm_derive_content_key(k1.private_key, k2.public_key, key1) != 0
        || dm_derive_content_key(k2.private_key, k1.public_key, key2) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "content_key derive failed");
        failures++;
    } else if (memcmp(key1, key2, DM_CONTENT_KEY_SIZE) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "content_key not symmetric (ECDH mismatch)");
        failures++;
    } else {
        DEBUG_INFO(DEBUG_CATEGORY_DM, "content_key symmetric (ECDH ok)");
    }

    /* ── 3. AES-256-CCM round-trip ── */
    const char* text = "hello, direct chat!";
    size_t text_len = strlen(text);
    uint8_t nonce[DM_NONCE_SIZE];
    dm_build_nonce(id_ab, a, 42, nonce);

    uint8_t enc[256];
    size_t enc_len = 0;
    if (dm_encrypt(key1, nonce, (const uint8_t*)text, text_len, enc, &enc_len) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "encrypt failed");
        failures++;
    } else if (enc_len != text_len + DM_TAG_SIZE) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "encrypt length mismatch: %zu", enc_len);
        failures++;
    } else {
        uint8_t dec[256];
        size_t dec_len = 0;
        if (dm_decrypt(key2, nonce, enc, enc_len, dec, &dec_len) != 0) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "decrypt failed");
            failures++;
        } else if (dec_len != text_len || memcmp(dec, text, text_len) != 0) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "decrypt content mismatch");
            failures++;
        } else {
            DEBUG_INFO(DEBUG_CATEGORY_DM, "AES-256-CCM round-trip ok");
        }
    }

    /* ── 4. детекция неверного ключа ── */
    {
        uint8_t wrong[DM_CONTENT_KEY_SIZE];
        memset(wrong, 0x5A, sizeof(wrong));
        uint8_t dec[256];
        size_t dec_len = 0;
        if (enc_len > 0 && dm_decrypt(wrong, nonce, enc, enc_len, dec, &dec_len) == 0) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "decrypt with wrong key should fail");
            failures++;
        } else {
            DEBUG_INFO(DEBUG_CATEGORY_DM, "wrong-key tamper detection ok");
        }
    }

    /* Границы CCM, включая пустое тело и порции медиа; повреждённый tag не принимается. */
    uint8_t* large = u_malloc(DM_CCM_PLAIN_MAX + DM_TAG_SIZE + 1);
    uint8_t* cipher = u_malloc(DM_CCM_PLAIN_MAX + DM_TAG_SIZE + 1);
    uint8_t* restored = u_malloc(DM_CCM_PLAIN_MAX + 1);
    if (!large || !cipher || !restored) return 1;
    memset(large, 0x91, DM_CCM_PLAIN_MAX + 1);
    const size_t sizes[] = {0, DM_MEDIA_PLAIN_SIZE, DM_CCM_PLAIN_MAX};
    for (size_t i = 0; i < sizeof(sizes) / sizeof(sizes[0]); i++) {
        size_t n = sizes[i], clen = 0, plen = 0;
        dm_build_nonce(id_ab, a, 100 + i, nonce);
        if (dm_encrypt(key1, nonce, n ? large : NULL, n, cipher, &clen) != 0 || clen != n + DM_TAG_SIZE ||
            dm_decrypt(key2, nonce, cipher, clen, restored, &plen) != 0 || plen != n || memcmp(large, restored, n)) failures++;
        cipher[clen - 1] ^= 1;
        if (dm_decrypt(key2, nonce, cipher, clen, restored, &plen) == 0 || plen != 0) failures++;
    }
    if (dm_encrypt(key1, nonce, large, DM_CCM_PLAIN_MAX + 1, cipher, &enc_len) == 0) failures++;
    uint8_t media_id[16] = {1}, media_nonce[DM_NONCE_SIZE], other_nonce[DM_NONCE_SIZE];
    dm_build_media_nonce(a, media_id, 0, media_nonce);
    dm_build_media_nonce(a, media_id, 1, other_nonce);
    if (!memcmp(media_nonce, other_nonce, DM_NONCE_SIZE)) failures++;
    dm_build_media_nonce(b, media_id, 0, other_nonce);
    if (!memcmp(media_nonce, other_nonce, DM_NONCE_SIZE)) failures++;
    media_id[1] = 1;
    dm_build_media_nonce(a, media_id, 0, other_nonce);
    if (!memcmp(media_nonce, other_nonce, DM_NONCE_SIZE)) failures++;
    dm_build_nonce(id_ab, a, 0, other_nonce);
    if (!memcmp(media_nonce, other_nonce, DM_NONCE_SIZE)) failures++;
    u_free(large);
    u_free(cipher);
    u_free(restored);

    if (failures == 0) {
        printf("TEST PASSED\n");
        return 0;
    }
    printf("TEST FAILED (%d failures)\n", failures);
    return 1;
}