/* ============================================================================
 * Pactor64 — Interactive Shell (x86-64)
 * Serial-based line editor with built-in commands
 * Includes "run" command for executing ELF64 userspace programs
 * ============================================================================ */
#include "../include/pactor64.h"
#include "ata.h"
/* Forward declaration */
static void cmd_run_elf(const unsigned char *elf_data, unsigned int elf_len, const char *name);
#include "hello_elf.h"
#include "pactor64_llc_bin.h"
#include "raw_test_bin.h"
#include "raw_stack_test_bin.h"
#include "hello_v2_bin.h"
#include "test_prog_bin.h"
#include "app_binaries.h"
#include "ext2.h"
#include "elf.h"

/* Embedded hello.elf binary (generated by xxd) */

/* Assembly functions for userspace transitions */
extern void enter_userspace(uint64_t user_rip, uint64_t user_rsp);
extern uint64_t kernel_stack_top;

#define CMD_BUF_SIZE 256

/* === Page table constants === */
#define PAGE_PRESENT    0x01
#define PAGE_WRITE      0x02
#define PAGE_USER       0x04    /* U/S bit: 1 = user accessible */
#define PAGE_SIZE_FLAG  0x80    /* PS bit for 2MB pages */

/* User stack: 64KB at virtual address 0x10000000 (256MB, well above kernel/heap) */
#define USER_STACK_BASE  0x10000000
#define USER_STACK_SIZE  0x10000  /* 64KB */
#define USER_STACK_TOP   (USER_STACK_BASE + USER_STACK_SIZE)

/* ============================================================================
 * Page table helpers — make pages user-accessible for ring 3 execution
 * The kernel uses 2MB pages with identity mapping (first 1GB).
 * We need to set the U/S bit on PML4, PDPT, and specific PD entries.
 * ============================================================================ */

/* Set the U/S bit on a 2MB PD entry covering a given virtual address */
static void pt_make_user_accessible(uint64_t vaddr) {
    uint64_t cr3;
    __asm__ volatile ("mov %%cr3, %0" : "=r"(cr3));

    /* PML4 index: bits [47:39] of vaddr */
    uint64_t pml4_idx = (vaddr >> 39) & 0x1FF;
    /* PDPT index: bits [38:30] */
    uint64_t pdpt_idx = (vaddr >> 30) & 0x1FF;
    /* PD index: bits [29:21] */
    uint64_t pd_idx   = (vaddr >> 21) & 0x1FF;

    uint64_t *pml4 = (uint64_t *)(cr3 & ~0xFFFULL);

    /* Set U/S on PML4 entry */
    pml4[pml4_idx] |= PAGE_USER;

    /* Get PDPT and set U/S */
    uint64_t *pdpt = (uint64_t *)(pml4[pml4_idx] & ~0xFFFULL);
    pdpt[pdpt_idx] |= PAGE_USER;

    /* Get PD and set U/S on the 2MB page entry */
    uint64_t *pd = (uint64_t *)(pdpt[pdpt_idx] & ~0xFFFULL);
    pd[pd_idx] |= PAGE_USER;

    kprintf("    PD[%d] for vaddr 0x%x → flags 0x%x\n", pd_idx, vaddr, pd[pd_idx]);
}

/* Flush the TLB by reloading CR3 */
static void tlb_flush(void) {
    uint64_t cr3;
    __asm__ volatile ("mov %%cr3, %0" : "=r"(cr3));
    __asm__ volatile ("mov %0, %%cr3" : : "r"(cr3) : "memory");
}

/* === Built-in Commands === */
static void cmd_ata_test(void) {
    uint8_t buf[512];
    kprintf("  Reading sector 0 (MBR)...\n");
    if (ata_read_sectors(0, 1, buf) != 0) {
        kprintf("  ERROR: Failed to read sector 0\n");
        return;
    }
    kprintf("  First 32 bytes of sector 0:\n    ");
    for (int i = 0; i < 32; i++) {
        if (buf[i] < 0x10) serial_putc('0');
        kprintf("%x ", buf[i]);
        if ((i & 15) == 15 && i < 31) kprintf("\n    ");
    }
    kprintf("\n");
    kprintf("  MBR signature: 0x%x 0x%x\n", buf[510], buf[511]);
}

static void cmd_diskinfo(void) {
    if (!ata_is_present()) {
        kprintf("  No ATA drive detected.\n");
        return;
    }
    kprintf("  ATA Drive Info:\n");
    kprintf("    Model:    %s\n", ata_get_model());
    kprintf("    Serial:   %s\n", ata_get_serial());
    kprintf("    Sectors:  %d\n", ata_get_sector_count());
    kprintf("    Capacity: %d MB\n", ata_get_sector_count() * 512 / (1024 * 1024));
}

