/* sc_lib.c - Secure Channel library implementation using TinyCrypt */

#include "secure_channel.h"
#include "../tinycrypt/lib/include/tinycrypt/ecc.h"
#include "../tinycrypt/lib/include/tinycrypt/ecc_dh.h"
#include "../tinycrypt/lib/include/tinycrypt/aes.h"
#include "../tinycrypt/lib/include/tinycrypt/ccm_mode.h"
#include "../tinycrypt/lib/include/tinycrypt/constants.h"
#include "../tinycrypt/lib/include/tinycrypt/ecc_platform_specific.h"
#include "../tinycrypt/lib/include/tinycrypt/sha256.h"
#include <string.h>
#include <stddef.h>
#include <sys/types.h>
#include <unistd.h>
#include <sys/time.h>
#include <stdio.h>

// Simple debug macros
#define DEBUG_CATEGORY_CRYPTO 1
#define DEBUG_ERROR(category, fmt, ...) fprintf(stderr, "ERROR: " fmt "\n", ##__VA_ARGS__)
#define DEBUG_INFO(category, fmt, ...) fprintf(stdout, "INFO: " fmt "\n", ##__VA_ARGS__)
#include <stdio.h>
#include <fcntl.h>
#include "crc32.h"

static const struct uECC_Curve_t *curve = NULL;
static uint8_t sc_urandom_seed[8] = {0};
static int sc_urandom_initialized = 0;

static void sc_init_random_seed(void)
{
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd >= 0) {
        ssize_t ret = read(fd, sc_urandom_seed, 8);
        close(fd);
        if (ret == 8) {
            sc_urandom_initialized = 1;
        }
    }
}


static int sc_rng(uint8_t *dest, unsigned size)
{
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0) {
        return 0;
    }

    ssize_t ret = read(fd, dest, size);
    close(fd);
    if (ret != size) {
        return 0;
    }

    /* Mix in PID and microtime for additional entropy */
    pid_t pid = getpid();
    struct timeval tv;
    gettimeofday(&tv, NULL);

    for (unsigned i = 0; i < size; i++) {
        dest[i] ^= ((pid >> (i % (sizeof(pid) * 8))) & 0xFF);
        dest[i] ^= ((tv.tv_sec >> (i % (sizeof(tv.tv_sec) * 8))) & 0xFF);
        dest[i] ^= ((tv.tv_usec >> (i % (sizeof(tv.tv_usec) * 8))) & 0xFF);
    }

    return 1;
}

static int sc_validate_key(const uint8_t *public_key)
{
    if (!curve) {
        curve = uECC_secp256r1();
    }
    int result = uECC_valid_public_key(public_key, curve);
    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "sc_validate_key: uECC_valid_public_key returned %d", result);
    return result;
}

sc_status_t sc_generate_keypair(struct SC_MYKEYS *pk)
{
    if (!pk) {
        return SC_ERR_INVALID_ARG;
    }

    if (!curve) {
        curve = uECC_secp256r1();
    }

    /* Set custom RNG function */
    uECC_set_rng(sc_rng);

    if (!uECC_make_key(pk->public_key, pk->private_key, curve)) {
        return SC_ERR_CRYPTO;
    }
    return SC_OK;
}

// Конвертация hex строки в бинарный формат
static int hex_to_binary(const char *hex_str, uint8_t *binary, size_t binary_len) {
    if (!hex_str || !binary || strlen(hex_str) != binary_len * 2) return -1;
    
    for (size_t i = 0; i < binary_len; i++) {
        unsigned int byte;
        if (sscanf(hex_str + i * 2, "%2x", &byte) != 1) return -1;
        binary[i] = (uint8_t)byte;
    }
    return 0;
}

sc_status_t sc_init_local_keys(struct SC_MYKEYS *mykeys, const char *public_key, const char *private_key) {
    if (!mykeys || !public_key || !private_key) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: invalid arguments");
        return SC_ERR_INVALID_ARG;
    }

    if (!curve) {
        curve = uECC_secp256r1();
    }

    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: public_key len=%zu, private_key len=%zu", 
               strlen(public_key), strlen(private_key));

    /* Convert hex to binary first */
    if (hex_to_binary(public_key, mykeys->public_key, SC_PUBKEY_SIZE)) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: failed to convert public key from hex");
        return SC_ERR_INVALID_ARG;
    }
    if (hex_to_binary(private_key, mykeys->private_key, SC_PRIVKEY_SIZE)) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: failed to convert private key from hex");
        return SC_ERR_INVALID_ARG;
    }

    /* Validate the converted binary public key */
    if (sc_validate_key(mykeys->public_key) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: public key validation failed");
        return SC_ERR_INVALID_ARG;
    }

    DEBUG_INFO(DEBUG_CATEGORY_CRYPTO, "sc_init_local_keys: keys initialized successfully");
    return SC_OK;
}

