/* uart_ota.c — UART OTA sender: Balance→IO board (bd-21hv)
 *
 * Downloads io-firmware.bin from Gitea, then sends to IO board via UART1.
 * IO board must update itself BEFORE Balance self-update (per spec).
 */

#include "uart_ota.h"
#include "gitea_ota.h"
#include "esp_log.h"
#include "esp_http_client.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "mbedtls/sha256.h"
#include <string.h>
#include <stdio.h>

static const char *TAG = "uart_ota";

volatile uart_ota_send_state_t g_uart_ota_state    = UART_OTA_S_IDLE;
volatile uint8_t               g_uart_ota_progress  = 0;

/* ── CRC8-SMBUS ── */
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;
}

/* ── Build and send one UART OTA frame ── */
static void send_frame(uint8_t type, uint16_t seq,
                        const uint8_t *payload, uint16_t plen)
{
    /* [TYPE:1][SEQ:2 BE][LEN:2 BE][PAYLOAD][CRC8:1] */
    uint8_t hdr[5];
    hdr[0] = type;
    hdr[1] = (uint8_t)(seq >> 8u);
    hdr[2] = (uint8_t)(seq);
    hdr[3] = (uint8_t)(plen >> 8u);
    hdr[4] = (uint8_t)(plen);

    /* CRC over hdr + payload */
    uint8_t crc_buf[5 + OTA_UART_CHUNK_SIZE];
    memcpy(crc_buf, hdr, 5);
    if (plen > 0 && payload) memcpy(crc_buf + 5, payload, plen);
    uint8_t crc = crc8(crc_buf, (uint16_t)(5 + plen));

    uart_write_bytes(UART_OTA_PORT, (char *)hdr, 5);
    if (plen > 0 && payload)
        uart_write_bytes(UART_OTA_PORT, (char *)payload, plen);
    uart_write_bytes(UART_OTA_PORT, (char *)&crc, 1);
}

/* ── Wait for ACK/NACK from IO board ── */
static bool wait_ack(uint16_t expected_seq)
{
    /* Response frame: [TYPE:1][SEQ:2][LEN:2][PAYLOAD][CRC:1] */
    uint8_t buf[16];
    int timeout = OTA_UART_ACK_TIMEOUT_MS;
    int got = 0;

    while (timeout > 0 && got < 6) {
        int r = uart_read_bytes(UART_OTA_PORT, buf + got, 1, pdMS_TO_TICKS(50));
        if (r > 0) got++;
        else timeout -= 50;
    }

    if (got < 3) return false;

    uint8_t type = buf[0];
    uint16_t seq = (uint16_t)((buf[1] << 8u) | buf[2]);

    if (type == UART_OTA_ACK && seq == expected_seq) return true;
    if (type == UART_OTA_NACK) {
        uint8_t err = (got >= 6) ? buf[5] : 0;
        ESP_LOGW(TAG, "NACK seq=%u err=%u", seq, err);
    }
    return false;
}

/* ── Download firmware to RAM buffer (max 1.75 MB) ── */
static uint8_t *download_io_firmware(uint32_t *out_size)
{
    const char *url = g_io_update.download_url;
    ESP_LOGI(TAG, "downloading IO fw: %s", url);

    esp_http_client_config_t cfg = {
        .url = url, .timeout_ms = 30000,
        .skip_cert_common_name_check = true,
    };
    esp_http_client_handle_t client = esp_http_client_init(&cfg);
    if (esp_http_client_open(client, 0) != ESP_OK) {
        esp_http_client_cleanup(client);
        return NULL;
    }

    int content_len = esp_http_client_fetch_headers(client);
    if (content_len <= 0 || content_len > (int)(0x1B0000)) {
        ESP_LOGE(TAG, "bad content-length: %d", content_len);
        esp_http_client_cleanup(client);
        return NULL;
    }

    uint8_t *buf = malloc(content_len);
    if (!buf) {
        ESP_LOGE(TAG, "malloc %d failed", content_len);
        esp_http_client_cleanup(client);
        return NULL;
    }

    int total = 0, rd;
    while ((rd = esp_http_client_read(client, (char *)buf + total,
                                       content_len - total)) > 0) {
        total += rd;
        g_uart_ota_progress = (uint8_t)((total * 50) / content_len); /* 0-50% for download */
    }
    esp_http_client_cleanup(client);

    if (total != content_len) {
        free(buf);
        return NULL;
    }
    *out_size = (uint32_t)total;
    return buf;
}

