#include <efi.h>
#include "commands.h"
#include "string_utils.h"
#include "memory.h"
#include "task.h"

#define SAFE_PRINT(Boot, ...)                           \
    do {                                                \
        if ((Boot) != NULL && (Boot)->print != NULL) {  \
            (Boot)->print(__VA_ARGS__);                 \
        }                                               \
    } while (0)

// Forward declarations
static void cmd_shutdown(BootInfo *Boot, CHAR16 *Args);
static void cmd_help(BootInfo *Boot, CHAR16 *Args);
static void cmd_man(BootInfo *Boot, CHAR16 *Args);
static void cmd_clear(BootInfo *Boot, CHAR16 *Args);
static void cmd_about(BootInfo *Boot, CHAR16 *Args);
static void cmd_mem(BootInfo *Boot, CHAR16 *Args);
static void cmd_ps(BootInfo *Boot, CHAR16 *Args);
static void cmd_spawn(BootInfo *Boot, CHAR16 *Args);
static void cmd_memtest(BootInfo *Boot, CHAR16 *Args);
static void cmd_tasktest(BootInfo *Boot, CHAR16 *Args);

// Command registry
static Command commands[] = {
    {
        L"shutdown",
        L"Shutdown the system",
        L"Usage: shutdown\n\r  Initiates a system shutdown using UEFI runtime services.",
        cmd_shutdown
    },
    {
        L"help",
        L"Display available commands",
        L"Usage: help\n\r  Lists all available commands with brief descriptions.",
        cmd_help
    },
    {
        L"man",
        L"Display manual page for a command",
        L"Usage: man <command>\n\r  Shows detailed help for the specified command.",
        cmd_man
    },
    {
        L"clear",
        L"Clear the screen",
        L"Usage: clear\n\r  Clears the console screen.",
        cmd_clear
    },
    {
        L"about",
        L"Display system information",
        L"Usage: about\n\r  Shows information about this operating system.",
        cmd_about
    },
    {
        L"mem",
        L"Display memory statistics",
        L"Usage: mem\n\r  Shows physical memory, heap, and paging information.",
        cmd_mem
    },
    {
        L"ps",
        L"List running tasks",
        L"Usage: ps\n\r  Displays all active tasks with PID, state, and name.",
        cmd_ps
    },
    {
        L"spawn",
        L"Spawn a demo background task",
        L"Usage: spawn [name]\n\r  Creates a cooperative demo task.\n\r  Optional argument sets the task name.",
        cmd_spawn
    },
    {
        L"memtest",
        L"Test memory allocation and deallocation",
        L"Usage: memtest\n\r  Allocates and frees heap and page memory to verify\n\r  the memory manager is working correctly.",
        cmd_memtest
    },
    {
        L"tasktest",
        L"Test task scheduler with multiple tasks",
        L"Usage: tasktest\n\r  Spawns several concurrent tasks that run cooperatively\n\r  and report their progress, demonstrating context switching.",
        cmd_tasktest
    },
    {NULL, NULL, NULL, NULL}  // Sentinel
};

static void request_shutdown(BootInfo *Boot)
{
    if (Boot == NULL) {
        return;
    }

    if (Boot->shutdown != NULL) {
        Boot->shutdown();
        return;
    }

    if (Boot->SystemTable != NULL &&
        Boot->SystemTable->RuntimeServices != NULL &&
        Boot->SystemTable->RuntimeServices->ResetSystem != NULL) {
        Boot->SystemTable->RuntimeServices->ResetSystem(
            EfiResetShutdown, EFI_SUCCESS, 0, NULL);
        return;
    }

    SAFE_PRINT(Boot, L"Shutdown service unavailable.\n\r");
}

static void cmd_shutdown(BootInfo *Boot, CHAR16 *Args)
{
    (void)Args;  // Unused
    SAFE_PRINT(Boot, L"Shutting down...\n\r");
    request_shutdown(Boot);
}

static void cmd_help(BootInfo *Boot, CHAR16 *Args)
{
    (void)Args;  // Unused
    show_help(Boot);
}

