#include "aes67_ota.h"

#include <stdio.h>
#include <string.h>

#include "aes67_web.h"
#include "esp_app_desc.h"
#include "esp_app_format.h"
#include "esp_http_client.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_ota_ops.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "hal/efuse_hal.h"

#define CHUNK               4096
#define SELFTEST_TIMEOUT_S  60
// Image header, first segment header, then the app description.
#define DESC_OFFSET         (sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t))
#define HEAD_LEN            (DESC_OFFSET + sizeof(esp_app_desc_t))

static const char *TAG = "ota";

#define MAX_UPDATE_CBS 4
static aes67_ota_update_cb_t s_update_cbs[MAX_UPDATE_CBS];
static int s_update_n;

esp_err_t aes67_ota_on_update(aes67_ota_update_cb_t cb)
{
    if (s_update_n >= MAX_UPDATE_CBS) {
        return ESP_ERR_NO_MEM;
    }
    s_update_cbs[s_update_n++] = cb;
    return ESP_OK;
}

static void notify_update(bool starting)
{
    for (int i = 0; i < s_update_n; i++) {
        s_update_cbs[i](starting);
    }
}

static bool is_pending(void)
{
    esp_ota_img_states_t st;
    return esp_ota_get_state_partition(esp_ota_get_running_partition(), &st) == ESP_OK &&
           st == ESP_OTA_IMG_PENDING_VERIFY;
}

/* ----- GET /api/ota ----- */

static esp_err_t ota_get(httpd_req_t *req)
{
    const esp_app_desc_t *app = esp_app_get_description();
    const esp_partition_t *run = esp_ota_get_running_partition();
    const esp_partition_t *other = esp_ota_get_next_update_partition(NULL);
    char built[40];
    snprintf(built, sizeof(built), "%s %s", app->date, app->time);

    cJSON *o = cJSON_CreateObject();
    cJSON_AddStringToObject(o, "version", app->version);
    cJSON_AddStringToObject(o, "project", app->project_name);
    cJSON_AddStringToObject(o, "build_date", built);
    cJSON_AddStringToObject(o, "idf", app->idf_ver);
    cJSON_AddStringToObject(o, "running", run->label);
    // Only a bootable image counts as "previous" (not a half-written upload or a rolled-back
    // one). With two slots, rollback is possible exactly when the other slot is bootable.
    bool can_rollback = esp_ota_check_rollback_is_possible();
    esp_app_desc_t prev;
    if (can_rollback && other && esp_ota_get_partition_description(other, &prev) == ESP_OK) {
        cJSON_AddStringToObject(o, "previous", other->label);
        cJSON_AddStringToObject(o, "previous_version", prev.version);
    }
    cJSON_AddBoolToObject(o, "pending_verify", is_pending());
    cJSON_AddBoolToObject(o, "can_rollback", can_rollback);
    esp_err_t err = web_send_json(req, o);
    cJSON_Delete(o);
    return err;
}

/* ----- POST /api/ota ----- */

// Checks on the first bytes, before anything is written to flash.
static const char *check_head(const uint8_t *buf)
{
    const esp_image_header_t *h = (const esp_image_header_t *)buf;
    const esp_app_desc_t *d = (const esp_app_desc_t *)(buf + DESC_OFFSET);
    static char msg[96];

    if (h->magic != ESP_IMAGE_HEADER_MAGIC || d->magic_word != ESP_APP_DESC_MAGIC_WORD) {
        return "not an ESP-IDF app image";
    }
    if (h->chip_id != CONFIG_IDF_FIRMWARE_CHIP_ID) {
        return "image is for a different chip";
    }
    if (strncmp(d->project_name, esp_app_get_description()->project_name, sizeof(d->project_name)) != 0) {
        snprintf(msg, sizeof(msg), "wrong project '%.32s'", d->project_name);
        return msg;
    }
    unsigned rev = efuse_hal_chip_revision();
    if (rev < h->min_chip_rev_full || rev > h->max_chip_rev_full) {
        snprintf(msg, sizeof(msg), "image supports chip rev v%u.%u-v%u.%u, this chip is v%u.%u",
                 h->min_chip_rev_full / 100, h->min_chip_rev_full % 100,
                 h->max_chip_rev_full / 100, h->max_chip_rev_full % 100, rev / 100, rev % 100);
        return msg;
    }
    return NULL;
}

