#include "esp_http_client.h"
#include "esp_ota_ops.h"
#include "esp_system.h"
#include "esp_log.h"
#include "ota.h"
#include "sender_task.h"
#include "sampler_task.h"
#include "ws.h"
#include "heartbeat_task.h"
#include "fw_command.h"
#include "psa/crypto.h"

static TaskHandle_t ota_task_handle = nullptr;
static char url_copy[256];
static char sha256[65];

#define TAG "OTA"

void perform_ota(void* par)
{
    sampler_task_stop();
    sender_task_stop();
    heartbeat_task_stop();
    ws_disconnect();

    ESP_LOGI(TAG, "OTA task started");
    esp_ota_handle_t ota_handle;
    uint8_t buf[1024];
    int read_bytes;
    const esp_partition_t* partition = nullptr;
    int content_length = 0;
    psa_hash_operation_t op = PSA_HASH_OPERATION_INIT;

    const char* url = static_cast<const char*>(par);

    ESP_LOGI(TAG, "URL: %s", url);

    esp_http_client_config_t config = {};
    config.url = url;
    config.timeout_ms = 5000;
    config.buffer_size = 1024;

    esp_http_client_handle_t client = esp_http_client_init(&config);

    ESP_LOGI(TAG, "HTTP init complete");


    if (esp_http_client_open(client, 0) != ESP_OK)
    {
        goto cleanup;
    }

    ESP_LOGI(TAG, "HTTP connection established");

    content_length = esp_http_client_fetch_headers(client);
    ESP_LOGI(TAG, "content_length=%d", content_length);

    partition = esp_ota_get_next_update_partition(nullptr);

    ESP_LOGI(TAG, "OTA update partition selected %s (offset 0x%08x)", partition->label, partition->address);

    if (esp_ota_begin(partition, OTA_SIZE_UNKNOWN, &ota_handle) != ESP_OK)
    {
        goto cleanup;
    }

    psa_crypto_init();
    psa_hash_setup(&op, PSA_ALG_SHA_256);

    ESP_LOGI(TAG, "OTA update started");

    while (1)
    {
        int data_read = esp_http_client_read(client, (char*)buf, sizeof(buf));

        if (data_read < 0)
        {
            ESP_LOGE(TAG, "read error");
            break;
        }
        else if (data_read == 0)
        {
            ESP_LOGI(TAG, "download finished");
            break;
        }

        ESP_LOGI(TAG, "write chunk: %d bytes", data_read);

        if (esp_ota_write(ota_handle, buf, data_read) != ESP_OK)
        {
            ESP_LOGE(TAG, "flash write failed");
            goto cleanup;
        }
        if (psa_hash_update(&op, buf, data_read) != PSA_SUCCESS) {
            ESP_LOGE(TAG, "hash update failed");
        }
    }

    ESP_LOGI(TAG, "finalizing...");

    if (esp_ota_end(ota_handle) == ESP_OK)
    {
        uint8_t hash[32];
        size_t hash_len;

        psa_hash_finish(&op, hash, sizeof(hash), &hash_len);

        // переводим в hex
        char hash_str[65];
        for (int i = 0; i < 32; i++)
        {
            sprintf(&hash_str[i * 2], "%02x", hash[i]);
        }
        hash_str[64] = '\0';

        ESP_LOGI(TAG, "sha256 = %s. Comparing...", hash_str);

        // сравнение
        if (strncmp(hash_str, sha256, 64) != 0)
        {
            ESP_LOGE(TAG, "SHA256 mismatch!");
            goto cleanup;
        }

        ESP_LOGI(TAG, "set boot partition");
        esp_ota_set_boot_partition(partition);
        ESP_LOGI(TAG, "restart pending");
        esp_restart();
    }
    else
    {
        ESP_LOGE(TAG, "ota_end failed");
    }

cleanup:
    esp_http_client_cleanup(client);
    vTaskDelete(nullptr);
    ota_task_handle = nullptr;
}

void ota_task_start(const fw_cmd_t* cmd)
{
    strncpy(url_copy, cmd->url, sizeof(url_copy));
    strncpy(sha256, cmd->sha256, sizeof(sha256));
    if (ota_task_handle == nullptr)
    {
        xTaskCreate(
            perform_ota,
            "ota",
            4096,
            url_copy,
            5,
            &ota_task_handle
        );
    }
}

void ota_task_stop()
{
    if (ota_task_handle != nullptr)
    {
        vTaskDelete(ota_task_handle);
        ota_task_handle = nullptr;
    }
}