#pragma once

// Shared AES-256-GCM frame envelope — the single wire format used by the
// BLE variables/live channels AND the ESP-NOW mesh payloads. Mirrors
// ble_frame.py on the host exactly:
//
//   [0]              tag_len
//   [1..tag_len]     device tag, e.g. "M5Stack|AA:BB:CC:DD:EE:FF" (also AAD)
//   [+12]            nonce (random, per frame)
//   [...]            ciphertext + 16-byte GCM tag
//
// Extracted from BLEManager so the mesh layer can encrypt under a
// per-session group key while BLE keeps using the per-device key — the
// only difference between the callers is which 32-byte key they pass in.

#include <Arduino.h>
#include <mbedtls/gcm.h>
#include <esp_random.h>
#include "config.h"

static constexpr size_t FRAME_NONCE_LEN = 12;
static constexpr size_t FRAME_GCM_TAG_LEN = 16;
static constexpr size_t FRAME_KEY_LEN = 32;  // AES-256

// Encrypt `plaintext` under `key` with `tag` as both header and AAD.
inline bool frameCryptoBuild(const uint8_t* key, const char* tag,
                             const uint8_t* plaintext, size_t plainLen,
                             uint8_t* out, size_t outCap, size_t* outLen) {
    if (!key || !tag) return false;
    size_t tagStrLen = strlen(tag);
    if (tagStrLen == 0 || tagStrLen > 255) return false;

    size_t total = 1 + tagStrLen + FRAME_NONCE_LEN + plainLen + FRAME_GCM_TAG_LEN;
    if (total > outCap) return false;

    out[0] = (uint8_t)tagStrLen;
    memcpy(out + 1, tag, tagStrLen);
    uint8_t* nonce = out + 1 + tagStrLen;
    uint8_t* ct = nonce + FRAME_NONCE_LEN;
    uint8_t* gcmTag = ct + plainLen;
    // 12-byte nonce sourced from esp_random (CSPRNG).
    for (size_t i = 0; i < FRAME_NONCE_LEN; i += 4) {
        uint32_t r = esp_random();
        for (size_t b = 0; b < 4 && i + b < FRAME_NONCE_LEN; b++) {
            nonce[i + b] = (uint8_t)(r >> (b * 8));
        }
    }

    mbedtls_gcm_context gcm;
    mbedtls_gcm_init(&gcm);
    int rc = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, key,
                                FRAME_KEY_LEN * 8);
    if (rc == 0) {
        rc = mbedtls_gcm_crypt_and_tag(
            &gcm, MBEDTLS_GCM_ENCRYPT, plainLen,
            nonce, FRAME_NONCE_LEN,
            (const uint8_t*)tag, tagStrLen,
            plaintext, ct,
            FRAME_GCM_TAG_LEN, gcmTag);
    }
    mbedtls_gcm_free(&gcm);
    if (rc != 0) return false;
    *outLen = total;
    return true;
}

// Verify and decrypt an inbound frame under `key`. Writes the recovered
// tag (NUL-terminated) and plaintext into the caller's buffers. Returns
// false silently on any malformed or auth-failed input.
inline bool frameCryptoParse(const uint8_t* key,
                             const uint8_t* in, size_t inLen,
                             char* outTag, size_t outTagCap,
                             uint8_t* outPlain, size_t outPlainCap,
                             size_t* outPlainLen) {
    if (!key) return false;
    if (inLen < 1 + FRAME_NONCE_LEN + FRAME_GCM_TAG_LEN) return false;
    size_t tagLen = in[0];
    if (tagLen == 0 || tagLen >= outTagCap) return false;
    if (inLen < 1 + tagLen + FRAME_NONCE_LEN + FRAME_GCM_TAG_LEN) return false;
    memcpy(outTag, in + 1, tagLen);
    outTag[tagLen] = '\0';
    // Reject anything not starting with our prefix early so we don't
    // burn cycles on adversarial input.
    if (strncmp(outTag, BLE_DEVICE_TAG_PREFIX,
                sizeof(BLE_DEVICE_TAG_PREFIX) - 1) != 0) {
        return false;
    }

    const uint8_t* nonce = in + 1 + tagLen;
    size_t ctLen = inLen - 1 - tagLen - FRAME_NONCE_LEN - FRAME_GCM_TAG_LEN;
    if (ctLen >= outPlainCap) return false;
    const uint8_t* ct = nonce + FRAME_NONCE_LEN;
    const uint8_t* gcmTag = ct + ctLen;

    mbedtls_gcm_context gcm;
    mbedtls_gcm_init(&gcm);
    int rc = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, key,
                                FRAME_KEY_LEN * 8);
    if (rc == 0) {
        rc = mbedtls_gcm_auth_decrypt(&gcm, ctLen,
                                      nonce, FRAME_NONCE_LEN,
                                      (const uint8_t*)outTag, tagLen,
                                      gcmTag, FRAME_GCM_TAG_LEN,
                                      ct, outPlain);
    }
    mbedtls_gcm_free(&gcm);
    if (rc != 0) return false;
    *outPlainLen = ctLen;
    return true;
}