static esp_err_t reject(httpd_req_t *req, const char *msg)
{
    ESP_LOGW(TAG, "upload rejected: %s", msg);
    return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, msg);
}

static esp_err_t ota_post(httpd_req_t *req)
{
    const esp_partition_t *dst = esp_ota_get_next_update_partition(NULL);
    if (!dst) {
        return reject(req, "no OTA partition");
    }
    if (req->content_len < HEAD_LEN || req->content_len > dst->size) {
        return reject(req, "image size out of range");
    }
    uint8_t *buf = malloc(CHUNK);
    if (!buf) {
        return httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "out of memory");
    }

    ESP_LOGI(TAG, "upload %u bytes to %s", (unsigned)req->content_len, dst->label);
    esp_ota_handle_t ota = 0;
    size_t total = 0, fill = 0;
    const char *err_msg = NULL;
    bool notified = false;
    int64_t t0 = esp_timer_get_time();

    while (total < req->content_len) {
        int r = httpd_req_recv(req, (char *)buf + fill, CHUNK - fill);
        if (r == HTTPD_SOCK_ERR_TIMEOUT) {
            continue;
        }
        if (r <= 0) {
            err_msg = "connection lost";
            break;
        }
        fill += r;
        total += r;
        if (!ota) {
            // Collect the header before deciding anything.
            if (fill < HEAD_LEN && total < req->content_len) {
                continue;
            }
            if ((err_msg = check_head(buf)) != NULL) {
                break;
            }
            const esp_app_desc_t *d = (const esp_app_desc_t *)(buf + DESC_OFFSET);
            ESP_LOGI(TAG, "image %.32s %.32s", d->project_name, d->version);
            notify_update(true);    // project quiets streaming/decoding before the flash writes
            notified = true;
            if (esp_ota_begin(dst, OTA_WITH_SEQUENTIAL_WRITES, &ota) != ESP_OK) {
                err_msg = "ota begin failed";
                break;
            }
        }
        if (esp_ota_write(ota, buf, fill) != ESP_OK) {
            err_msg = "flash write failed";
            break;
        }
        fill = 0;
    }
    free(buf);

    if (err_msg) {
        if (ota) {
            esp_ota_abort(ota);
        }
        if (notified) {
            notify_update(false);
        }
        // Rejected before the whole body was read: close instead of draining it.
        httpd_resp_set_hdr(req, "Connection", "close");
        return reject(req, err_msg);
    }
    esp_err_t err = esp_ota_end(ota);  // verifies the image (checksum/hash, chip)
    if (err != ESP_OK) {
        notify_update(false);
        return reject(req, err == ESP_ERR_OTA_VALIDATE_FAILED ? "image verification failed" : "ota end failed");
    }
    if (esp_ota_set_boot_partition(dst) != ESP_OK) {
        notify_update(false);
        return httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "could not set boot partition");
    }
    ESP_LOGW(TAG, "installed to %s in %.1f s, rebooting", dst->label,
             (esp_timer_get_time() - t0) / 1e6);
    httpd_resp_sendstr(req, "ok");
    web_reboot_later(500);
    return ESP_OK;
}

/* ----- confirm / rollback ----- */

static esp_err_t confirm_post(httpd_req_t *req)
{
    if (!is_pending()) {
        return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "running firmware is already confirmed");
    }
    esp_ota_mark_app_valid_cancel_rollback();
    ESP_LOGI(TAG, "firmware confirmed via API");
    return httpd_resp_sendstr(req, "");
}

