#include <stdio.h>
#include <string.h>
#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_heap_caps.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "nvs_flash.h"
#include "esp_log.h"
#include "esp_vfs.h"
#include "esp_littlefs.h"
#include <dirent.h>

#include "paper.h"
#include "server.h"

#define WIFI_SCAN_LIST_SIZE 10
#define LUA_FILE_PATH "/assets"

#ifndef MIN
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
#endif

static const char *TAG = "inkpot";

// Function to log memory usage with the message at the end
void log_memory_usage(const char *message) {
    ESP_LOGI(TAG, "Free heap: %d, Min free heap: %d, Largest free block: %d, %s",
             heap_caps_get_free_size(MALLOC_CAP_DEFAULT),
             heap_caps_get_minimum_free_size(MALLOC_CAP_DEFAULT),
             heap_caps_get_largest_free_block(MALLOC_CAP_DEFAULT),
             message);
}

// Initialize and mount the filesystem
void init_filesystem() {
    ESP_LOGI(TAG, "Initializing File System");

    esp_vfs_littlefs_conf_t conf = {
        .base_path = LUA_FILE_PATH,
        .partition_label = "assets",
        .format_if_mount_failed = false,
        .dont_mount = false,
    };

    esp_err_t err = esp_vfs_littlefs_register(&conf);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to mount or format filesystem");
    } else {
        ESP_LOGI(TAG, "Filesystem mounted at %s", LUA_FILE_PATH);
    }
}

// Function to run a Lua script from file
void run_lua_file(const char *file_name, const char *test_name) {
    ESP_LOGI(TAG, "Starting Lua test from file: %s", test_name);

    log_memory_usage("Start of test");

    lua_State *L = luaL_newstate();
    if (L == NULL) {
        ESP_LOGE(TAG, "Failed to create new Lua state");
        return;
    }
    log_memory_usage("After luaL_newstate");

    luaL_openlibs(L);

    // Set the Lua module search path to include the assets directory
    if (luaL_dostring(L, "package.path = package.path .. ';./?.lua;/assets/?.lua'")) {
        ESP_LOGE(TAG, "Failed to set package.path: %s", lua_tostring(L, -1));
        lua_pop(L, 1);  // Remove error message from the stack
    }

    log_memory_usage("After luaL_openlibs");

    // load `paper` library and expose it globally
    luaL_requiref(L, "paper", luaopen_paper, 1);
    lua_pop(L, 1);
    log_memory_usage("After adding paper");

    // Construct the full file path
    char full_path[128];
    snprintf(full_path, sizeof(full_path), LUA_FILE_PATH"/%s", file_name);

    if (luaL_dofile(L, full_path) == LUA_OK) {
        lua_pop(L, lua_gettop(L));
    } else {
        ESP_LOGE(TAG, "Error running Lua script from file '%s': %s", full_path, lua_tostring(L, -1));
        lua_pop(L, 1);  // Remove error message from the stack
    }
    log_memory_usage("After executing Lua script from file");

    lua_close(L);
    log_memory_usage("After lua_close");

    ESP_LOGI(TAG, "End of Lua test from file: %s", test_name);
}

// Function to run an embedded Lua script
void run_lua_string(const char *lua_script, const char *script_name) {
    ESP_LOGI(TAG, "Running lua script: %s", script_name);
    ESP_LOGI(TAG, "%s", lua_script);

    log_memory_usage("Start running script");

    lua_State *L = luaL_newstate();
    if (L == NULL) {
        ESP_LOGE(TAG, "Failed to create new Lua state");
        return;
    }
    log_memory_usage("After luaL_newstate");

    luaL_openlibs(L);
    log_memory_usage("After luaL_openlibs");

    // load `paper` library and expose it globally
    luaL_requiref(L, "paper", luaopen_paper, 1);
    lua_pop(L, 1);
    log_memory_usage("After adding paper");

    if (luaL_dostring(L, lua_script) == LUA_OK) {
        lua_pop(L, lua_gettop(L));
    } else {
        ESP_LOGE(TAG, "Error running embedded Lua script: %s", lua_tostring(L, -1));
        lua_pop(L, 1);  // Remove error message from the stack
    }
    log_memory_usage("After executing Lua script");

    lua_close(L);
    log_memory_usage("After lua_close");

    ESP_LOGI(TAG, "End of Lua script: %s", script_name);
}

