static const char *TAG = "ota";
#include <esp_log.h>

#include <mdns.h>
#include <esp_ota_ops.h>

#include <freertos/FreeRTOS.h>
#include <freertos/task.h>

#include "ota.hpp"


// The real version probably lives somewhere else.  Too lazy to find it
template <class T> static inline T MIN(T a, T b) { return a < b ? a : b; }

// TODO: Stick this somewhere useful?
static void reboot_soon_task(void*) {
    int delay_s = 1;
    ESP_LOGI(TAG, "Rebooting in %d second%s",
        delay_s, delay_s == 1 ? "" : "s");
    vTaskDelay(pdMS_TO_TICKS(delay_s * 1000));
    esp_restart();
}

static void reboot_soon() {
    xTaskCreate(
            reboot_soon_task, "reboot_task",
            configMINIMAL_STACK_SIZE * 2, // TODO: Choose the best size
            NULL, 1, NULL
        );
}

// HTTP Server Helpers.
// TODO: Share these?
static inline esp_err_t serv_err(httpd_req_t *req, const char *msg) {
    ESP_LOGE(TAG, "%s", msg);
    return httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, msg);
}

static inline esp_err_t serv_progress(httpd_req_t *req, int progress) {
    ESP_LOGI(TAG, "%d%% complete", progress);
    /*
    char msg[100];
    snprintf(msg, sizeof(msg), "%d%% complete\n", progress);
    httpd_resp_set_status(req, "102 Processing");
    httpd_resp_set_type(req, HTTPD_TYPE_TEXT);
    return httpd_resp_send(req, msg, HTTPD_RESP_USE_STRLEN);
    */
    return ESP_OK;
}


// OTA POST
static esp_err_t ota_post_handler(httpd_req_t *req)
{
    esp_err_t err;

    const esp_partition_t *part = esp_ota_get_next_update_partition(NULL);
    if (!part) return serv_err(req, "Could not find update partition");

    int remaining = req->content_len;
    ESP_LOGI(TAG, "Starting OTA: %d bytes", remaining);

    esp_ota_handle_t ota;
    err = esp_ota_begin(part, remaining, &ota);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_ota_begin: %d", err);
        return serv_err(req, "Could not start OTA");
    }

    constexpr int progress_step = 10;
    int next_progress = 0;
    constexpr size_t CHUNK_MAX = 4096; 
    char *chunk_buf = new char[CHUNK_MAX];
    while (remaining > 0) {
        int chunk_size = MIN((unsigned int)remaining, CHUNK_MAX);

        int nread = httpd_req_recv(req, chunk_buf, chunk_size);
        if (nread <= 0) { /* 0 return value indicates connection closed */
            ESP_LOGE(TAG, "httpd_req_recv(%d): %d", chunk_size, nread);
            esp_ota_abort(ota);
            delete[] chunk_buf;

            /* Check if timeout occurred */
            if (nread == HTTPD_SOCK_ERR_TIMEOUT) {
                /* In case of timeout one can choose to retry calling
                 * httpd_req_recv(), but to keep it simple, here we
                 * respond with an HTTP 408 (Request Timeout) error */
                httpd_resp_send_408(req);
            }

            /* In case of error, returning ESP_FAIL will
             * ensure that the underlying socket is closed */
            return ESP_FAIL;
        }

        err = esp_ota_write(ota, chunk_buf, nread);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "esp_ota_write(%d): %d", nread, err);
            esp_ota_abort(ota);
            delete[] chunk_buf;
            return serv_err(req, "Failed to write OTA chunk");
        }

        remaining -= nread;

        int progress = 100 * (req->content_len - remaining) / req->content_len;
        if (progress >= next_progress) {
            serv_progress(req, progress);
            next_progress += progress_step;
        }
    }
    delete[] chunk_buf;

    ESP_LOGI(TAG, "Finalizing OTA");

    err = esp_ota_end(ota);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_ota_end: %d", err);
        esp_ota_abort(ota);
        return serv_err(req, "Failed to finalize OTA");
    }

    ESP_LOGI(TAG, "Setting new boot partition");
    err = esp_ota_set_boot_partition(part);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_ota_set_boot_partition: %d", err);
        return serv_err(req, "Failed to set partition");
    }

    ESP_LOGI(TAG, "OTA Successful");
    httpd_resp_sendstr(req, "OTA Successful\n");

    reboot_soon();

    return ESP_FAIL;
}


int start_ota_serv(httpd_handle_t server) {
    ESP_LOGI(TAG, "Starting OTA server");

    // Register handlers (move these to components)
    httpd_uri_t uri_ota_post = {
        .uri      = "/ota",
        .method   = HTTP_POST,
        .handler  = ota_post_handler,
        .user_ctx = NULL
    };
    ESP_ERROR_CHECK(httpd_register_uri_handler(server, &uri_ota_post));

    // mDNS Entry
    mdns_service_add(NULL, "_http", "_tcp", 80, NULL, 0);
    mdns_service_instance_name_set("_http", "_tcp", "ESP32S2 OTA");

    return 0;
}


int stop_ota_serv() {
     // TODO
    return -1;
}