#include <efi.h>
#include <graphics/context.h>
#include <graphics/draw.h>
#include <fonts/pixel_font.h>
#include <serial.h>
#include <common.h>
#include <memory/pmm.h>
#include <memory/heap.h>
#include <scheduler.h>
#include <fs.h>
#include <gdt.h>
#include <idt.h>
#include <pic.h>
#include <pit.h>

extern EFI_SYSTEM_TABLE *ST;

inline void init_gop() {
    EFI_GUID gop_guid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    EFI_GRAPHICS_OUTPUT_PROTOCOL *GOP;
    uefi_call_wrapper(ST->BootServices->SetWatchdogTimer, 4, 0, 0, 0, NULL);
    uefi_call_wrapper(ST->BootServices->LocateProtocol, 3, &gop_guid, NULL, (void **)&GOP);
    gfx_init(GOP);
}

inline void init_serial() {
    EFI_SERIAL_IO_PROTOCOL *SerialIo = NULL;
    EFI_GUID gEfiSerialIoProtocolGuid = EFI_SERIAL_IO_PROTOCOL_GUID;
    EFI_HANDLE *SerialHandles = NULL;
    UINTN NumSerials = 0;
    EFI_STATUS status = uefi_call_wrapper(ST->BootServices->LocateHandleBuffer, 5,
        ByProtocol,
        &gEfiSerialIoProtocolGuid,
        NULL,
        &NumSerials,
        &SerialHandles
    );
    if (status == EFI_SUCCESS && NumSerials > 0) {
        status = uefi_call_wrapper(ST->BootServices->HandleProtocol, 3,
            SerialHandles[0], &gEfiSerialIoProtocolGuid, (void **)&SerialIo);
        if (status == EFI_SUCCESS) {
            serial_init(SerialIo);
        }
    }
    if (SerialHandles) {
        ST->BootServices->FreePool(SerialHandles);
    }
}

// External: PIT IRQ handler defined in pit.cpp
extern "C" void pit_irq_handler(void);

// PIC IRQ handler — dispatches IRQ 0 (timer)
static void irq_handler(trap_frame* frame) {
    UINT8 vector = (UINT8)frame->vector;
    UINT8 irq = vector - PIC_IRQ_BASE;

    // Send EOI BEFORE handling, so PIC can deliver new interrupts
    // immediately after a context switch inside the handler.
    pic_send_eoi(irq);

    switch (irq) {
        case 0:  // PIT timer
            pit_irq_handler();
            break;
        default:
            break;
    }
}

extern "C" void kernel_main() {
    init_gop();
    init_serial();

    uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, 5);
    uefi_call_wrapper(ST->ConOut->OutputString, 2, ST->ConOut, (CHAR16*)L"Kernel is running!\n");

    uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
    serial_write("\n\n");

    // init memory managers
    serial_write("Sylva: init PMM...\n");
    EFI_STATUS st = pmm_init();
    if (EFI_ERROR(st)) {
        serial_write("Sylva: PMM init FAILED!\n");
    } else {
        serial_write("Sylva: PMM init OK\n");
        serial_write("Sylva: free pages = ");
        serial_write_hex(pmm_get_free_count());
        serial_write("\n");
    }

    serial_write("Sylva: init heap...\n");
    init_heap();

    // test kmalloc/kfree
    serial_write("Sylva: kmalloc test...\n");
    void* p1 = kmalloc(64);
    void* p2 = kmalloc(128);
    void* p3 = kmalloc(256);
    serial_write("Sylva: p1 = ");
    serial_write_hex((UINTN)p1);
    serial_write("  p2 = ");
    serial_write_hex((UINTN)p2);
    serial_write("  p3 = ");
    serial_write_hex((UINTN)p3);
    serial_write("\n");

    serial_write("Sylva: kfree test...\n");
    kfree(p2);
    kfree(p1);
    kfree(p3);

    void* p4 = kmalloc(32);
    serial_write("Sylva: realloc p4 = ");
    serial_write_hex((UINTN)p4);
    serial_write("\n");
    kfree(p4);

    serial_write("Sylva: memory init done.\n");

    serial_write("Sylva: FS init...\n");
    EFI_STATUS fs_st = fs_init();
    if (EFI_ERROR(fs_st)) {
        serial_write("Sylva: FS init FAILED!\n");
    } else {
        serial_write("Sylva: root directory listing:\n");
        fs_list();
    }

    pf_print("Welcome to Sylva OS!\n");
    serial_write(" Kernel prepared well.\n");

    // --- Interrupt infrastructure ---
    serial_write("Sylva: init GDT...\n");
    gdt_init();

    serial_write("Sylva: init IDT...\n");
    idt_init();

    serial_write("Sylva: init PIC...\n");
    pic_init();

    // Register IRQ handler (vector 0x20 = PIC_IRQ_BASE + 0)
    idt_set_handler(PIC_IRQ_BASE + 0, irq_handler);

    serial_write("Sylva: init PIT...\n");
    pit_init();
    pit_set_tick_handler(scheduler_tick);

    // Enable interrupts
    ASM("sti");
    serial_write("Sylva: interrupts enabled\n");

    // --- Multitasking demo ---
    serial_write("Sylva: creating tasks...\n");

    // Task A: prints messages — preemption handles time slicing
    task_create("taskA", []() {
        for (int i = 0; i < 10; i++) {
            serial_write("[taskA] running iteration ");
            serial_write_hex(i);
            serial_write("\n");
            for (volatile int j = 0; j < 50000000; j++) {}
        }
        serial_write("[taskA] done\n");
    });

    // Task B: prints messages
    task_create("taskB", []() {
        for (int i = 0; i < 5; i++) {
            serial_write("[taskB] hello from taskB #");
            serial_write_hex(i);
            serial_write("\n");
            for (volatile int j = 0; j < 50000000; j++) {}
        }
        serial_write("[taskB] done\n");
    });

    // Task C: short task
    task_create("taskC", []() {
        serial_write("[taskC] quick task\n");
        for (volatile int j = 0; j < 50000000; j++) {}
        serial_write("[taskC] finished\n");
    });

    serial_write("Sylva: starting preemptive scheduler\n");
    scheduler_run();  // never returns
}