// crypto.c – ChaCha20-Poly1305 AEAD with HKDF key derivation
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "esp_log.h"
#include "esp_random.h"
#include "nvs_flash.h"
#include "nvs.h"

#include "mbedtls/chachapoly.h"
#include "mbedtls/hkdf.h"
#include "mbedtls/md.h"
#include "mbedtls/base64.h"
#include "mbedtls/platform_util.h"

#include "pb_decode.h"
#include "c2.pb.h"

#include "utils.h"
#include "command.h"

static const char *TAG = "CRYPTO";

#define NONCE_LEN   12
#define TAG_LEN     16
#define KEY_LEN     32
#define OVERHEAD    (NONCE_LEN + TAG_LEN)   /* 28 bytes */

static uint8_t derived_key[KEY_LEN];
static bool    crypto_ready = false;

/* ============================================================
 * crypto_init – read master key from factory NVS, derive via HKDF
 * ============================================================ */
bool crypto_init(void)
{
    esp_err_t err;

    /* 1) Init the factory NVS partition */
    err = nvs_flash_init_partition("fctry");
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "nvs_flash_init_partition(fctry) failed: %s",
                 esp_err_to_name(err));
        return false;
    }

    /* 2) Open the crypto namespace (read-only) */
    nvs_handle_t handle;
    err = nvs_open_from_partition(
        "fctry",
        CONFIG_CRYPTO_FCTRY_NS,
        NVS_READONLY,
        &handle
    );
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "nvs_open_from_partition(fctry/%s) failed: %s",
                 CONFIG_CRYPTO_FCTRY_NS, esp_err_to_name(err));
        return false;
    }

    /* 3) Read the 32-byte master key blob */
    uint8_t master_key[KEY_LEN];
    size_t  mk_len = sizeof(master_key);

    err = nvs_get_blob(handle, CONFIG_CRYPTO_FCTRY_KEY, master_key, &mk_len);
    nvs_close(handle);

    if (err != ESP_OK || mk_len != KEY_LEN) {
        ESP_LOGE(TAG, "nvs_get_blob(%s) failed: %s (len=%u)",
                 CONFIG_CRYPTO_FCTRY_KEY, esp_err_to_name(err),
                 (unsigned)mk_len);
        mbedtls_platform_zeroize(master_key, sizeof(master_key));
        return false;
    }

    /* 4) HKDF-SHA256: derive the encryption key */
    const char *info = "espilon-c2-v1";
    const char *salt = CONFIG_DEVICE_ID;

    int ret = mbedtls_hkdf(
        mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
        (const uint8_t *)salt, strlen(salt),
        master_key, KEY_LEN,
        (const uint8_t *)info, strlen(info),
        derived_key, KEY_LEN
    );

    /* Wipe master key from RAM immediately */
    mbedtls_platform_zeroize(master_key, sizeof(master_key));

    if (ret != 0) {
        ESP_LOGE(TAG, "HKDF failed (%d)", ret);
        return false;
    }

    crypto_ready = true;
    ESP_LOGI(TAG, "Crypto ready (ChaCha20-Poly1305 + HKDF)");
    return true;
}

/* ============================================================
 * crypto_encrypt – ChaCha20-Poly1305 AEAD
 *
 * Output layout: nonce[12] || ciphertext[plain_len] || tag[16]
 * Returns total output length, or -1 on error.
 * ============================================================ */
int crypto_encrypt(const uint8_t *plain, size_t plain_len,
                   uint8_t *out, size_t out_cap)
{
    if (!crypto_ready) {
        ESP_LOGE(TAG, "crypto_encrypt: not initialized");
        return -1;
    }
    if (!plain || plain_len == 0 || !out) {
        ESP_LOGE(TAG, "crypto_encrypt: invalid args");
        return -1;
    }

    size_t needed = plain_len + OVERHEAD;
    if (out_cap < needed) {
        ESP_LOGE(TAG, "crypto_encrypt: buffer too small (%u < %u)",
                 (unsigned)out_cap, (unsigned)needed);
        return -1;
    }

    /* Random nonce in the first 12 bytes */
    esp_fill_random(out, NONCE_LEN);

    mbedtls_chachapoly_context ctx;
    mbedtls_chachapoly_init(&ctx);
    mbedtls_chachapoly_setkey(&ctx, derived_key);

    int ret = mbedtls_chachapoly_encrypt_and_tag(
        &ctx,
        plain_len,
        out,                        /* nonce  */
        NULL, 0,                    /* no AAD */
        plain,                      /* input  */
        out + NONCE_LEN,            /* output (ciphertext) */
        out + NONCE_LEN + plain_len /* tag    */
    );

    mbedtls_chachapoly_free(&ctx);

    if (ret != 0) {
        ESP_LOGE(TAG, "chachapoly encrypt failed (%d)", ret);
        return -1;
    }

    return (int)needed;
}

/* ============================================================
 * crypto_decrypt – ChaCha20-Poly1305 AEAD
 *
 * Input layout: nonce[12] || ciphertext[N] || tag[16]
 * Returns plaintext length, or -1 on error / auth failure.
 * ============================================================ */