/* Directory listing callback for ext2 */
static void dir_list_cb(const char *name, uint8_t name_len,
                         uint32_t inode, uint8_t file_type, void *ctx) {
    (void)ctx;
    const char *type_str = "?";
    switch (file_type) {
        case EXT2_FT_REG_FILE: type_str = "F"; break;
        case EXT2_FT_DIR:      type_str = "D"; break;
        case EXT2_FT_SYMLINK:  type_str = "L"; break;
    }
    /* kprintf doesn't support %.*s, so null-terminate into a temp buffer */
    char name_buf[256];
    kmemcpy(name_buf, name, name_len);
    name_buf[name_len] = '\0';
    kprintf("    [%s] inode=%d  %s\n", type_str, inode, name_buf);
}

static void cmd_mount(void) {
    static ext2_fs_t fs;
    /* MBR partition 1 starts at sector 2048 (default in mkdisk.sh) */
    if (ext2_mount(&fs, 2048) != 0) {
        kprintf("  Failed to mount ext2\n");
        return;
    }
    kprintf("  ext2 mounted successfully!\n");
    kprintf("  Volume: %.16s\n", fs.sb.s_volume_name);

    /* Try to list root directory */
    kprintf("  Root directory:\n");
    ext2_readdir(&fs, EXT2_ROOT_INO, dir_list_cb, NULL);
}

static void cmd_help(void) {
    kprintf("  Available commands:\n");
    kprintf("    help     - Show this help\n");
    kprintf("    run      - Run hello.elf in userspace (ring 3)\n");
    kprintf("    meminfo  - Memory information\n");
    kprintf("    ticks    - Timer tick count\n");
    kprintf("    ata      - ATA drive info\n");
    kprintf("    ata_test - Read sector 0 (MBR)\n");
    kprintf("    mount    - Mount ext2 filesystem\n");
    kprintf("    reboot   - Reboot the system\n");
    kprintf("    halt     - Halt the system\n");
}

static void cmd_meminfo(void) {
    uint64_t total = pmm_get_total_pages();
    uint64_t free  = pmm_get_free_pages();
    uint64_t used  = total - free;
    kprintf("  Physical Memory:\n");
    kprintf("    Total pages: %d (%d MB)\n", total, total * 4 / 1024);
    kprintf("    Free pages:  %d (%d MB)\n", free,  free * 4 / 1024);
    kprintf("    Used pages:  %d (%d MB)\n", used,  used * 4 / 1024);
}

static void cmd_ticks(void) {
    kprintf("  Timer ticks:   %d\n", timer_get_ticks());
    kprintf("  Uptime (sec):  %d\n", timer_get_seconds());
}

static void cmd_reboot(void) {
    kprintf("  Rebooting...\n");
    /* Triple fault: load an invalid IDT then trigger an interrupt */
    struct idtr bad = { 0, 0 };
    idt_load(&bad);
    __asm__ volatile ("int $0x03");
    /* Fallback: keyboard controller reset */
    outb(0x64, 0xFE);
    for (;;) hlt();
}

static void cmd_halt(void) {
    kprintf("  System halted.\n");
    cli();
    for (;;) hlt();
}

/* ============================================================================
 * cmd_run — Load and execute hello.elf in userspace (ring 3)
 *
 * Steps:
 *   1. Parse and load ELF segments
 *   2. Set up page table U/S bits for user code + stack
 *   3. Zero user stack pages
 *   4. Enter ring 3 via IRETQ
 *   5. When program calls exit(), return_to_kernel() brings us back
 * ============================================================================ */
