#include "esp_console.h"
#include "esp_log.h"
#include "nvs.h"
#include "nvs_flash.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "linenoise/linenoise.h"
#include "esp_vfs_dev.h"
#include <string.h>
#include <stdio.h>

#define LOG_NVS_NS "logcfg"
#define MAX_TAGS 16
#define MAX_TAGLEN 16

typedef struct
{
    char tag[MAX_TAGLEN];
    uint8_t level;
} tag_entry_t;

static esp_log_level_t Saved_Global_Level = ESP_LOG_INFO;
static tag_entry_t Tag_Table[MAX_TAGS];
static size_t Tag_Count = 0;

/* ---------------- Utility ---------------- */

static bool parse_level(const char *s, esp_log_level_t *out)
{
    if (!strcmp(s, "none"))
        *out = ESP_LOG_NONE;
    else if (!strcmp(s, "error"))
        *out = ESP_LOG_ERROR;
    else if (!strcmp(s, "warn"))
        *out = ESP_LOG_WARN;
    else if (!strcmp(s, "info"))
        *out = ESP_LOG_INFO;
    else if (!strcmp(s, "debug"))
        *out = ESP_LOG_DEBUG;
    else if (!strcmp(s, "verbose"))
        *out = ESP_LOG_VERBOSE;
    else
        return false;
    return true;
}

static const char *level_str(esp_log_level_t lvl)
{
    switch (lvl)
    {
    case ESP_LOG_NONE:
        return "none";
    case ESP_LOG_ERROR:
        return "error";
    case ESP_LOG_WARN:
        return "warn";
    case ESP_LOG_INFO:
        return "info";
    case ESP_LOG_DEBUG:
        return "debug";
    case ESP_LOG_VERBOSE:
        return "verbose";
    default:
        return "?";
    }
}

/* ---------------- NVS ---------------- */

static void nvs_save(void)
{
    nvs_handle_t h;
    esp_err_t err = nvs_open(LOG_NVS_NS, NVS_READWRITE, &h);
    
    if (err == ESP_ERR_NVS_NOT_INITIALIZED) {
        ESP_LOGW("Console", "NVS not initialized, cannot save");
        return;
    }
    else if (err != ESP_OK) {
        ESP_LOGE("Console", "Failed to open NVS: %s", esp_err_to_name(err));
        return;
    }

    uint8_t lvl = esp_log_get_default_level();
    nvs_set_u8(h, "global", lvl);

    nvs_set_u8(h, "tagcnt", Tag_Count);
    nvs_set_blob(h, "tags", Tag_Table,
                 Tag_Count * sizeof(tag_entry_t));

    nvs_commit(h);
    nvs_close(h);
}

static void log_config_load(void)
{
    nvs_handle_t h;
    esp_err_t err = nvs_open(LOG_NVS_NS, NVS_READONLY, &h);
    
    if (err == ESP_ERR_NVS_NOT_INITIALIZED) {
        ESP_LOGW("Console", "NVS not initialized, cannot load");
        return;
    }
    else if (err != ESP_OK) {
        ESP_LOGE("Console", "Failed to open NVS: %s", esp_err_to_name(err));
        return;
    }

    uint8_t lvl;
    if (nvs_get_u8(h, "global", &lvl) == ESP_OK)
    {
        esp_log_level_set("*", lvl);
        Saved_Global_Level = lvl;
    }

    uint8_t cnt = 0;
    if (nvs_get_u8(h, "tagcnt", &cnt) == ESP_OK && cnt <= MAX_TAGS)
    {
        size_t sz = cnt * sizeof(tag_entry_t);
        if (nvs_get_blob(h, "tags", Tag_Table, &sz) == ESP_OK)
        {
            Tag_Count = cnt;
            for (size_t i = 0; i < Tag_Count; i++)
            {
                esp_log_level_set(Tag_Table[i].tag,
                                  Tag_Table[i].level);
            }
        }
    }

    nvs_close(h);
}

static void nvs_clear_tags(void)
{
    nvs_handle_t h;
    esp_err_t err = nvs_open(LOG_NVS_NS, NVS_READWRITE, &h);
    
    if (err == ESP_ERR_NVS_NOT_INITIALIZED) {
        ESP_LOGW("Console", "NVS not initialized, cannot clear tags");
        return;
    }
    else if (err != ESP_OK) {
        ESP_LOGE("Console", "Failed to open NVS: %s", esp_err_to_name(err));
        return;
    }

    nvs_erase_key(h, "tagcnt");
    nvs_erase_key(h, "tags");
    nvs_commit(h);
    nvs_close(h);
}

