/*
 * 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/>.
 */

// From https://github.com/espressif/esp-idf/blob/master/examples/peripherals/usb/host/msc/main/msc_example_main.c

#include <SystemK.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <sys/stat.h>
#include <dirent.h>
#include <inttypes.h>
#include <errno.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/queue.h"
#include "usb/usb_host.h"
#include "usb/msc_host.h"
#include "usb/msc_host_vfs.h"
#include "esp_timer.h"
#include "esp_err.h"
#include "Key_Value.h"
#include "USB.h"

#define STACK_SIZE 4096
static StaticTask_t xTaskBuffer;
static StackType_t xStack[STACK_SIZE];

#define USB_HOST_TASK_PRIORITY 2
#define MSC_HOST_TASK_PRIORITY 3

#define MNT_PATH "/usb"  // Path in the Virtual File System, where the USB flash drive is going to be mounted
#define BUFFER_SIZE 4096 // The read/write performance can be improved with larger buffer for the cost of RAM; 4kB is enough for most use cases.

static const char *TAG = "USB";

static const char *OTA_FILE = "/usb/esp/OTA_URL.txt";
const char *CONFIG_FILE = "/usb/esp/config.txt";

typedef enum
{
    USB_STATE_UNINITIALIZED,
    USB_STATE_IDLE,
    USB_STATE_DEVICE_CONNECTED,
    USB_STATE_VFS_REGISTERED,
    USB_STATE_PROCESSING_DISCONNECTION
} USB_State_T;

static QueueHandle_t usb_queue;
typedef struct
{
    enum
    {
        USB_HOST_INITIALIZED,
        USB_DEVICE_CONNECTED,    // USB device connect event
        USB_DEVICE_DISCONNECTED, // USB device disconnect event
    } id;
    union
    {
        uint8_t new_dev_address; // Address of new USB device for APP_DEVICE_CONNECTED event
    } data;
} usb_message_t;

/**
 * @brief MSC driver callback
 *
 * Signal device connection/disconnection to the main task
 *
 * @param[in] event MSC event
 * @param[in] arg   MSC event data
 */
static void MSC_Event_Callback(const msc_host_event_t *event, void *arg)
{
    if (event->event == MSC_DEVICE_CONNECTED)
    {
        usb_message_t message = {
            .id = USB_DEVICE_CONNECTED,
            .data.new_dev_address = event->device.address,
        };
        xQueueSend(usb_queue, &message, portMAX_DELAY);
    }
    else if (event->event == MSC_DEVICE_DISCONNECTED)
    {
        usb_message_t message = {
            .id = USB_DEVICE_DISCONNECTED,
        };
        xQueueSend(usb_queue, &message, portMAX_DELAY);
    }
    else
    {
        KLOG_INFO(TAG, "Event: %d", event->event);
    }
}

__attribute__((unused)) static void print_device_info(msc_host_device_info_t *info)
{
    const size_t megabyte = 1024 * 1024;
    uint64_t capacity = ((uint64_t)info->sector_size * info->sector_count) / megabyte;

    printf("Device info:\n");
    printf("\t Capacity: %llu MB\n", capacity);
    printf("\t Sector size: %" PRIu32 "\n", info->sector_size);
    printf("\t Sector count: %" PRIu32 "\n", info->sector_count);
    printf("\t PID: 0x%04X \n", info->idProduct);
    printf("\t VID: 0x%04X \n", info->idVendor);
    wprintf(L"\t iProduct: %S \n", info->iProduct);
    wprintf(L"\t iManufacturer: %S \n", info->iManufacturer);
    wprintf(L"\t iSerialNumber: %S \n", info->iSerialNumber);
}

__attribute__((unused)) static void file_operations(void)
{
    const char *directory = "/usb/esp";
    const char *file_path = "/usb/esp/test.txt";

    // Create /usb/esp directory
    struct stat s = {0};
    bool directory_exists = stat(directory, &s) == 0;
    if (!directory_exists)
    {
        if (mkdir(directory, 0775) != 0)
        {
            KLOG_ERROR(TAG, "mkdir failed with errno: %s", strerror(errno));
        }
    }

    // Create /usb/esp/test.txt file, if it doesn't exist
    if (stat(file_path, &s) != 0)
    {
        KLOG_INFO(TAG, "Creating file");
        FILE *f = fopen(file_path, "w");
        if (f == NULL)
        {
            KLOG_ERROR(TAG, "Failed to open file for writing");
            return;
        }
        fprintf(f, "Hello World!\n");
        fclose(f);
    }

    // Read back the file
    FILE *f;
    KLOG_INFO(TAG, "Reading file");
    f = fopen(file_path, "r");
    if (f == NULL)
    {
        KLOG_ERROR(TAG, "Failed to open file for reading");
        return;
    }
    char line[64];
    fgets(line, sizeof(line), f);
    fclose(f);
    // strip newline
    char *pos = strchr(line, '\n');
    if (pos)
    {
        *pos = '\0';
    }
    KLOG_INFO(TAG, "Read from file '%s': '%s'", file_path, line);
}