static void cmd_man(BootInfo *Boot, CHAR16 *Args)
{
    UINTN i = 0;
    
    if (Args == NULL || Args[0] == L'\0') {
        SAFE_PRINT(Boot, L"Usage: man <command>\n\r");
        SAFE_PRINT(Boot, L"Try 'help' for a list of available commands.\n\r");
        return;
    }

    trim_spaces_inplace(Args);
    
    for (i = 0; commands[i].name != NULL; i++) {
        if (ascii_streq_ci(Args, commands[i].name)) {
            SAFE_PRINT(Boot, L"\n\r");
            SAFE_PRINT(Boot, L"NAME\n\r");
            SAFE_PRINT(Boot, L"  %s - %s\n\r\n\r", commands[i].name, commands[i].description);
            SAFE_PRINT(Boot, L"DESCRIPTION\n\r");
            SAFE_PRINT(Boot, L"  %s\n\r", commands[i].usage);
            SAFE_PRINT(Boot, L"\n\r");
            return;
        }
    }

    SAFE_PRINT(Boot, L"No manual entry for '%s'\n\r", Args);
    SAFE_PRINT(Boot, L"Try 'help' for a list of available commands.\n\r");
}

static void cmd_clear(BootInfo *Boot, CHAR16 *Args)
{
    EFI_STATUS Status;
    (void)Args;  // Unused
    
    if (Boot != NULL && Boot->clear_screen != NULL) {
        Status = Boot->clear_screen();
        if (EFI_ERROR(Status)) {
            SAFE_PRINT(Boot, L"Failed to clear screen: %r\n\r", Status);
        }
    } else {
        SAFE_PRINT(Boot, L"Clear screen function unavailable.\n\r");
    }
}

static void cmd_about(BootInfo *Boot, CHAR16 *Args)
{
    (void)Args;  // Unused
    
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"================================================\n\r");
    SAFE_PRINT(Boot, L"  Simple UEFI Operating System\n\r");
    SAFE_PRINT(Boot, L"================================================\n\r");
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"A minimal bootable UEFI operating system written in C.\n\r");
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Features:\n\r");
    SAFE_PRINT(Boot, L"  - UEFI Boot\n\r");
    SAFE_PRINT(Boot, L"  - Console I/O\n\r");
    SAFE_PRINT(Boot, L"  - ELF64 Kernel Loader\n\r");
    SAFE_PRINT(Boot, L"  - Memory Management (PMM + Heap)\n\r");
    SAFE_PRINT(Boot, L"  - Cooperative Multitasking\n\r");
    SAFE_PRINT(Boot, L"  - Interactive Command Interface\n\r");
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Type 'help' for available commands.\n\r");
    SAFE_PRINT(Boot, L"\n\r");
}

static void cmd_mem(BootInfo *Boot, CHAR16 *Args)
{
    (void)Args;
    memory_print_stats(Boot);
}

/* ----------------------------------------------------------------
 *  Demo task – runs cooperatively, prints progress, then exits
 * ---------------------------------------------------------------- */

static void demo_task_fn(void *arg)
{
    BootInfo *Boot = (BootInfo *)arg;
    Task *self = task_current();
    UINTN i;

    SAFE_PRINT(Boot, L"[%s:%d] started\n\r", self->name, self->pid);

    for (i = 0; i < 5; i++) {
        SAFE_PRINT(Boot, L"[%s:%d] tick %d/5\n\r",
                   self->name, self->pid, i + 1);
        task_yield();
    }

    SAFE_PRINT(Boot, L"[%s:%d] finished\n\r", self->name, self->pid);
}

static void cmd_ps(BootInfo *Boot, CHAR16 *Args)
{
    (void)Args;
    task_print_list(Boot);
}

static void cmd_spawn(BootInfo *Boot, CHAR16 *Args)
{
    Task *t;
    const CHAR16 *name;

    if (Args != NULL && Args[0] != L'\0') {
        trim_spaces_inplace(Args);
        name = Args;
    } else {
        name = L"demo";
    }

    t = task_create(name, demo_task_fn, Boot);
    if (t == NULL) {
        SAFE_PRINT(Boot, L"Failed to create task (out of slots or memory).\n\r");
        return;
    }

    SAFE_PRINT(Boot, L"Created task '%s' (PID %d)\n\r", t->name, t->pid);
}

/* ----------------------------------------------------------------
 *  memtest – exercise the heap and PMM allocators
 * ---------------------------------------------------------------- */