sc_status_t sc_init_ctx(sc_context_t *ctx, struct SC_MYKEYS *mykeys) {

    ctx->pk=mykeys;
    ctx->initialized = 1;
    ctx->peer_key_set = 0;
    ctx->session_ready = 0;
    ctx->tx_counter = 0;
    ctx->rx_counter = 0;

    return SC_OK;
}

sc_status_t sc_set_peer_public_key(sc_context_t *ctx, const char *peer_public_key_h, int mode) {
    uint8_t shared_secret[SC_SHARED_SECRET_SIZE];
    uint8_t peer_public_key[SC_PUBKEY_SIZE];

    if (mode) {
      if (hex_to_binary(peer_public_key_h, peer_public_key, SC_PUBKEY_SIZE)) {
          DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: invalid hex key format");
          return SC_ERR_INVALID_ARG;
          }
      }
    else memcpy(peer_public_key, peer_public_key_h, SC_PUBKEY_SIZE);

    if (!ctx) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: invalid ctx");
        return SC_ERR_INVALID_ARG;
    }

    if (!ctx->initialized) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: ctx not initialized");
        return SC_ERR_NOT_INITIALIZED;
    }

    if (!curve) {
        curve = uECC_secp256r1();
    }

    /* Validate peer public key */
    if (sc_validate_key(peer_public_key) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: invalid key");
        return SC_ERR_INVALID_ARG;
    }

    /* Compute shared secret using ECDH */
    if (!ctx->pk) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: no private key");
        return SC_ERR_NOT_INITIALIZED;
    }
    if (!uECC_shared_secret(peer_public_key, ctx->pk->private_key,
                            shared_secret, curve)) {
        DEBUG_ERROR(DEBUG_CATEGORY_CRYPTO, "sc_set_peer_public_key: shared secret error");
        return SC_ERR_CRYPTO;
    }

    /* Derive session key from shared secret (simple copy for demo) */
    memcpy(ctx->session_key, shared_secret, SC_SESSION_KEY_SIZE);

    /* Store peer public key */
    memcpy(ctx->peer_public_key, peer_public_key, SC_PUBKEY_SIZE);
    ctx->peer_key_set = 1;

    ctx->session_ready = 1;

    return SC_OK;
}

static void sc_build_nonce(uint64_t counter, uint8_t *nonce_out)
{
    struct tc_sha256_state_struct sha_ctx;
    uint8_t hash[32];
    struct timeval tv;
    uint8_t data[8 + 8 + 4];

    if (!sc_urandom_initialized) {
        sc_init_random_seed();
    }

    gettimeofday(&tv, NULL);

    memcpy(data, sc_urandom_seed, 8);
    data[8] = (counter >> 0) & 0xFF;
    data[9] = (counter >> 8) & 0xFF;
    data[10] = (counter >> 16) & 0xFF;
    data[11] = (counter >> 24) & 0xFF;
    data[12] = (counter >> 32) & 0xFF;
    data[13] = (counter >> 40) & 0xFF;
    data[14] = (counter >> 48) & 0xFF;
    data[15] = (counter >> 56) & 0xFF;
    data[16] = (tv.tv_sec >> 0) & 0xFF;
    data[17] = (tv.tv_sec >> 8) & 0xFF;
    data[18] = (tv.tv_sec >> 16) & 0xFF;
    data[19] = (tv.tv_sec >> 24) & 0xFF;

    tc_sha256_init(&sha_ctx);
    tc_sha256_update(&sha_ctx, data, 20);
    tc_sha256_final(hash, &sha_ctx);

    memcpy(nonce_out, hash, SC_NONCE_SIZE);
}

