/* uart_ota_recv.c — IO board OTA receiver (bd-21hv)
 *
 * Listens on UART0 for OTA frames from Balance board.
 * Writes incoming chunks to the inactive OTA partition, verifies SHA256,
 * then reboots into new firmware.
 */

#include "uart_ota_recv.h"
#include "config.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "driver/uart.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "mbedtls/sha256.h"
#include <string.h>

static const char *TAG = "io_ota";

volatile io_ota_state_t g_io_ota_state    = IO_OTA_IDLE;
volatile uint8_t        g_io_ota_progress = 0;

/* Frame type bytes (same as uart_ota.h sender side) */
#define OTA_BEGIN   0xC0u
#define OTA_DATA    0xC1u
#define OTA_END     0xC2u
#define OTA_ABORT   0xC3u
#define OTA_ACK     0xC4u
#define OTA_NACK    0xC5u

#define CHUNK_MAX   1024

static uint8_t crc8(const uint8_t *d, uint16_t len)
{
    uint8_t crc = 0;
    for (uint16_t i = 0; i < len; i++) {
        crc ^= d[i];
        for (uint8_t b = 0; b < 8; b++)
            crc = (crc & 0x80u) ? (uint8_t)((crc << 1u) ^ 0x07u) : (uint8_t)(crc << 1u);
    }
    return crc;
}

static void send_ack(uint16_t seq)
{
    uint8_t frame[6];
    frame[0] = OTA_ACK;
    frame[1] = (uint8_t)(seq >> 8u);
    frame[2] = (uint8_t)(seq);
    frame[3] = 0; frame[4] = 0;  /* LEN=0 */
    uint8_t crc = crc8(frame, 5);
    frame[5] = crc;
    uart_write_bytes(IO_UART_PORT, (char *)frame, 6);
}

static void send_nack(uint16_t seq, uint8_t err)
{
    uint8_t frame[8];
    frame[0] = OTA_NACK;
    frame[1] = (uint8_t)(seq >> 8u);
    frame[2] = (uint8_t)(seq);
    frame[3] = 0; frame[4] = 1;  /* LEN=1 */
    frame[5] = err;
    uint8_t crc = crc8(frame, 6);
    frame[6] = crc;
    uart_write_bytes(IO_UART_PORT, (char *)frame, 7);
}

/* Read exact n bytes with timeout */
static bool uart_read_exact(uint8_t *buf, int n, int timeout_ms)
{
    int got = 0;
    while (got < n && timeout_ms > 0) {
        int r = uart_read_bytes(IO_UART_PORT, buf + got, n - got,
                                pdMS_TO_TICKS(50));
        if (r > 0) got += r;
        else timeout_ms -= 50;
    }
    return got == n;
}

static void ota_recv_task(void *arg)
{
    esp_ota_handle_t       handle    = 0;
    const esp_partition_t *ota_part  = esp_ota_get_next_update_partition(NULL);
    mbedtls_sha256_context sha;
    mbedtls_sha256_init(&sha);
    uint32_t expected_size = 0;
    uint8_t  expected_digest[32] = {0};
    uint32_t received = 0;
    bool     ota_started = false;
    static uint8_t payload[CHUNK_MAX];

    for (;;) {
        /* Read frame header: TYPE(1) + SEQ(2) + LEN(2) = 5 bytes */
        uint8_t hdr[5];
        if (!uart_read_exact(hdr, 5, 5000)) continue;

        uint8_t  type = hdr[0];
        uint16_t seq  = (uint16_t)((hdr[1] << 8u) | hdr[2]);
        uint16_t plen = (uint16_t)((hdr[3] << 8u) | hdr[4]);

        if (plen > CHUNK_MAX + 36) {
            ESP_LOGW(TAG, "oversized frame plen=%u", plen);
            continue;
        }

        /* Read payload + CRC */
        if (plen > 0 && !uart_read_exact(payload, plen, 2000)) continue;
        uint8_t crc_rx;
        if (!uart_read_exact(&crc_rx, 1, 500)) continue;

        /* Verify CRC over hdr+payload */
        uint8_t crc_buf[5 + CHUNK_MAX + 36];
        memcpy(crc_buf, hdr, 5);
        if (plen > 0) memcpy(crc_buf + 5, payload, plen);
        uint8_t expected_crc = crc8(crc_buf, (uint16_t)(5 + plen));
        if (crc_rx != expected_crc) {
            ESP_LOGW(TAG, "CRC fail seq=%u", seq);
            send_nack(seq, 0x01u);  /* OTA_ERR_BAD_CRC */
            continue;
        }

        switch (type) {
        case OTA_BEGIN:
            if (plen < 36) { send_nack(seq, 0x03u); break; }
            expected_size = ((uint32_t)payload[0] << 24u) |
                            ((uint32_t)payload[1] << 16u) |
                            ((uint32_t)payload[2] <<  8u) |
                             (uint32_t)payload[3];
            memcpy(expected_digest, &payload[4], 32);

            if (!ota_part || esp_ota_begin(ota_part, OTA_WITH_SEQUENTIAL_WRITES,
                                            &handle) != ESP_OK) {
                send_nack(seq, 0x02u);
                break;
            }
            mbedtls_sha256_starts(&sha, 0);
            received   = 0;
            ota_started = true;
            g_io_ota_state    = IO_OTA_RECEIVING;
            g_io_ota_progress = 0;
            ESP_LOGI(TAG, "OTA begin: %lu bytes", (unsigned long)expected_size);
            send_ack(seq);
            break;

        case OTA_DATA:
            if (!ota_started) { send_nack(seq, 0x02u); break; }
            if (esp_ota_write(handle, payload, plen) != ESP_OK) {
                send_nack(seq, 0x02u);
                esp_ota_abort(handle);
                ota_started = false;
                g_io_ota_state = IO_OTA_FAILED;
                break;
            }
            mbedtls_sha256_update(&sha, payload, plen);
            received += plen;
            if (expected_size > 0)
                g_io_ota_progress = (uint8_t)((received * 100u) / expected_size);
            send_ack(seq);
            break;

        case OTA_END: {
            if (!ota_started) { send_nack(seq, 0x02u); break; }
            g_io_ota_state = IO_OTA_VERIFYING;

            uint8_t digest[32];
            mbedtls_sha256_finish(&sha, digest);
            if (memcmp(digest, expected_digest, 32) != 0) {
                ESP_LOGE(TAG, "SHA256 mismatch");
                esp_ota_abort(handle);
                send_nack(seq, 0x01u);
                g_io_ota_state = IO_OTA_FAILED;
                break;
            }

            if (esp_ota_end(handle) != ESP_OK ||
                esp_ota_set_boot_partition(ota_part) != ESP_OK) {
                send_nack(seq, 0x02u);
                g_io_ota_state = IO_OTA_FAILED;
                break;
            }

            g_io_ota_state    = IO_OTA_REBOOTING;
            g_io_ota_progress = 100;
            ESP_LOGI(TAG, "OTA done — rebooting");
            send_ack(seq);
            vTaskDelay(pdMS_TO_TICKS(500));
            esp_restart();
            break;
        }

        case OTA_ABORT:
            if (ota_started) { esp_ota_abort(handle); ota_started = false; }
            g_io_ota_state = IO_OTA_IDLE;
            ESP_LOGW(TAG, "OTA aborted");
            break;

        default:
            break;
        }
    }
}

void uart_ota_recv_init(void)
{
    /* UART0 already initialized for inter-board comms; just create the task */
    xTaskCreate(ota_recv_task, "io_ota_recv", 8192, NULL, 6, NULL);
    ESP_LOGI(TAG, "OTA receiver task started");
}