#include <multiboot2.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include "gdt.h"
#include "idt.h"
#include "pic.h"
#include "port_io.h"
#include "pmm.h"
#include "paging.h"
#include "kmalloc.h"
#include "driver.h"
#include "vga.h"
#include "tss.h"
#include "syscall.h"
#include "process.h"
#include "cpio.h"
#include "vfs.h"
#include "initrd_fs.h"
#include "keyboard.h"
#include "framebuffer.h"

/* Global framebuffer info, parsed from multiboot2 tags. */
framebuffer_info_t fb_info;

/**
 * Initialize COM1 serial port for debug output.
 * Baud rate: 115200, 8N1.
 */
static void serial_init(void) {
    outb(0x3F8 + 1, 0x00);  /* Disable interrupts */
    outb(0x3F8 + 3, 0x80);  /* Enable DLAB */
    outb(0x3F8 + 0, 0x01);  /* Divisor low: 115200 baud */
    outb(0x3F8 + 1, 0x00);  /* Divisor high */
    outb(0x3F8 + 3, 0x03);  /* 8 bits, no parity, 1 stop */
    outb(0x3F8 + 2, 0xC7);  /* Enable FIFO */
    outb(0x3F8 + 4, 0x03);  /* RTS/DSR set */
}

static void serial_putc(char c) {
    /* Wait for transmit buffer empty */
    while (!(inb(0x3F8 + 5) & 0x20));
    outb(0x3F8, c);
}

void offset_print(const char *str)
{
    while (*str) {
        outb(0xE9, *str);   /* debugcon */
        serial_putc(*str);  /* COM1 serial */
        str++;
    }
}

void print_hex(uint32_t val)
{
    const char *hex = "0123456789ABCDEF";
    outb(0xE9, '0'); serial_putc('0');
    outb(0xE9, 'x'); serial_putc('x');
    for (int i = 28; i >= 0; i -= 4) {
        char c = hex[(val >> i) & 0xF];
        outb(0xE9, c);
        serial_putc(c);
    }
    outb(0xE9, '\n'); serial_putc('\n');
}