static void rollback_cb(void *arg)
{
    esp_ota_mark_app_invalid_rollback_and_reboot();
    ESP_LOGE(TAG, "rollback failed");
}

static esp_err_t rollback_post(httpd_req_t *req)
{
    if (!esp_ota_check_rollback_is_possible()) {
        return httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "no previous firmware to roll back to");
    }
    ESP_LOGW(TAG, "rollback requested via API");
    httpd_resp_sendstr(req, "");
    const esp_timer_create_args_t args = { .callback = rollback_cb, .name = "rollback" };
    esp_timer_handle_t t;
    if (esp_timer_create(&args, &t) == ESP_OK) {
        esp_timer_start_once(t, 500 * 1000);
    }
    return ESP_OK;
}

/* ----- Self-test after an update ----- */

static bool web_answers(const esp_netif_ip_info_t *ip)
{
    char url[48];
    snprintf(url, sizeof(url), "http://" IPSTR "/api/status", IP2STR(&ip->ip));
    esp_http_client_config_t cfg = { .url = url, .timeout_ms = 3000 };
    esp_http_client_handle_t c = esp_http_client_init(&cfg);
    bool ok = c && esp_http_client_perform(c) == ESP_OK && esp_http_client_get_status_code(c) == 200;
    esp_http_client_cleanup(c);
    return ok;
}

static void selftest_task(void *arg)
{
    const char *fail = NULL;
    int64_t deadline = esp_timer_get_time() + SELFTEST_TIMEOUT_S * 1000000LL;
    esp_netif_t *netif = esp_netif_get_handle_from_ifkey("ETH_DEF");
    bool passed = false;

    while (!passed && esp_timer_get_time() < deadline) {
        vTaskDelay(pdMS_TO_TICKS(1000));
        esp_netif_ip_info_t ip;
        if (!netif || esp_netif_get_ip_info(netif, &ip) != ESP_OK || !ip.ip.addr) {
            fail = "no IP address";
            continue;
        }
        if (!web_answers(&ip)) {
            fail = "web server not answering";
            continue;
        }
        passed = true;
    }
#if CONFIG_AES67_OTA_SELFTEST_FORCE_FAIL
    passed = false;
    fail = "forced failure (CONFIG_AES67_OTA_SELFTEST_FORCE_FAIL)";
#endif
    if (passed) {
        esp_ota_mark_app_valid_cancel_rollback();
        ESP_LOGI(TAG, "self-test passed, firmware confirmed");
    } else if (is_pending()) {  // not confirmed manually in the meantime
        ESP_LOGE(TAG, "self-test failed (%s), rolling back", fail);
        esp_ota_mark_app_invalid_rollback_and_reboot();
    }
    vTaskDelete(NULL);
}

esp_err_t aes67_ota_init(void)
{
    static const httpd_uri_t uris[] = {
        { .uri = "/api/ota",          .method = HTTP_GET,  .handler = ota_get },
        { .uri = "/api/ota",          .method = HTTP_POST, .handler = ota_post },
        { .uri = "/api/ota/confirm",  .method = HTTP_POST, .handler = confirm_post },
        { .uri = "/api/ota/rollback", .method = HTTP_POST, .handler = rollback_post },
    };
    for (int i = 0; i < sizeof(uris) / sizeof(uris[0]); i++) {
        esp_err_t err = web_register_uri(&uris[i]);
        if (err != ESP_OK) {
            return err;
        }
    }

    const esp_app_desc_t *app = esp_app_get_description();
    ESP_LOGI(TAG, "running %s from %s", app->version, esp_ota_get_running_partition()->label);
    if (is_pending()) {
        ESP_LOGW(TAG, "new firmware on trial: self-test (IP + web) within %d s", SELFTEST_TIMEOUT_S);
        xTaskCreate(selftest_task, "ota_selftest", 4096, NULL, 2, NULL);
    }
    return ESP_OK;
}