__attribute__((unused)) void speed_test(void)
{
#define TEST_FILE "/usb/esp/dummy"
#define ITERATIONS 256 // 256 * 4kb = 1MB
    int64_t test_start, test_end;

    FILE *f = fopen(TEST_FILE, "wb+");
    if (f == NULL)
    {
        KLOG_ERROR(TAG, "Failed to open file for writing");
        return;
    }
    // Set larger buffer for this file. It results in larger and more effective USB transfers
    setvbuf(f, NULL, _IOFBF, BUFFER_SIZE);

    // Allocate application buffer used for read/write
    uint8_t *data = malloc(BUFFER_SIZE);
    assert(data);

    KLOG_INFO(TAG, "Writing to file %s", TEST_FILE);
    test_start = esp_timer_get_time();
    for (int i = 0; i < ITERATIONS; i++)
    {
        if (fwrite(data, BUFFER_SIZE, 1, f) == 0)
        {
            return;
        }
    }
    test_end = esp_timer_get_time();
    KLOG_INFO(TAG, "Write speed %1.2f MiB/s", (BUFFER_SIZE * ITERATIONS) / (float)(test_end - test_start));
    rewind(f);

    KLOG_INFO(TAG, "Reading from file %s", TEST_FILE);
    test_start = esp_timer_get_time();
    for (int i = 0; i < ITERATIONS; i++)
    {
        if (0 == fread(data, BUFFER_SIZE, 1, f))
        {
            return;
        }
    }
    test_end = esp_timer_get_time();
    KLOG_INFO(TAG, "Read speed %1.2f MiB/s", (BUFFER_SIZE * ITERATIONS) / (float)(test_end - test_start));

    fclose(f);
    free(data);
}

static void usb_host_task(void *args)
{
    usb_host_config_t host_config = {
        .skip_phy_setup = false,
        .intr_flags = ESP_INTR_FLAG_LEVEL1,
    };
    ESP_ERROR_CHECK(usb_host_install(&host_config));

    const msc_host_driver_config_t msc_config = {
        .create_backround_task = true,
        .task_priority = MSC_HOST_TASK_PRIORITY,
        .stack_size = 4096,
        .callback = MSC_Event_Callback,
    };
    ESP_ERROR_CHECK(msc_host_install(&msc_config));

    vTaskDelay(10); // Short delay to let MSC client task spin up.

    usb_message_t message = {
        .id = USB_HOST_INITIALIZED};
    xQueueSend(usb_queue, &message, portMAX_DELAY);

    while (true)
    {
        uint32_t event_flags;

        // This function handles all of the USB Host Library's processing and should be called repeatedly in a loop.
        ESP_ERROR_CHECK(usb_host_lib_handle_events(portMAX_DELAY, &event_flags));

        // Release devices once all clients have deregistered.
        if (event_flags & USB_HOST_LIB_EVENT_FLAGS_NO_CLIENTS)
        {
            if (usb_host_device_free_all() == ESP_OK)
            {
                KLOG_INFO(TAG, "USB_HOST_LIB_EVENT_FLAGS_NO_CLIENTS");
                // break;
            };
        }
        if (event_flags & USB_HOST_LIB_EVENT_FLAGS_ALL_FREE)
        {
            KLOG_INFO(TAG, "USB_HOST_LIB_EVENT_FLAGS_ALL_FREE");
            // break;
        }
    }
}

