#include "servergprs.h"

#include "psa_crypto_driver_esp_aes_gcm.h"

char linebufIn[1024];
char *linebufInPtr = linebufIn;

#define MAX_FRAME 512

static uint8_t rx_buf[MAX_FRAME];
static size_t rx_len = 0;
static uint16_t expected_len = 0;
static int state = 0;

void process_encrypted_frame(uint8_t *data, size_t len)
{
    // layout:
    // [4B device_id][4B counter][12B nonce][ciphertext...][16B tag]

    if (len < 36) return;

    uint8_t *device_id = data;
    uint8_t *counter   = data + 4;
    uint8_t *nonce     = data + 8;
    uint8_t *cipher    = data + 20;
    size_t cipher_len  = len - 36;
    uint8_t *tag       = data + len - 16;

    uint8_t plaintext[512];
    size_t plaintext_len = 0;

    psa_key_attributes_t attr = PSA_KEY_ATTRIBUTES_INIT;

    psa_status_t st = esp_crypto_aes_gcm_decrypt(
        &attr,
        shared_key,
        sizeof(shared_key),
        PSA_ALG_GCM,
        nonce, 12,
        NULL, 0,
        cipher, cipher_len,
        plaintext, sizeof(plaintext),
        &plaintext_len
    );

    if (st != PSA_SUCCESS) {
        return;
    }

    handle_message(plaintext, plaintext_len);
}

void send_secure_packet(uint8_t *msg, size_t msg_len)
{
    uint8_t nonce[12];
    generate_random(nonce, 12);

    uint8_t cipher[512];
    size_t cipher_len = 0;

    uint8_t tag[16];
    size_t tag_len = 0;

    esp_crypto_aes_gcm_encrypt(
        &attr,
        shared_key,
        sizeof(shared_key),
        PSA_ALG_GCM,
        nonce, 12,
        NULL, 0,
        msg, msg_len,
        cipher, sizeof(cipher),
        &cipher_len
    );

    // build frame:
    // [LEN][NONCE][CIPHER][TAG]

    uint16_t total =
        12 + cipher_len + 16;

    uint8_t out[600];

    out[0] = (total >> 8) & 0xFF;
    out[1] = total & 0xFF;

    memcpy(out + 2, nonce, 12);
    memcpy(out + 14, cipher, cipher_len);
    memcpy(out + 14 + cipher_len, tag, 16);

    sim800_tcp_send(out, total + 2);
}


void tcp_on_byte(char c)
{
    switch (state)
    {
        // WAIT FOR LENGTH (2 bytes)
        case 0:
            rx_buf[rx_len++] = c;
            if (rx_len == 2)
            {
                expected_len = (rx_buf[0] << 8) | rx_buf[1];
                rx_len = 0;

                if (expected_len > MAX_FRAME) {
                    rx_len = 0;
                    state = 0;
                    return;
                }

                state = 1;
            }
            break;

        // READ FRAME
        case 1:
            rx_buf[rx_len++] = c;

            if (rx_len >= expected_len)
            {
                process_encrypted_frame(rx_buf, rx_len);

                rx_len = 0;
                state = 0;
            }
            break;
    }
}