sc_status_t sc_encrypt(sc_context_t *ctx,
                       const uint8_t *plaintext,
                       size_t plaintext_len,
                       uint8_t *ciphertext,
                       size_t *ciphertext_len)
{
    uint8_t nonce[SC_NONCE_SIZE];
    struct tc_aes_key_sched_struct sched;
    struct tc_ccm_mode_struct ccm_state;
    size_t total_plaintext_len = plaintext_len + SC_CRC32_SIZE;
    uint8_t plaintext_with_crc[total_plaintext_len];
    uint8_t combined_output[total_plaintext_len + SC_TAG_SIZE];

    if (!ctx || !plaintext || !ciphertext || !ciphertext_len) {
        return SC_ERR_INVALID_ARG;
    }

    if (!ctx->session_ready) {
        return SC_ERR_NOT_INITIALIZED;
    }

    if (plaintext_len == 0) {
        return SC_ERR_INVALID_ARG;
    }

    /* Добавляем CRC32 к данным */
    memcpy(plaintext_with_crc, plaintext, plaintext_len);
    uint32_t crc = crc32_calc(plaintext, plaintext_len);
    plaintext_with_crc[plaintext_len] = (crc >> 0) & 0xFF;
    plaintext_with_crc[plaintext_len + 1] = (crc >> 8) & 0xFF;
    plaintext_with_crc[plaintext_len + 2] = (crc >> 16) & 0xFF;
    plaintext_with_crc[plaintext_len + 3] = (crc >> 24) & 0xFF;

    /* Initialize AES key schedule */
    if (tc_aes128_set_encrypt_key(&sched, ctx->session_key) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Build nonce from counter */
    sc_build_nonce(ctx->tx_counter, nonce);

    /* Configure CCM mode */
    if (tc_ccm_config(&ccm_state, &sched, nonce, SC_NONCE_SIZE, SC_TAG_SIZE) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Encrypt and generate tag */
    if (tc_ccm_generation_encryption(combined_output, sizeof(combined_output),
                                     NULL, 0, /* no associated data */
                                     plaintext_with_crc, total_plaintext_len,
                                     &ccm_state) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Copy ciphertext + tag to output buffer */
    memcpy(ciphertext, combined_output, total_plaintext_len + SC_TAG_SIZE);
    *ciphertext_len = total_plaintext_len + SC_TAG_SIZE;

    ctx->tx_counter++;

    return SC_OK;
}

sc_status_t sc_decrypt(sc_context_t *ctx,
                       const uint8_t *ciphertext,
                       size_t ciphertext_len,
                       uint8_t *plaintext,
                       size_t *plaintext_len)
{
    uint8_t nonce[SC_NONCE_SIZE];
    struct tc_aes_key_sched_struct sched;
    struct tc_ccm_mode_struct ccm_state;
    TCCcmMode_t c = &ccm_state;
    size_t total_plaintext_len = ciphertext_len - SC_TAG_SIZE;
    uint8_t plaintext_with_crc[total_plaintext_len];

    if (!ctx || !ciphertext || !plaintext || !plaintext_len) {
        return SC_ERR_INVALID_ARG;
    }

    if (!ctx->session_ready) {
        return SC_ERR_NOT_INITIALIZED;
    }

    if (ciphertext_len < SC_TAG_SIZE + SC_CRC32_SIZE) {
        return SC_ERR_INVALID_ARG;
    }

    /* Initialize AES key schedule */
    if (tc_aes128_set_encrypt_key(&sched, ctx->session_key) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Build nonce from counter */
    sc_build_nonce(ctx->rx_counter, nonce);

    /* Configure CCM mode */
    if (tc_ccm_config(c, &sched, nonce, SC_NONCE_SIZE, SC_TAG_SIZE) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Decrypt and verify tag */
    if (tc_ccm_decryption_verification(plaintext_with_crc, total_plaintext_len,
                                       NULL, 0, /* no associated data */
                                       ciphertext, ciphertext_len,
                                       c) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_AUTH_FAILED;
    }

    /* Проверяем CRC32 */
    size_t data_len = total_plaintext_len - SC_CRC32_SIZE;
    uint32_t expected_crc = crc32_calc(plaintext_with_crc, data_len);
    uint32_t received_crc = (plaintext_with_crc[data_len] << 0) |
                           (plaintext_with_crc[data_len + 1] << 8) |
                           (plaintext_with_crc[data_len + 2] << 16) |
                           (plaintext_with_crc[data_len + 3] << 24);
    
    if (expected_crc != received_crc) {
        return SC_ERR_CRC_FAILED;
    }

    /* Копируем данные без CRC32 */
    memcpy(plaintext, plaintext_with_crc, data_len);
    *plaintext_len = data_len;

    ctx->rx_counter++;

    return SC_OK;
}

sc_status_t sc_compute_public_key_from_private(const uint8_t *private_key, uint8_t *public_key) {
    if (!private_key || !public_key) {
        return SC_ERR_INVALID_ARG;
    }

    if (!curve) {
        curve = uECC_secp256r1();
    }

    if (!uECC_compute_public_key(private_key, public_key, curve)) {
        return SC_ERR_CRYPTO;
    }
    return SC_OK;
}