/**
 * @file crypto.c
 * @brief AES-256-GCM Crypto Operations (mbedtls)
 */

#include <string.h>
#include <stdlib.h>
#include "crypto.h"
#include "esp_log.h"

#include "mbedtls/sha256.h"
#include "mbedtls/cipher.h"
#include "mbedtls/error.h"

static const char *TAG = "nexus_crypto";

/* AES-256-GCM constants */
#define GCM_NONCE_SIZE   12
#define GCM_TAG_SIZE     16
#define KEY_SIZE         32

esp_err_t crypto_init(void)
{
    /* mbedtls auto-initializes on first use, nothing needed */
    ESP_LOGI(TAG, "Crypto module initialized (mbedtls)");
    return ESP_OK;
}

esp_err_t crypto_derive_key(const char *hwid, const char *salt, uint8_t key[32])
{
    if (hwid == NULL || salt == NULL || key == NULL) {
        return ESP_ERR_INVALID_ARG;
    }

    /* Build input: hwid + salt (concatenation, NOT HMAC) */
    size_t hwid_len = strlen(hwid);
    size_t salt_len = strlen(salt);
    size_t input_len = hwid_len + salt_len;

    char *input = malloc(input_len);
    if (input == NULL) {
        ESP_LOGE(TAG, "Failed to allocate key derivation input");
        return ESP_ERR_NO_MEM;
    }

    memcpy(input, hwid, hwid_len);
    memcpy(input + hwid_len, salt, salt_len);

    /* SHA-256 hash */
    int ret = mbedtls_sha256((const unsigned char *)input, input_len, key, 0);
    free(input);

    if (ret != 0) {
        ESP_LOGE(TAG, "SHA-256 failed: -0x%04x", -ret);
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "Key derived from HWID+salt");
    return ESP_OK;
}

esp_err_t crypto_decrypt(const uint8_t *ciphertext, size_t len,
                        const uint8_t key[32], uint8_t **plaintext,
                        size_t *plaintext_len)
{
    if (ciphertext == NULL || key == NULL || plaintext == NULL || plaintext_len == NULL) {
        return ESP_ERR_INVALID_ARG;
    }

    /* Minimum: nonce(12) + tag(16) = 28 bytes minimum */
    if (len < GCM_NONCE_SIZE + GCM_TAG_SIZE) {
        ESP_LOGE(TAG, "Ciphertext too short: %zu bytes (min %d)", len, GCM_NONCE_SIZE + GCM_TAG_SIZE);
        return ESP_ERR_INVALID_SIZE;
    }

    /* Split: nonce(12) + encrypted_data + tag(16) */
    const uint8_t *nonce = ciphertext;
    const uint8_t *encrypted = ciphertext + GCM_NONCE_SIZE;
    size_t encrypted_len = len - GCM_NONCE_SIZE - GCM_TAG_SIZE;
    const uint8_t *tag = ciphertext + len - GCM_TAG_SIZE;

    /* Allocate output buffer */
    uint8_t *out = malloc(encrypted_len + 1); /* +1 for null terminator (JSON payload) */
    if (out == NULL) {
        ESP_LOGE(TAG, "Failed to allocate plaintext buffer");
        return ESP_ERR_NO_MEM;
    }

    /* Initialize mbedtls cipher context */
    mbedtls_cipher_context_t ctx;
    mbedtls_cipher_init(&ctx);

    ret = mbedtls_cipher_setup(&ctx, mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_256_GCM));
    if (ret != 0) {
        ESP_LOGE(TAG, "Cipher setup failed: -0x%04x", -ret);
        mbedtls_cipher_free(&ctx);
        free(out);
        return ESP_FAIL;
    }

    ret = mbedtls_cipher_setkey(&ctx, key, KEY_SIZE * 8, MBEDTLS_DECRYPT);
    if (ret != 0) {
        ESP_LOGE(TAG, "Cipher setkey failed: -0x%04x", -ret);
        mbedtls_cipher_free(&ctx);
        free(out);
        return ESP_FAIL;
    }

    ret = mbedtls_cipher_set_iv(&ctx, nonce, GCM_NONCE_SIZE);
    if (ret != 0) {
        ESP_LOGE(TAG, "Cipher setiv failed: -0x%04x", -ret);
        mbedtls_cipher_free(&ctx);
        free(out);
        return ESP_FAIL;
    }

    /* Start decrypt */
    ret = mbedtls_cipher_auth_decrypt(&ctx, tag, GCM_TAG_SIZE,
                                       NULL, 0,
                                       encrypted, encrypted_len,
                                       out, plaintext_len);
    mbedtls_cipher_free(&ctx);

    if (ret != 0) {
        ESP_LOGE(TAG, "AES-256-GCM decrypt failed: -0x%04x", -ret);
        free(out);
        return ESP_ERR_AUTH_CRYPTO;
    }

    /* Null-terminate for JSON parsing */
    out[*plaintext_len] = '\0';
    *plaintext = out;

    ESP_LOGI(TAG, "Decrypted %zu bytes", *plaintext_len);
    return ESP_OK;
}