static void cmd_memtest(BootInfo *Boot, CHAR16 *Args)
{
    void *ptrs[8];
    UINTN sizes[] = { 16, 64, 128, 256, 512, 1024, 2048, 4096 };
    UINTN i;
    UINT64 page;
    UINTN h_total, h_used, h_free, h_blocks;
    (void)Args;

    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Memory Test\n\r");
    SAFE_PRINT(Boot, L"================================================\n\r");
    SAFE_PRINT(Boot, L"\n\r");

    /* --- Heap allocation test --- */
    SAFE_PRINT(Boot, L"[1] Heap allocation test\n\r");
    for (i = 0; i < 8; i++) {
        ptrs[i] = kmalloc(sizes[i]);
        if (ptrs[i] != NULL) {
            SAFE_PRINT(Boot, L"  kmalloc(%4d) = 0x%lx  OK\n\r",
                       sizes[i], (UINT64)(UINTN)ptrs[i]);
        } else {
            SAFE_PRINT(Boot, L"  kmalloc(%4d) = NULL    FAIL\n\r", sizes[i]);
        }
    }

    heap_get_stats(&h_total, &h_used, &h_free, &h_blocks);
    SAFE_PRINT(Boot, L"  Heap after alloc: used=%d  free=%d  blocks=%d\n\r",
               h_used, h_free, h_blocks);

    /* --- Heap free test --- */
    SAFE_PRINT(Boot, L"\n\r[2] Heap free test\n\r");
    for (i = 0; i < 8; i++) {
        if (ptrs[i] != NULL) {
            kfree(ptrs[i]);
            SAFE_PRINT(Boot, L"  kfree(0x%lx)  OK\n\r",
                       (UINT64)(UINTN)ptrs[i]);
        }
    }

    heap_get_stats(&h_total, &h_used, &h_free, &h_blocks);
    SAFE_PRINT(Boot, L"  Heap after free:  used=%d  free=%d  blocks=%d\n\r",
               h_used, h_free, h_blocks);

    /* --- PMM page allocation test --- */
    SAFE_PRINT(Boot, L"\n\r[3] PMM page allocation test\n\r");
    SAFE_PRINT(Boot, L"  Free pages before: %d\n\r", pmm_get_free_pages());

    page = pmm_alloc_page();
    if (page != 0) {
        SAFE_PRINT(Boot, L"  pmm_alloc_page() = 0x%lx  OK\n\r", page);
        SAFE_PRINT(Boot, L"  Free pages after alloc: %d\n\r", pmm_get_free_pages());
        pmm_free_page(page);
        SAFE_PRINT(Boot, L"  pmm_free_page()  OK\n\r");
        SAFE_PRINT(Boot, L"  Free pages after free:  %d\n\r", pmm_get_free_pages());
    } else {
        SAFE_PRINT(Boot, L"  pmm_alloc_page() = 0  FAIL (out of pages)\n\r");
    }

    /* --- Multi-page allocation test --- */
    SAFE_PRINT(Boot, L"\n\r[4] PMM multi-page allocation test (4 pages)\n\r");
    page = pmm_alloc_pages(4);
    if (page != 0) {
        SAFE_PRINT(Boot, L"  pmm_alloc_pages(4) = 0x%lx  OK\n\r", page);
        SAFE_PRINT(Boot, L"  Free pages after alloc: %d\n\r", pmm_get_free_pages());
        pmm_free_pages(page, 4);
        SAFE_PRINT(Boot, L"  pmm_free_pages()  OK\n\r");
        SAFE_PRINT(Boot, L"  Free pages after free:  %d\n\r", pmm_get_free_pages());
    } else {
        SAFE_PRINT(Boot, L"  pmm_alloc_pages(4) = 0  FAIL\n\r");
    }

    SAFE_PRINT(Boot, L"\n\rAll memory tests completed.\n\r\n\r");
}

/* ----------------------------------------------------------------
 *  tasktest – spawn multiple concurrent tasks to exercise scheduler
 * ---------------------------------------------------------------- */

static void worker_task_fn(void *arg)
{
    BootInfo *Boot = (BootInfo *)arg;
    Task *self = task_current();
    UINTN i;

    for (i = 0; i < 3; i++) {
        SAFE_PRINT(Boot, L"  [%s:%d] working... step %d/3\n\r",
                   self->name, self->pid, i + 1);
        task_yield();
    }

    SAFE_PRINT(Boot, L"  [%s:%d] done\n\r", self->name, self->pid);
}

