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

#include "sc_lib.h"
#include <tinycrypt/ecc.h>
#include <tinycrypt/ecc_dh.h>
#include <tinycrypt/aes.h>
#include <tinycrypt/ccm_mode.h>
#include <tinycrypt/constants.h>
#include <tinycrypt/ecc_platform_specific.h>
#include <string.h>
#include <stddef.h>
#include <sys/types.h>
#include <unistd.h>
#include <sys/time.h>
#include <fcntl.h>

static const struct uECC_Curve_t *curve = NULL;

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();
    }
    return uECC_valid_public_key(public_key, curve);
}

sc_status_t sc_generate_keypair(sc_context_t *ctx)
{
    if (!ctx) {
        return SC_ERR_INVALID_ARG;
    }

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

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

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

    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_init_local_keys(sc_context_t *ctx,
                               const uint8_t *public_key,
                               const uint8_t *private_key)
{
    if (!ctx || !public_key || !private_key) {
        return SC_ERR_INVALID_ARG;
    }

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

    /* Validate public key */
    if (sc_validate_key(public_key) != 0) {
        return SC_ERR_INVALID_ARG;
    }

    memcpy(ctx->public_key, public_key, SC_PUBKEY_SIZE);
    memcpy(ctx->private_key, private_key, SC_PRIVKEY_SIZE);

    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 uint8_t *peer_public_key)
{
    uint8_t shared_secret[SC_SHARED_SECRET_SIZE];

    if (!ctx || !peer_public_key) {
        return SC_ERR_INVALID_ARG;
    }

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

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

    /* Validate peer public key */
    if (sc_validate_key(peer_public_key) != 0) {
        return SC_ERR_INVALID_ARG;
    }

    /* Compute shared secret using ECDH */
    if (!uECC_shared_secret(peer_public_key, ctx->private_key,
                            shared_secret, curve)) {
        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(uint32_t counter, uint8_t *nonce_out)
{
    /* Simple nonce construction: counter in little-endian + zeros */
    memset(nonce_out, 0, SC_NONCE_SIZE);
    nonce_out[0] = (counter >> 0) & 0xFF;
    nonce_out[1] = (counter >> 8) & 0xFF;
    nonce_out[2] = (counter >> 16) & 0xFF;
    nonce_out[3] = (counter >> 24) & 0xFF;
}

sc_status_t sc_encrypt(sc_context_t *ctx,
                       const uint8_t *plaintext,
                       size_t plaintext_len,
                       uint8_t *ciphertext,
                       uint8_t *tag)
{
    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;
    uint8_t combined_output[plaintext_len + SC_TAG_SIZE];

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

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

    if (plaintext_len == 0) {
        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->tx_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;
    }

    /* Encrypt and generate tag */
    if (tc_ccm_generation_encryption(combined_output, sizeof(combined_output),
                                     NULL, 0, /* no associated data */
                                     plaintext, plaintext_len,
                                     c) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_CRYPTO;
    }

    /* Split combined output into ciphertext and tag */
    memcpy(ciphertext, combined_output, plaintext_len);
    memcpy(tag, combined_output + 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,
                       const uint8_t *tag,
                       uint8_t *plaintext)
{
    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;
    uint8_t combined_input[ciphertext_len + SC_TAG_SIZE];

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

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

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

    /* Combine ciphertext and tag for CCM input */
    memcpy(combined_input, ciphertext, ciphertext_len);
    memcpy(combined_input + ciphertext_len, tag, SC_TAG_SIZE);

    /* 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, ciphertext_len,
                                       NULL, 0, /* no associated data */
                                       combined_input, ciphertext_len + SC_TAG_SIZE,
                                       c) != TC_CRYPTO_SUCCESS) {
        return SC_ERR_AUTH_FAILED;
    }

    ctx->rx_counter++;

    return SC_OK;
}