/* ── UART OTA send task ── */
static void uart_ota_task(void *arg)
{
    g_uart_ota_state    = UART_OTA_S_DOWNLOADING;
    g_uart_ota_progress = 0;

    uint32_t fw_size = 0;
    uint8_t *fw = download_io_firmware(&fw_size);
    if (!fw) {
        ESP_LOGE(TAG, "download failed");
        g_uart_ota_state = UART_OTA_S_FAILED;
        vTaskDelete(NULL);
        return;
    }

    /* Compute SHA256 of downloaded firmware */
    uint8_t digest[32];
    mbedtls_sha256_context sha;
    mbedtls_sha256_init(&sha);
    mbedtls_sha256_starts(&sha, 0);
    mbedtls_sha256_update(&sha, fw, fw_size);
    mbedtls_sha256_finish(&sha, digest);
    mbedtls_sha256_free(&sha);

    g_uart_ota_state = UART_OTA_S_SENDING;

    /* Send OTA_BEGIN: uint32 size + uint8[32] sha256 */
    uint8_t begin_payload[36];
    begin_payload[0] = (uint8_t)(fw_size >> 24u);
    begin_payload[1] = (uint8_t)(fw_size >> 16u);
    begin_payload[2] = (uint8_t)(fw_size >>  8u);
    begin_payload[3] = (uint8_t)(fw_size);
    memcpy(&begin_payload[4], digest, 32);

    for (int retry = 0; retry < OTA_UART_MAX_RETRIES; retry++) {
        send_frame(UART_OTA_BEGIN, 0, begin_payload, 36);
        if (wait_ack(0)) goto send_data;
        ESP_LOGW(TAG, "BEGIN retry %d", retry);
    }
    ESP_LOGE(TAG, "BEGIN failed");
    free(fw);
    g_uart_ota_state = UART_OTA_S_FAILED;
    vTaskDelete(NULL);
    return;

send_data: {
    uint32_t offset = 0;
    uint16_t seq    = 1;
    while (offset < fw_size) {
        uint16_t chunk = (uint16_t)((fw_size - offset) < OTA_UART_CHUNK_SIZE
                          ? (fw_size - offset) : OTA_UART_CHUNK_SIZE);
        bool acked = false;
        for (int retry = 0; retry < OTA_UART_MAX_RETRIES; retry++) {
            send_frame(UART_OTA_DATA, seq, fw + offset, chunk);
            if (wait_ack(seq)) { acked = true; break; }
            ESP_LOGW(TAG, "DATA seq=%u retry=%d", seq, retry);
        }
        if (!acked) {
            ESP_LOGE(TAG, "DATA seq=%u failed", seq);
            send_frame(UART_OTA_ABORT, seq, NULL, 0);
            free(fw);
            g_uart_ota_state = UART_OTA_S_FAILED;
            vTaskDelete(NULL);
            return;
        }
        offset += chunk;
        seq++;
        /* 50-100% for sending phase */
        g_uart_ota_progress = (uint8_t)(50u + (offset * 50u) / fw_size);
    }

    /* Send OTA_END */
    for (int retry = 0; retry < OTA_UART_MAX_RETRIES; retry++) {
        send_frame(UART_OTA_END, seq, NULL, 0);
        if (wait_ack(seq)) break;
    }
    }

    free(fw);
    g_uart_ota_progress = 100;
    g_uart_ota_state    = UART_OTA_S_DONE;
    ESP_LOGI(TAG, "IO OTA complete — %lu bytes sent", (unsigned long)fw_size);
    vTaskDelete(NULL);
}

bool uart_ota_trigger(void)
{
    if (!g_io_update.available) {
        ESP_LOGW(TAG, "no IO update available");
        return false;
    }
    if (g_uart_ota_state != UART_OTA_S_IDLE) {
        ESP_LOGW(TAG, "UART OTA busy (state=%d)", g_uart_ota_state);
        return false;
    }
    /* Init UART1 for OTA */
    uart_config_t ucfg = {
        .baud_rate  = UART_OTA_BAUD,
        .data_bits  = UART_DATA_8_BITS,
        .parity     = UART_PARITY_DISABLE,
        .stop_bits  = UART_STOP_BITS_1,
        .flow_ctrl  = UART_HW_FLOWCTRL_DISABLE,
    };
    uart_param_config(UART_OTA_PORT, &ucfg);
    uart_set_pin(UART_OTA_PORT, UART_OTA_TX_GPIO, UART_OTA_RX_GPIO,
                  UART_PIN_NO_CHANGE, UART_PIN_NO_CHANGE);
    uart_driver_install(UART_OTA_PORT, 2048, 0, 0, NULL, 0);

    xTaskCreate(uart_ota_task, "uart_ota", 16384, NULL, 4, NULL);
    return true;
}