int crypto_decrypt(const uint8_t *in, size_t in_len,
                   uint8_t *out, size_t out_cap)
{
    if (!crypto_ready) {
        ESP_LOGE(TAG, "crypto_decrypt: not initialized");
        return -1;
    }
    if (!in || in_len < OVERHEAD || !out) {
        ESP_LOGE(TAG, "crypto_decrypt: invalid args (in_len=%u)",
                 (unsigned)in_len);
        return -1;
    }

    size_t ct_len = in_len - OVERHEAD;
    if (out_cap < ct_len) {
        ESP_LOGE(TAG, "crypto_decrypt: buffer too small");
        return -1;
    }

    const uint8_t *nonce = in;
    const uint8_t *ct    = in + NONCE_LEN;
    const uint8_t *tag   = in + NONCE_LEN + ct_len;

    mbedtls_chachapoly_context ctx;
    mbedtls_chachapoly_init(&ctx);
    mbedtls_chachapoly_setkey(&ctx, derived_key);

    int ret = mbedtls_chachapoly_auth_decrypt(
        &ctx,
        ct_len,
        nonce,
        NULL, 0,    /* no AAD */
        tag,
        ct,
        out
    );

    mbedtls_chachapoly_free(&ctx);

    if (ret != 0) {
        ESP_LOGE(TAG, "AEAD auth/decrypt failed (%d)", ret);
        return -1;
    }

    return (int)ct_len;
}

/* ============================================================
 * Base64 encode
 * ============================================================ */
char *base64_encode(const unsigned char *input, size_t input_len)
{
    if (!input || input_len == 0) {
        ESP_LOGE(TAG, "Invalid input to base64_encode");
        return NULL;
    }

    size_t out_len = 4 * ((input_len + 2) / 3);
    char *out = (char *)malloc(out_len + 1);
    if (!out) {
        ESP_LOGE(TAG, "malloc failed in base64_encode");
        return NULL;
    }

    size_t written = 0;
    int ret = mbedtls_base64_encode(
        (unsigned char *)out,
        out_len + 1,
        &written,
        input,
        input_len
    );

    if (ret != 0) {
        ESP_LOGE(TAG, "base64 encode failed (%d)", ret);
        free(out);
        return NULL;
    }

    out[written] = '\0';
    return out;
}

/* ============================================================
 * Base64 decode
 * ============================================================ */
char *base64_decode(const char *input, size_t *output_len)
{
    if (!input || !output_len) {
        ESP_LOGE(TAG, "Invalid input to base64_decode");
        return NULL;
    }

    size_t in_len = strlen(input);
    size_t est_len = (in_len * 3) / 4;

    unsigned char *out = (unsigned char *)malloc(est_len + 1);
    if (!out) {
        ESP_LOGE(TAG, "malloc failed in base64_decode");
        return NULL;
    }

    int ret = mbedtls_base64_decode(
        out,
        est_len + 1,
        output_len,
        (const unsigned char *)input,
        in_len
    );

    if (ret != 0) {
        ESP_LOGE(TAG, "base64 decode failed (%d)", ret);
        free(out);
        return NULL;
    }

    out[*output_len] = '\0';
    return (char *)out;
}

/* ============================================================
 * C2: Decode + decrypt + protobuf + exec (COMMON WIFI/GPRS)
 * ============================================================ */
bool c2_decode_and_exec(const char *frame)
{
    if (!frame || !frame[0]) {
        ESP_LOGW(TAG, "Empty C2 frame");
        return false;
    }

    /* Trim CR/LF/spaces at end (SIM800 sometimes adds \r) */
    char tmp[1024];
    size_t n = strnlen(frame, sizeof(tmp) - 1);
    memcpy(tmp, frame, n);
    tmp[n] = '\0';
    while (n > 0 && (tmp[n - 1] == '\r' || tmp[n - 1] == '\n' || tmp[n - 1] == ' ')) {
        tmp[--n] = '\0';
    }

    ESP_LOGD(TAG, "C2 RX b64 (%u bytes)", (unsigned)n);

    /* 1) Base64 decode */
    size_t decoded_len = 0;
    char *decoded = base64_decode(tmp, &decoded_len);
    if (!decoded || decoded_len == 0) {
        ESP_LOGE(TAG, "Base64 decode failed");
        free(decoded);
        return false;
    }

    /* 2) Decrypt + authenticate (AEAD) */
    uint8_t plain[1024];
    int plain_len = crypto_decrypt(
        (const uint8_t *)decoded, decoded_len,
        plain, sizeof(plain)
    );
    free(decoded);

    if (plain_len < 0) {
        ESP_LOGE(TAG, "Decrypt/auth failed – tampered or wrong key");
        return false;
    }

    /* 3) Protobuf decode -> c2_Command */
    c2_Command cmd = c2_Command_init_zero;
    pb_istream_t is = pb_istream_from_buffer(plain, (size_t)plain_len);

    if (!pb_decode(&is, c2_Command_fields, &cmd)) {
        ESP_LOGE(TAG, "PB decode error: %s", PB_GET_ERROR(&is));
        return false;
    }

    /* 4) Log + dispatch */
#ifdef CONFIG_ESPILON_LOG_C2_VERBOSE
    ESP_LOGI(TAG, "==== C2 COMMAND ====");
    ESP_LOGI(TAG, "name: %s", cmd.command_name);
    ESP_LOGI(TAG, "argc: %d", cmd.argv_count);
    if (cmd.request_id[0]) ESP_LOGI(TAG, "req : %s", cmd.request_id);
    for (int i = 0; i < cmd.argv_count; i++) {
        ESP_LOGI(TAG, "arg[%d]=%s", i, cmd.argv[i]);
    }
    ESP_LOGI(TAG, "====================");
#else
    ESP_LOGI(
        TAG,
        "C2 CMD: %s argc=%d req=%s",
        cmd.command_name,
        cmd.argv_count,
        cmd.request_id[0] ? cmd.request_id : "-"
    );
    for (int i = 0; i < cmd.argv_count; i++) {
        ESP_LOGD(TAG, "arg[%d]=%s", i, cmd.argv[i]);
    }
#endif

    process_command(&cmd);
    return true;
}