static void clear_tag_levels(void)
{
    esp_log_level_t default_level = esp_log_get_default_level();
    
    for (size_t i = 0; i < Tag_Count; i++)
    {
        esp_log_level_set(Tag_Table[i].tag, default_level);
    }
    Tag_Count = 0;
}

/* ---------------- Tag table ---------------- */

static void set_tag_level(const char *tag, esp_log_level_t level)
{
    for (size_t i = 0; i < Tag_Count; i++)
    {
        if (!strcmp(Tag_Table[i].tag, tag))
        {
            Tag_Table[i].level = level;
            return;
        }
    }

    if (Tag_Count < MAX_TAGS)
    {
        strncpy(Tag_Table[Tag_Count].tag, tag, MAX_TAGLEN - 1);
        Tag_Table[Tag_Count].tag[MAX_TAGLEN - 1] = 0;
        Tag_Table[Tag_Count].level = level;
        Tag_Count++;
    }
}

/* ---------------- Console command ---------------- */

static void print_usage(void)
{
    printf(
        "Usage:\n"
        "  log show\n"
        "  log clear\n"
        "  log off | on\n"
        "  log <level>\n"
        "  log tag <tag|*> <level|off>\n"
        "Levels: none error warn info debug verbose\n");
}

static int cmd_log(int argc, char **argv)
{
    esp_log_level_t level;

    if (argc < 2)
    {
        print_usage();
        return 0;
    }

    if (!strcmp(argv[1], "show"))
    {
        printf("Global level: %s\n",
               level_str(esp_log_get_default_level()));

        if (Tag_Count == 0)
        {
            printf("No tag overrides\n");
        }
        else
        {
            printf("Tag overrides:\n");
            for (size_t i = 0; i < Tag_Count; i++)
            {
                printf("  %-10s : %s\n",
                       Tag_Table[i].tag,
                       level_str(Tag_Table[i].level));
            }
        }
        return 0;
    }

    if (!strcmp(argv[1], "clear"))
    {
        clear_tag_levels();
        nvs_clear_tags();
        printf("All tag overrides cleared\n");
        return 0;
    }

    if (!strcmp(argv[1], "off"))
    {
        Saved_Global_Level = esp_log_get_default_level();
        esp_log_level_set("*", ESP_LOG_NONE);
        nvs_save();
        printf("Global logging disabled\n");
        return 0;
    }

    if (!strcmp(argv[1], "on"))
    {
        esp_log_level_set("*", Saved_Global_Level);
        printf("Global logging restored\n");
        return 0;
    }

    if (!strcmp(argv[1], "tag"))
    {
        if (argc != 4)
        {
            print_usage();
            return 0;
        }

        const char *tag = argv[2];
        const char *lvl = argv[3];

        if (!strcmp(lvl, "off"))
        {
            level = ESP_LOG_NONE;
        }
        else if (!parse_level(lvl, &level))
        {
            printf("Invalid level\n");
            return 0;
        }

        esp_log_level_set(tag, level);
        set_tag_level(tag, level);
        nvs_save();

        printf("Tag '%s' set to %s\n", tag, lvl);
        return 0;
    }

    if (!parse_level(argv[1], &level))
    {
        print_usage();
        return 0;
    }

    esp_log_level_set("*", level);
    nvs_save();
    printf("Global level set to %s\n", argv[1]);
    return 0;
}

/* ---------------- Registration ---------------- */

static void register_log_command(void)
{
    const esp_console_cmd_t cmd = {
        .command = "log",
        .help = "Control logging levels",
        .func = &cmd_log,
    };

    ESP_ERROR_CHECK(esp_console_cmd_register(&cmd));
}

void Initialize_Console(void)
{
    esp_console_repl_t *repl = NULL;
    esp_console_repl_config_t repl_config = ESP_CONSOLE_REPL_CONFIG_DEFAULT();
    repl_config.prompt = "KTag>";
    repl_config.max_cmdline_length = 1024;

    log_config_load();
    register_log_command();
    esp_console_register_help_command();

    esp_console_dev_uart_config_t hw_config = ESP_CONSOLE_DEV_UART_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_console_new_repl_uart(&hw_config, &repl_config, &repl));
    ESP_ERROR_CHECK(esp_console_start_repl(repl));
}