static void cmd_tasktest(BootInfo *Boot, CHAR16 *Args)
{
    Task *t1, *t2, *t3;
    UINTN i;
    (void)Args;

    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Task Scheduler Test\n\r");
    SAFE_PRINT(Boot, L"================================================\n\r");
    SAFE_PRINT(Boot, L"\n\r");

    SAFE_PRINT(Boot, L"Spawning 3 worker tasks...\n\r\n\r");

    t1 = task_create(L"worker-A", worker_task_fn, Boot);
    t2 = task_create(L"worker-B", worker_task_fn, Boot);
    t3 = task_create(L"worker-C", worker_task_fn, Boot);

    if (t1 == NULL || t2 == NULL || t3 == NULL) {
        SAFE_PRINT(Boot, L"Failed to create one or more tasks.\n\r");
        return;
    }

    SAFE_PRINT(Boot, L"Created: %s (PID %d), %s (PID %d), %s (PID %d)\n\r",
               t1->name, t1->pid, t2->name, t2->pid, t3->name, t3->pid);
    SAFE_PRINT(Boot, L"\n\rYielding to let workers run:\n\r\n\r");

    /* Yield enough times for all workers to complete (3 tasks x 3 steps) */
    for (i = 0; i < 12; i++) {
        task_yield();
    }

    SAFE_PRINT(Boot, L"\n\rTask list after test:\n\r");
    task_print_list(Boot);

    SAFE_PRINT(Boot, L"Task scheduler test completed.\n\r\n\r");
}

void show_help(BootInfo *Boot)
{
    UINTN i = 0;
    UINTN max_len = 0;
    UINTN name_len = 0;
    UINTN padding = 0;
    
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"Available commands:\n\r");
    SAFE_PRINT(Boot, L"\n\r");
    
    // Find longest command name for formatting
    for (i = 0; commands[i].name != NULL; i++) {
        name_len = 0;
        while (commands[i].name[name_len] != L'\0') {
            name_len++;
        }
        if (name_len > max_len) {
            max_len = name_len;
        }
    }
    
    for (i = 0; commands[i].name != NULL; i++) {
        SAFE_PRINT(Boot, L"  %s", commands[i].name);
        
        // Add padding
        name_len = 0;
        while (commands[i].name[name_len] != L'\0') {
            name_len++;
        }
        for (padding = 0; padding < (max_len - name_len + 2); padding++) {
            SAFE_PRINT(Boot, L" ");
        }
        
        SAFE_PRINT(Boot, L"- %s\n\r", commands[i].description);
    }
    
    SAFE_PRINT(Boot, L"\n\r");
    SAFE_PRINT(Boot, L"For detailed help on a command, type: man <command>\n\r");
    SAFE_PRINT(Boot, L"\n\r");
}

void execute_command(BootInfo *Boot, CHAR16 *Input)
{
    CHAR16 *cmd_start = NULL;
    CHAR16 *args_start = NULL;
    UINTN i = 0;
    
    if (Boot == NULL || Input == NULL) {
        return;
    }

    trim_spaces_inplace(Input);
    
    if (Input[0] == L'\0') {
        return;
    }

    // Split command and arguments
    cmd_start = Input;
    args_start = Input;
    
    // Find first space
    while (*args_start != L'\0' && !is_space16(*args_start)) {
        args_start++;
    }
    
    // If we found a space, null-terminate command and advance to args
    if (*args_start != L'\0') {
        *args_start = L'\0';
        args_start++;
        
        // Skip leading spaces in args
        while (*args_start != L'\0' && is_space16(*args_start)) {
            args_start++;
        }
    }
    
    // Search for command
    for (i = 0; commands[i].name != NULL; i++) {
        if (ascii_streq_ci(cmd_start, commands[i].name)) {
            commands[i].handler(Boot, args_start);
            return;
        }
    }

    // Command not found
    SAFE_PRINT(Boot, L"Unknown command: %s\n\r", cmd_start);
    SAFE_PRINT(Boot, L"Type 'help' for a list of available commands.\n\r");
}