/* ota_self.c — Balance self-OTA (bd-18nb)
 *
 * Uses esp_https_ota / esp_ota_ops to download from Gitea release URL,
 * stream-verify SHA256 with mbedTLS, set new boot partition, and reboot.
 * CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE in sdkconfig allows auto-rollback
 * if the new image doesn't call esp_ota_mark_app_valid_cancel_rollback()
 * within OTA_ROLLBACK_WINDOW_S seconds.
 */

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

static const char *TAG = "ota_self";

volatile ota_self_state_t g_ota_self_state   = OTA_SELF_IDLE;
volatile uint8_t          g_ota_self_progress = 0;

#define OTA_CHUNK_SIZE  4096

/* ── SHA256 verify helper ── */
static bool sha256_matches(const uint8_t *digest, const char *expected_hex)
{
    if (!expected_hex || expected_hex[0] == '\0') {
        ESP_LOGW(TAG, "no SHA256 to verify — skipping");
        return true;
    }
    char got[65] = {0};
    for (int i = 0; i < 32; i++) {
        snprintf(&got[i*2], 3, "%02x", digest[i]);
    }
    bool ok = (strncasecmp(got, expected_hex, 64) == 0);
    if (!ok) {
        ESP_LOGE(TAG, "SHA256 mismatch: got=%s exp=%s", got, expected_hex);
    }
    return ok;
}

/* ── OTA download + flash task ── */
static void ota_self_task(void *arg)
{
    const char *url    = g_balance_update.download_url;
    const char *sha256 = g_balance_update.sha256;

    g_ota_self_state    = OTA_SELF_DOWNLOADING;
    g_ota_self_progress = 0;

    ESP_LOGI(TAG, "OTA start: %s", url);

    esp_ota_handle_t       handle     = 0;
    const esp_partition_t *ota_part   = esp_ota_get_next_update_partition(NULL);
    if (!ota_part) {
        ESP_LOGE(TAG, "no OTA partition");
        g_ota_self_state = OTA_SELF_FAILED;
        vTaskDelete(NULL);
        return;
    }

    esp_err_t err = esp_ota_begin(ota_part, OTA_WITH_SEQUENTIAL_WRITES, &handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "ota_begin: %s", esp_err_to_name(err));
        g_ota_self_state = OTA_SELF_FAILED;
        vTaskDelete(NULL);
        return;
    }

    /* Setup HTTP client */
    esp_http_client_config_t hcfg = {
        .url        = url,
        .timeout_ms = 30000,
        .buffer_size = OTA_CHUNK_SIZE,
        .skip_cert_common_name_check = true,
    };
    esp_http_client_handle_t client = esp_http_client_init(&hcfg);
    err = esp_http_client_open(client, 0);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "http_open: %s", esp_err_to_name(err));
        esp_ota_abort(handle);
        esp_http_client_cleanup(client);
        g_ota_self_state = OTA_SELF_FAILED;
        vTaskDelete(NULL);
        return;
    }

    int content_len = esp_http_client_fetch_headers(client);
    ESP_LOGI(TAG, "content-length: %d", content_len);

    mbedtls_sha256_context sha_ctx;
    mbedtls_sha256_init(&sha_ctx);
    mbedtls_sha256_starts(&sha_ctx, 0);  /* 0 = SHA-256 */

    static uint8_t buf[OTA_CHUNK_SIZE];
    int total = 0;
    int rd;
    while ((rd = esp_http_client_read(client, (char *)buf, sizeof(buf))) > 0) {
        mbedtls_sha256_update(&sha_ctx, buf, rd);
        err = esp_ota_write(handle, buf, rd);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "ota_write: %s", esp_err_to_name(err));
            esp_ota_abort(handle);
            goto cleanup;
        }
        total += rd;
        if (content_len > 0) {
            g_ota_self_progress = (uint8_t)((total * 100) / content_len);
        }
    }
    esp_http_client_close(client);

    /* Verify SHA256 */
    g_ota_self_state = OTA_SELF_VERIFYING;
    uint8_t digest[32];
    mbedtls_sha256_finish(&sha_ctx, digest);
    if (!sha256_matches(digest, sha256)) {
        ESP_LOGE(TAG, "SHA256 verification failed");
        esp_ota_abort(handle);
        g_ota_self_state = OTA_SELF_FAILED;
        goto cleanup;
    }

    /* Finalize + set boot partition */
    g_ota_self_state = OTA_SELF_APPLYING;
    err = esp_ota_end(handle);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "ota_end: %s", esp_err_to_name(err));
        g_ota_self_state = OTA_SELF_FAILED;
        goto cleanup;
    }

    err = esp_ota_set_boot_partition(ota_part);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "set_boot_partition: %s", esp_err_to_name(err));
        g_ota_self_state = OTA_SELF_FAILED;
        goto cleanup;
    }

    g_ota_self_state    = OTA_SELF_REBOOTING;
    g_ota_self_progress = 100;
    ESP_LOGI(TAG, "OTA success — rebooting");
    vTaskDelay(pdMS_TO_TICKS(500));
    esp_restart();

cleanup:
    mbedtls_sha256_free(&sha_ctx);
    esp_http_client_cleanup(client);
    handle = 0;
    vTaskDelete(NULL);
}

bool ota_self_trigger(void)
{
    if (!g_balance_update.available) {
        ESP_LOGW(TAG, "no update available");
        return false;
    }
    if (g_ota_self_state != OTA_SELF_IDLE) {
        ESP_LOGW(TAG, "OTA already in progress (state=%d)", g_ota_self_state);
        return false;
    }
    xTaskCreate(ota_self_task, "ota_self", 8192, NULL, 5, NULL);
    return true;
}

void ota_self_health_check(void)
{
    /* Mark running image as valid — prevents rollback */
    esp_err_t err = esp_ota_mark_app_valid_cancel_rollback();
    if (err == ESP_OK) {
        ESP_LOGI(TAG, "image marked valid");
    } else if (err == ESP_ERR_NOT_SUPPORTED) {
        /* Not an OTA image (e.g., flashed via JTAG) — ignore */
    } else {
        ESP_LOGW(TAG, "mark_valid: %s", esp_err_to_name(err));
    }
}