static void app_usb_task(void *args)
{
    SemaphoreHandle_t init_complete = args;

    static USB_State_T Current_State = USB_STATE_UNINITIALIZED;
    static msc_host_device_handle_t msc_device = NULL;
    static msc_host_vfs_handle_t vfs_handle = NULL;
    static uint8_t device_address = 1;
    usb_message_t msg;

    while (true)
    {
        switch (Current_State)
        {
        case USB_STATE_UNINITIALIZED:
        {
            if (xQueueReceive(usb_queue, &msg, portMAX_DELAY) == pdTRUE)
            {
                if (msg.id == USB_HOST_INITIALIZED)
                {
                    KLOG_INFO(TAG, "Host initialized.");
                    Current_State = USB_STATE_IDLE;
                }
            }
        }
        break;

        case USB_STATE_IDLE:
        {
            KLOG_TRACE(TAG, "USB_STATE_IDLE");

            if (xQueueReceive(usb_queue, &msg, pdMS_TO_TICKS(1000)) == pdTRUE)
            {
                if (msg.id == USB_DEVICE_CONNECTED)
                {
                    device_address = msg.data.new_dev_address;
                    Current_State = USB_STATE_DEVICE_CONNECTED;
                    KLOG_INFO(TAG, "Device connected.");
                }
            }
            else
            {
                KLOG_ERROR(TAG, "No flash drive detected--rebooting.");
                vTaskDelay(pdMS_TO_TICKS(100));
                esp_restart();
            }
        }
        break;

        case USB_STATE_DEVICE_CONNECTED:
        {
            KLOG_TRACE(TAG, "USB_STATE_DEVICE_CONNECTED");

            // USB device connected. Open it and attempt to map it to Virtual File System.
            if (msc_host_install_device(device_address, &msc_device) != ESP_OK)
            {
                ESP_LOGW(TAG, "Problem connecting to flash drive.");
                Current_State = USB_STATE_PROCESSING_DISCONNECTION;
            }
            else
            {
                const esp_vfs_fat_mount_config_t mount_config = {
                    .format_if_mount_failed = false,
                    .max_files = 3,
                    .allocation_unit_size = 8192,
                };
                if (msc_host_vfs_register(msc_device, MNT_PATH, &mount_config, &vfs_handle) != ESP_OK)
                {
                    ESP_LOGW(TAG, "Problem mounting filesystem.");
                    Current_State = USB_STATE_PROCESSING_DISCONNECTION;
                }
                else
                {
                    Current_State = USB_STATE_VFS_REGISTERED;
                    xSemaphoreGive(init_complete);
                }
            }
        }
        break;

        case USB_STATE_VFS_REGISTERED:
        {
            KLOG_TRACE(TAG, "USB_STATE_VFS_REGISTERED");

            if (xQueueReceive(usb_queue, &msg, pdMS_TO_TICKS(1000)) == pdTRUE)
            {
                if (msg.id == USB_DEVICE_DISCONNECTED)
                {
                    Current_State = USB_STATE_PROCESSING_DISCONNECTION;
                }
            }
            vTaskDelay(pdMS_TO_TICKS(1000));
        }
        break;

        case USB_STATE_PROCESSING_DISCONNECTION:
        {
            KLOG_TRACE(TAG, "USB_STATE_PROCESSING_DISCONNECTION");

            if (vfs_handle)
            {
                ESP_ERROR_CHECK(msc_host_vfs_unregister(vfs_handle));
                vfs_handle = NULL;
            }
            if (msc_device)
            {
                ESP_ERROR_CHECK(msc_host_uninstall_device(msc_device));
                msc_device = NULL;
                device_address = 0;
            }

            Current_State = USB_STATE_IDLE;
        }
        break;
        }
    }
}

void Initialize_USB(SemaphoreHandle_t init_complete)
{
    KLOG_INFO(TAG, "Initializing USB file system...");

    // Create FreeRTOS primitives
    usb_queue = xQueueCreate(5, sizeof(usb_message_t));
    assert(usb_queue);

    BaseType_t usb_host_task_created = xTaskCreatePinnedToCore(usb_host_task, "USB Host", 4096, NULL, USB_HOST_TASK_PRIORITY, NULL, 1);
    assert(usb_host_task_created);

    TaskHandle_t xHandle = xTaskCreateStaticPinnedToCore(
        app_usb_task,          // Function that implements the task.
        "USB NVM",             // Text name for the task.
        STACK_SIZE,            // Stack size in bytes, not words.
        (void *)init_complete, // Parameter passed into the task.
        tskIDLE_PRIORITY + 2,  // Priority at which the task is created.
        xStack,                // Array to use as the task's stack.
        &xTaskBuffer,          // Variable to hold the task's data structure.
        APP_CPU_NUM);          // Specify the task's core affinity
    assert(xHandle);
}

bool OTA_File_Exists(void)
{
    struct stat s = {0};
    return (stat(OTA_FILE, &s) == 0);
}

char *Get_OTA_Image_URL(void)
{
    static char url[64] = "";
    FILE *f;
    f = fopen(OTA_FILE, "r");
    if (f == NULL)
    {
        KLOG_ERROR(TAG, "The OTA file is missing!");
        return url;
    }

    fgets(url, sizeof(url), f);
    fclose(f);
    // Strip the newline
    char *pos = strchr(url, '\n');
    if (pos)
    {
        *pos = '\0';
    }
    return url;
}

void Erase_OTA_File(void)
{
    if (remove(OTA_FILE) != 0)
    {
        KLOG_ERROR(TAG, "Error erasing the OTA file: %s", strerror(errno));
    }
}