static void cmd_run(void) { cmd_run_elf(hello_elf_data, hello_elf_data_len, "hello"); }
static void cmd_run_elf(const unsigned char *elf_data, unsigned int elf_len, const char *name) {
    kprintf("  Loading ELF64 binary: %s\n", name);

    /* Step 1: Load ELF */
    elf64_load_result_t result = elf64_load(elf_data, elf_len);
    if (!result.valid) {
        kprintf("  ERROR: Failed to load ELF\n");
        return;
    }

    kprintf("  Entry: 0x%x, Load: 0x%x-0x%x\n",
            result.entry, result.load_base, result.load_end);

    /* Step 2: Set up page table U/S bits for user pages */
    kprintf("  Setting up page tables for ring 3...\n");

    /* Make code pages user-accessible.
     * The ELF loads at 0x400000 which is in PD[2] (covers 4MB-6MB). */
    pt_make_user_accessible(result.load_base);

    /* Also make pages around the code user-accessible (in case it spans 2MB boundaries) */
    if (result.load_end > (result.load_base & ~0x1FFFFF) + 0x200000) {
        pt_make_user_accessible(result.load_end);
    }

    /* Make user stack pages user-accessible.
     * Stack at 0x10000000 is in PD[128] (covers 256MB-258MB). */
    pt_make_user_accessible(USER_STACK_BASE);

    /* Flush TLB so page table changes take effect */
    tlb_flush();

    /* Step 3: Zero user stack area */
    kmemset((void *)USER_STACK_BASE, 0, USER_STACK_SIZE);
    kprintf("  User stack: 0x%x - 0x%x (64KB)\n",
            USER_STACK_BASE, USER_STACK_TOP);

    /* Step 4: Update TSS RSP0 for SYSCALL entry from user mode */
    /* (Already set in entry.asm, but update to be safe) */
    extern char tss_struct[];
    *(uint64_t *)(tss_struct + 4) = (uint64_t)&kernel_stack_top;

    kprintf("  Entering userspace at RIP=0x%x RSP=0x%x...\n",
            result.entry, USER_STACK_TOP);
    kprintf("  ----------------------------------------\n");

    /* Step 5: Jump to ring 3!
     * enter_userspace saves kernel RSP and does IRETQ.
     * When the program calls exit(), return_to_kernel() will
     * restore the kernel stack and we'll return here. */
    outb(0x21, 0xFD);  /* Mask timer during userspace */
    enter_userspace(result.entry, USER_STACK_TOP);
    /* DEBUG: serial marker to confirm return */

    /* We get here when return_to_kernel() is called from sys_exit */
    kprintf("\n  ----------------------------------------\n");
    kprintf("  Returned to kernel shell.\n");
    outb(0x21, 0xFC);  /* Re-enable timer */
}

/* === Execute a command === */
static void execute_command(const char *cmd) {
    while (*cmd == ' ') cmd++;
    if (*cmd == '\0') return;

    if (kstrncmp(cmd, "execute ", 8) == 0) {
        cmd_execute(cmd + 8);
    } else if (kstrcmp(cmd, "apps") == 0) {
        cmd_apps();
    } else if (kstrcmp(cmd, "help") == 0) {
        cmd_help();
    } else if (kstrcmp(cmd, "run") == 0) {
        cmd_run();
    } else if (kstrcmp(cmd, "meminfo") == 0) {
        cmd_meminfo();
    } else if (kstrcmp(cmd, "ticks") == 0) {
        cmd_ticks();
    } else if (kstrcmp(cmd, "ata") == 0) {
        cmd_diskinfo();
    } else if (kstrcmp(cmd, "ata_test") == 0) {
        cmd_ata_test();
    } else if (kstrcmp(cmd, "mount") == 0) {
        cmd_mount();
    } else if (kstrcmp(cmd, "reboot") == 0) {
        cmd_reboot();
    } else if (kstrcmp(cmd, "halt") == 0) {
        cmd_halt();
    } else {
        kprintf("  Unknown: '%s'. Type 'apps' for applications.\n", cmd);
    }
}

/* === Initialize shell === */
void shell_init(void) {
    kprintf("\n  ========================================\n");
    kprintf("    Pactor64 Shell Ready\n");
    kprintf("    Type 'help' for available commands\n");
    kprintf("    Type 'run' to execute hello.elf\n");
    kprintf("  ========================================\n\n");
}

/* === Run the shell (main loop) === */
void shell_run(void) {
    char buf[CMD_BUF_SIZE];
    int  pos = 0;

    kprintf("pactor64> ");

    for (;;) {
        /* Poll serial for input */
        int c = serial_getc();
        if (c < 0) {
            /* Also check keyboard (PS/2) */
            c = keyboard_getchar();
            if (c < 0) {
                __asm__ volatile ("hlt");
                continue;
            }
        }

        char ch = (char)c;

        if (ch == '\n' || ch == '\r') {
            /* Enter: execute command */
            serial_puts("\r\n");
            buf[pos] = '\0';
            execute_command(buf);
            pos = 0;
            kprintf("pactor64> ");
        } else if (ch == '\b' || ch == 127) {
            /* Backspace */
            if (pos > 0) {
                pos--;
                serial_puts("\b \b");  /* erase character on terminal */
            }
        } else if (ch >= ' ' && ch < 127) {
            /* Printable character */
            if (pos < CMD_BUF_SIZE - 1) {
                buf[pos++] = ch;
                serial_putc(ch);  /* echo */
            }
        }
        /* Ignore other control characters */
    }
}