// Function to scan Wi-Fi networks
void scan_wifi_networks(void) {
    ESP_LOGI(TAG, "Starting Wi-Fi scan...");

    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_LOGW(TAG, "NVS Flash init error (%s), erasing...", esp_err_to_name(ret));
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_start());

    wifi_scan_config_t scan_config = {
        .ssid = NULL,
        .bssid = NULL,
        .channel = 0,
        .show_hidden = true
    };

    ESP_ERROR_CHECK(esp_wifi_scan_start(&scan_config, true));

    uint16_t ap_count = WIFI_SCAN_LIST_SIZE;
    wifi_ap_record_t ap_info[WIFI_SCAN_LIST_SIZE];
    ESP_ERROR_CHECK(esp_wifi_scan_get_ap_records(&ap_count, ap_info));

    ESP_LOGI(TAG, "Found %d access points:", ap_count);
    for (int i = 0; i < ap_count; i++) {
        ESP_LOGI(TAG, "SSID: %s, RSSI: %d", ap_info[i].ssid, ap_info[i].rssi);
    }

    ESP_ERROR_CHECK(esp_wifi_stop());
    ESP_ERROR_CHECK(esp_wifi_deinit());

    ESP_LOGI(TAG, "Wi-Fi scan completed.");
}

// HTTP Server

#define WRITE_HEADER(req, buffer, name, format, src) \
    sprintf(buffer, format, src);                    \
    ESP_ERROR_CHECK(httpd_resp_set_hdr(req, name, buffer));

static esp_err_t http_index(httpd_req_t* req) {
    // TODO: Serve HTML file with form POSTing to `/draw`
    const char* response = "Hello world!\n";
    httpd_resp_set_type(req, "text/plain");
    httpd_resp_set_status(req, "200");
    httpd_resp_send(req, response, HTTPD_RESP_USE_STRLEN);
    return ESP_OK;
}

esp_err_t http_draw(httpd_req_t* req) {
    // READING STREAM
    int max_post_size = 4 * 1024;
    char* buf = (char*)heap_caps_malloc(max_post_size, MALLOC_CAP_SPIRAM);
    for (int i = 0; i < max_post_size; i++) {
        buf[i] = 0;
    }

    if (buf == NULL) {
        char msg[50];
        sprintf(msg, "Failed to allocate %d chars\n", max_post_size);
        httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, msg);
        return ESP_ERR_INVALID_ARG;
    }

    // READING STREAM
    int ret, remaining = req->content_len;

    while (remaining > 0) {
        /* Read the data for the request */
        if ((ret = httpd_req_recv(req, buf,
                        MIN(remaining, max_post_size))) <= 0) {
            if (ret == HTTPD_SOCK_ERR_TIMEOUT) {
                /* Retry receiving if timeout occurred */
                continue;
            }
            return ESP_FAIL;
        }

        /* Send back the same data */
        httpd_resp_send_chunk(req, buf, ret);
        remaining -= ret;

        /* Log data received */
        ESP_LOGI(TAG, "=========== RECEIVED DATA ==========");
        ESP_LOGI(TAG, "%.*s", ret, buf);
        ESP_LOGI(TAG, "====================================");
    }

    // TODO: Error handling
    // TODO: Ensure that `run_lua_string` is executed in an uninterrupted task;
    // this means we need to set the task prio of the current request.
    // Execute the Lua script received as post body; this is executed as a
    // FreeRTOS task that will not be interrupted
    run_lua_string(buf, "E-Paper Script via HTTP");
    heap_caps_free(buf);

    // Done reading
    char response[100];
    sprintf(
        response, "script drawn!"
    );
    httpd_resp_set_type(req, "text/plain");
    httpd_resp_set_status(req, "200");
    httpd_resp_send(req, response, HTTPD_RESP_USE_STRLEN);
    return ESP_OK;
}

void register_http_routes(httpd_handle_t server) {
    {
        httpd_uri_t uri
            = { .uri = "/", .method = HTTP_GET, .handler = http_index, .user_ctx = NULL };
        httpd_register_uri_handler(server, &uri);
    }
    {
        httpd_uri_t uri
            = { .uri = "/draw", .method = HTTP_POST, .handler = http_draw, .user_ctx = NULL };
        httpd_register_uri_handler(server, &uri);
    }
}

// init

void app_main(void) {
    // Initialize and mount the filesystem
    init_filesystem();

    // Initialize NVS, needed for wifi
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    // FIXME: Wifi server causes some crash, failing to yield to watchdog in time
    // Set up HTTP server
    httpd_handle_t server = get_server();
    if (server != NULL) {
        register_http_routes(server);
    }

    // Run script in assets/boot.lua; this is executed as a FreeRTOS task
    // that may not be interrupted
    void runLuaFile (void* arg) {
        run_lua_file("boot.lua", "E-Paper Startup Script");
        vTaskDelete(NULL);
    }
    xTaskCreate(runLuaFile, "run_lua_file", 4096, NULL, tskIDLE_PRIORITY, NULL);

    ESP_LOGI(TAG, "End of application.");

    // Prevent the task from ending
    while (1) {
        vTaskDelay(pdMS_TO_TICKS(1000));
    }
}