/*
 * This program source code file is part of the KTag project.
 *
 * 🛡️ <https://ktag.clubk.club> 🃞
 *
 * Copyright © 2024-2025 Joseph P. Kearney and the KTag developers.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more
 * details.
 *
 * There should be a copy of the GNU Affero General Public License in the LICENSE
 * file in the root of this repository.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <SystemK.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <errno.h>
#include <NVM.h>
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_spiffs.h"
#include "esp_partition.h"

static const char *TAG = "Reprogramming";
static const char *FIRMWARE_DIR = "/usb/esp/firmware";
static const char *APP_IMAGE_PREFIX = "APP";
static const char *SPIFFS_IMAGE_PREFIX = "SPIFFS";

#define REPROGRAMMING_TASK_PRIORITY 1
static TaskHandle_t xReprogrammingTask = NULL;
static StaticTask_t xReprogrammingTaskBuffer;
static StackType_t xReprogrammingTaskStack[4096];

#define BUFFER_SIZE 1024

typedef struct
{
    unsigned long total_size;
    unsigned long remaining_size;
    void *data;
} data_manager_t;

typedef enum {
    IMAGE_TYPE_APP,
    IMAGE_TYPE_SPIFFS
} image_type_t;

static char *find_image_filename(const char *prefix)
{
    DIR *dir = opendir(FIRMWARE_DIR);
    if (dir == NULL)
    {
        KLOG_ERROR(TAG, "Failed to open directory '%s': %s!", FIRMWARE_DIR, strerror(errno));
        return NULL;
    }

    struct dirent *entry;
    char *filename = NULL;

    while ((entry = readdir(dir)) != NULL)
    {
        if (entry->d_type == DT_REG && strncmp(entry->d_name, prefix, strlen(prefix)) == 0)
        {
            filename = strdup(entry->d_name);
            if (filename == NULL)
            {
                KLOG_ERROR(TAG, "Memory allocation failed!");
            }
            else
            {
                KLOG_DEBUG(TAG, "Found image: %s", filename);
            }
            break;
        }
    }

    closedir(dir);

    if (filename == NULL)
    {
        KLOG_WARN(TAG, "Couldn't find an image with prefix '%s' in directory '%s'!", prefix, FIRMWARE_DIR);
        return NULL;
    }

    size_t needed = snprintf(NULL, 0, "%s/%s", FIRMWARE_DIR, filename) + 1;
    char *full_path = malloc(needed);
    if (full_path != NULL)
    {
        snprintf(full_path, needed, "%s/%s", FIRMWARE_DIR, filename);
    }

    free(filename);
    return full_path;
}

static SystemKResult_T program_image(const char *filename, image_type_t image_type)
{
    esp_err_t err;
    FILE *file = fopen(filename, "rb");
    if (file == NULL)
    {
        KLOG_ERROR(TAG, "Failed to open file %s", filename);
        return SYSTEMK_RESULT_FILE_NOT_FOUND;
    }

    fseek(file, 0, SEEK_END);
    long file_size = ftell(file);
    fseek(file, 0, SEEK_SET);

    KLOG_INFO(TAG, "File size: %ld bytes", file_size);

    data_manager_t data_manager = {
        .total_size = file_size,
        .remaining_size = file_size,
        .data = malloc(BUFFER_SIZE)
    };

    if (data_manager.data == NULL)
    {
        KLOG_ERROR(TAG, "Failed to allocate memory");
        fclose(file);
        return SYSTEMK_RESULT_NOT_ENOUGH_MEMORY;
    }

    if (image_type == IMAGE_TYPE_APP)
    {
        esp_ota_handle_t update_handle = 0;
        const esp_partition_t *update_partition = esp_ota_get_next_update_partition(NULL);
        KLOG_INFO(TAG, "Writing to partition subtype %d at offset 0x%" PRIx32, update_partition->subtype, update_partition->address);
        err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
        if (err != ESP_OK)
        {
            KLOG_ERROR(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
            free(data_manager.data);
            fclose(file);
            return SYSTEMK_RESULT_UNSPECIFIED_FAILURE;
        }

        while (data_manager.remaining_size > 0)
        {
            size_t size = (data_manager.remaining_size <= BUFFER_SIZE) ? data_manager.remaining_size : BUFFER_SIZE;
            if (fread(data_manager.data, 1, size, file) != size)
            {
                KLOG_ERROR(TAG, "fread failed");
                err = ESP_FAIL;
                break;
            }

            err = esp_ota_write(update_handle, data_manager.data, size);
            if (err != ESP_OK)
            {
                KLOG_ERROR(TAG, "esp_ota_write failed (%s)", esp_err_to_name(err));
                break;
            }

            data_manager.remaining_size -= size;
        }

        if (err == ESP_OK)
        {
            err = esp_ota_end(update_handle);
            if (err != ESP_OK)
            {
                KLOG_ERROR(TAG, "esp_ota_end failed (%s)", esp_err_to_name(err));
            }
            else
            {
                err = esp_ota_set_boot_partition(update_partition);
                if (err != ESP_OK)
                {
                    KLOG_ERROR(TAG, "esp_ota_set_boot_partition failed (%s)", esp_err_to_name(err));
                }
            }
        }
    }
    else if (image_type == IMAGE_TYPE_SPIFFS)
    {
        const esp_partition_t *spiffs_partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_SPIFFS, NULL);
        if (spiffs_partition == NULL)
        {
            KLOG_ERROR(TAG, "SPIFFS partition not found");
            free(data_manager.data);
            fclose(file);
            return SYSTEMK_RESULT_UNSPECIFIED_FAILURE;
        }

        err = esp_partition_erase_range(spiffs_partition, 0, spiffs_partition->size);
        if (err != ESP_OK)
        {
            KLOG_ERROR(TAG, "Failed to erase SPIFFS partition (%s)", esp_err_to_name(err));
            free(data_manager.data);
            fclose(file);
            return SYSTEMK_RESULT_UNSPECIFIED_FAILURE;
        }

        while (data_manager.remaining_size > 0)
        {
            size_t size = (data_manager.remaining_size <= BUFFER_SIZE) ? data_manager.remaining_size : BUFFER_SIZE;
            if (fread(data_manager.data, 1, size, file) != size)
            {
                KLOG_ERROR(TAG, "fread failed");
                err = ESP_FAIL;
                break;
            }

            err = esp_partition_write(spiffs_partition, spiffs_partition->size - data_manager.remaining_size, data_manager.data, size);
            if (err != ESP_OK)
            {
                KLOG_ERROR(TAG, "esp_partition_write failed (%s)", esp_err_to_name(err));
                break;
            }

            data_manager.remaining_size -= size;
        }
    }

    free(data_manager.data);
    fclose(file);

    return (err == ESP_OK) ? SYSTEMK_RESULT_SUCCESS : SYSTEMK_RESULT_UNSPECIFIED_FAILURE;
}

SystemKResult_T Maybe_Reprogram_From_USB(void)
{
    SystemKResult_T result = SYSTEMK_RESULT_SUCCESS;
    char *app_image_filename = find_image_filename(APP_IMAGE_PREFIX);
    char *spiffs_image_filename = find_image_filename(SPIFFS_IMAGE_PREFIX);

    if (app_image_filename)
    {
        KLOG_INFO(TAG, "Application image found: %s", app_image_filename);
        result = program_image(app_image_filename, IMAGE_TYPE_APP);
        free(app_image_filename);
        if (result != SYSTEMK_RESULT_SUCCESS)
        {
            return result;
        }
    }

    if (spiffs_image_filename)
    {
        KLOG_INFO(TAG, "SPIFFS image found: %s", spiffs_image_filename);
        result = program_image(spiffs_image_filename, IMAGE_TYPE_SPIFFS);
        free(spiffs_image_filename);
    }

    if ((app_image_filename == NULL) && (spiffs_image_filename == NULL))
    {
        KLOG_ERROR(TAG, "Reprogramming failed--there's nothing to reprogram!");
        result = SYSTEMK_RESULT_FILE_NOT_FOUND;
    }

    return result;
}

static void Reprogramming_Task(void *pvParameters)
{
    SystemKResult_T result = Maybe_Reprogram_From_USB();
    static All_On_Data_T data;

    if (result == SYSTEMK_RESULT_SUCCESS)
    {
        KLOG_INFO(TAG, "Reprogramming succeeded--rebooting...");
        data.color = ApplyMask(COLOR_GREEN, 0x70);
        data.style = DISPLAY_STYLE_SOLID;
    }
    else
    {
        KLOG_ERROR(TAG, "Reprogramming failed--rebooting...");
        data.color = ApplyMask(COLOR_RED, 0x70);
        data.style = DISPLAY_STYLE_SOLID;
    }

    NeoPixelsAction_T neopixels_action = {.ID = NEOPIXELS_ALL_ON, .Prominence = NEOPIXELS_FOREGROUND, .Data = (void *)&data};
    xQueueSend(xQueueNeoPixels, &neopixels_action, 0);
    vTaskDelay(5000 / portTICK_PERIOD_MS);
    esp_restart();

    vTaskDelete(NULL);
}

void Request_Reprogramming_From_USB(void)
{
    if (xReprogrammingTask == NULL)
    {
        KLOG_INFO(TAG, "Attempting USB reprogramming...");

        xReprogrammingTask = xTaskCreateStaticPinnedToCore(
            Reprogramming_Task,
            "Reprogramming",
            sizeof(xReprogrammingTaskStack) / sizeof(StackType_t),
            NULL,
            REPROGRAMMING_TASK_PRIORITY,
            xReprogrammingTaskStack,
            &xReprogrammingTaskBuffer,
            APP_CPU_NUM);
    }
    else
    {
        KLOG_WARN(TAG, "Reprogramming already in progress.");
    }
}