void kernel_main(uint32_t magic, uint32_t addr) {
    /* Initialize serial port first so all debug output goes to COM1 too */
    serial_init();

    /* Early canary: write directly to VGA text buffer at 0xB8000.
     * If the display is in text mode, this will show a bright magenta 'C'
     * in the top-left corner before any subsystem is initialized. */
    {
        volatile uint16_t *vga = (volatile uint16_t *)0xB8000;
        vga[0] = (uint16_t)'C' | (0x5F << 8); /* magenta on magenta = visible block */
        vga[1] = (uint16_t)'O' | (0x5F << 8);
        vga[2] = (uint16_t)'S' | (0x5F << 8);
    }

    if (magic != MULTIBOOT2_BOOTLOADER_MAGIC) {
        offset_print("Invalid magic number: ");
        print_hex(magic);
        return;
    }
    
    offset_print("Booting...\n");

    init_gdt();
    offset_print("GDT initialized\n");
    
    init_idt();
    offset_print("IDT initialized\n");

    init_pic();
    /* Unmask timer IRQ (IRQ0) explicitly */
    pic_clear_mask(0);
    offset_print("PIC initialized\n");
    
    init_pmm(addr);
    offset_print("PMM initialized\n");

    /* Scan Multiboot2 tags for the initrd module and framebuffer info */
    uint32_t initrd_start = 0, initrd_end = 0;
    memset(&fb_info, 0, sizeof(fb_info));
    {
        struct multiboot_tag *tag;
        for (tag = (struct multiboot_tag *)(addr + 8);
             tag->type != MULTIBOOT_TAG_TYPE_END;
             tag = (struct multiboot_tag *)((uint8_t *)tag + ((tag->size + 7) & ~7u))) {
            if (tag->type == MULTIBOOT_TAG_TYPE_MODULE) {
                struct multiboot_tag_module *mod = (struct multiboot_tag_module *)tag;
                initrd_start = mod->mod_start;
                initrd_end = mod->mod_end;
                offset_print("Initrd module at ");
                print_hex(initrd_start);
                offset_print(" to ");
                print_hex(initrd_end);
            }
            if (tag->type == MULTIBOOT_TAG_TYPE_FRAMEBUFFER) {
                struct multiboot_tag_framebuffer *fbt =
                    (struct multiboot_tag_framebuffer *)tag;
                fb_info.addr   = (uint32_t)fbt->common.framebuffer_addr;
                fb_info.pitch  = fbt->common.framebuffer_pitch;
                fb_info.width  = fbt->common.framebuffer_width;
                fb_info.height = fbt->common.framebuffer_height;
                fb_info.bpp    = fbt->common.framebuffer_bpp;
                fb_info.type   = fbt->common.framebuffer_type;

                if (fb_info.type == FB_TYPE_RGB) {
                    fb_info.red_pos    = fbt->framebuffer_red_field_position;
                    fb_info.red_size   = fbt->framebuffer_red_mask_size;
                    fb_info.green_pos  = fbt->framebuffer_green_field_position;
                    fb_info.green_size = fbt->framebuffer_green_mask_size;
                    fb_info.blue_pos   = fbt->framebuffer_blue_field_position;
                    fb_info.blue_size  = fbt->framebuffer_blue_mask_size;
                }

                offset_print("Framebuffer: type=");
                print_hex(fb_info.type);
                offset_print("  addr=");
                print_hex(fb_info.addr);
                offset_print("  ");
                print_hex(fb_info.width);
                offset_print("  x ");
                print_hex(fb_info.height);
                offset_print("  bpp=");
                print_hex(fb_info.bpp);
            }
        }
    }

    init_paging();
    offset_print("Paging initialized\n");

    /* If GRUB provided a graphical framebuffer, identity-map it so
     * the VGA driver can write pixels to it after paging is enabled. */
    if (fb_info.addr != 0 && fb_info.type == FB_TYPE_RGB) {
        uint32_t fb_size = fb_info.pitch * fb_info.height;
        uint32_t fb_base = fb_info.addr & ~0xFFFu; /* page-align down */
        uint32_t fb_end  = (fb_info.addr + fb_size + 0xFFF) & ~0xFFFu;
        offset_print("  Mapping framebuffer ");
        print_hex(fb_base);
        offset_print("  to ");
        print_hex(fb_end);
        for (uint32_t pa = fb_base; pa < fb_end; pa += 4096) {
            paging_map_page(pa, pa, PAGE_PRESENT | PAGE_WRITE | PAGE_WRITETHROUGH);
        }
    }

    /* Test paging: allocate a page and write to it */
    void *test_page = paging_alloc_page();
    if (test_page) {
        offset_print("Allocated virtual page at: ");
        print_hex((uint32_t)test_page);
        *((volatile uint32_t *)test_page) = 0xDEADBEEF;
        offset_print("Virtual page write/read OK\n");
        paging_free_page(test_page);
    } else {
        offset_print("FAILED to allocate virtual page\n");
    }

    init_kmalloc();
    offset_print("Memory allocator initialized\n");

    /* Initialize CPIO ramdisk if module was loaded */
    if (initrd_start != 0) {
        cpio_init((const void *)initrd_start, initrd_end - initrd_start);
        offset_print("CPIO ramdisk initialized\n");
    } else {
        offset_print("No initrd module found\n");
    }

    init_vfs();
    offset_print("VFS initialized\n");

    if (initrd_start != 0) {
        init_initrd_fs();
        offset_print("Initrd filesystem mounted\n");

        /* Test VFS: read a file from the initrd */
        int fd = vfs_open("/initrd/README", 0);
        if (fd >= 0) {
            char buf[64];
            int32_t n = vfs_read(fd, buf, sizeof(buf) - 1);
            if (n > 0) {
                buf[n] = '\0';
                offset_print("VFS read /initrd/README: ");
                offset_print(buf);
            }
            vfs_close(fd);
        }
    }

    init_tss();
    offset_print("TSS initialized\n");

    init_syscalls();
    offset_print("Syscalls initialized\n");

    keyboard_init();
    offset_print("Keyboard initialized\n");

    init_process();
    offset_print("Process subsystem initialized\n");

    init_drivers();
    offset_print("Drivers initialized\n");

    /* Show memory statistics and boot progress on VGA */
    vga_show_mem_stats();
    vga_puts("Boot complete.\n\n");

    /* Test kmalloc/kfree */
    uint32_t *test_alloc = (uint32_t *)kmalloc(64);
    if (test_alloc) {
        *test_alloc = 42;
        offset_print("kmalloc(64) = ");
        print_hex((uint32_t)test_alloc);
        kfree(test_alloc);
        offset_print("kfree OK\n");
    } else {
        offset_print("FAILED to kmalloc\n");
    }

    /* Load the initial program from the initrd and run it */
    cpio_entry_t app_entry;
    const char *init_app = "sh";
    if (cpio_find(init_app, &app_entry) == 0) {
        offset_print("Found ");
        offset_print(init_app);
        offset_print(" in initrd (");
        print_hex(app_entry.datasize);
        offset_print(" bytes)\n");

        int32_t pid = process_create(init_app,
                                      app_entry.data,
                                      app_entry.datasize);
        if (pid > 0) {
            offset_print("Created init process, pid=");
            print_hex((uint32_t)pid);

            /* Enable interrupts before entering user mode */
            asm volatile("sti");
            offset_print("Interrupts enabled\n");

            /* Enter user mode - does not return */
            process_run_first();
        } else {
            offset_print("FAILED to create init process\n");
        }
    } else {
        offset_print(init_app);
        offset_print(" not found in initrd\n");
    }

    /* Enable interrupts */
    asm volatile("sti");
    offset_print("Interrupts enabled\n");

    offset_print("Hello, world\n");
}