#include <scheduler.h>
#include <idt.h>
#include <pic.h>
#include <memory/heap.h>
#include <memory/pmm.h>
#include <common.h>
#include <serial.h>

// Assembly: context_switch(UINT64* old_rsp, UINT64 new_rsp)
extern "C" void context_switch(UINT64* old_rsp, UINT64 new_rsp);

static task_t  g_tasks[TASK_MAX];
static UINT32  g_task_count = 0;
static task_t* g_current = NULL;
static task_t* g_task_list = NULL;  // circular linked list head

// Trampoline: first thing a new task runs after context_switch.
static void (*g_task_entries[TASK_MAX])(void);

extern "C" void task_entry_trampoline() {
    ASM("sti");
    task_t* cur = scheduler_current();
    if (cur && g_task_entries[cur->id]) {
        g_task_entries[cur->id]();
    }
    task_exit();
}

task_t* task_create(const char* name, void (*entry)(void)) {
    if (g_task_count >= TASK_MAX) {
        serial_write("SCHED: task limit reached\n");
        return NULL;
    }

    UINT32 id = g_task_count++;
    task_t* task = &g_tasks[id];

    g_task_entries[id] = entry;

    // Allocate kernel stack
    UINTN stack_pages = TASK_STACK_SIZE / PAGE_SIZE;
    void* stack = pmm_alloc_pages(stack_pages);
    if (!stack) {
        serial_write("SCHED: stack alloc failed for task ");
        serial_write(name);
        serial_write("\n");
        return NULL;
    }

    task->id = id;
    task->state = TASK_STATE_READY;
    task->stack_base = stack;
    task->time_slice = TIME_SLICE_DEFAULT;

    // Copy name
    const char* s = name;
    char* d = task->name;
    for (int i = 0; i < TASK_NAME_LEN - 1 && *s; i++) {
        *d++ = *s++;
    }
    *d = '\0';

    // Set up initial stack for first context_switch into this task.
    // Stack grows downward. context_switch will pop 6 regs then ret.
    //
    // Layout (high addr -> low addr):
    //   [stack + TASK_STACK_SIZE]     <- top
    //   return addr = task_entry_trampoline  (ret goes here)
    //   rbx = 0
    //   rbp = 0
    //   r12 = 0
    //   r13 = 0
    //   r14 = 0
    //   r15 = 0                       <- RSP points here initially
    //
    // When preempted by timer IRQ, the ISR stub saves a full trap_frame
    // on the task's stack — that layout is only created by hardware+ISR.
    //
    UINT64* sp = (UINT64*)((UINT8*)stack + TASK_STACK_SIZE);

    *--sp = (UINT64)task_entry_trampoline;

    *--sp = 0;  // rbx
    *--sp = 0;  // rbp
    *--sp = 0;  // r12
    *--sp = 0;  // r13
    *--sp = 0;  // r14
    *--sp = 0;  // r15

    task->rsp = (UINT64)sp;

    // Insert into circular linked list
    if (g_task_list == NULL) {
        task->next = task;
        g_task_list = task;
    } else {
        task->next = g_task_list->next;
        g_task_list->next = task;
        g_task_list = task;
    }

    serial_write("SCHED: created task '");
    serial_write(task->name);
    serial_write("' id=");
    serial_write_hex(id);
    serial_write("\n");

    return task;
}

// Find next READY task in the circular list, starting from g_current->next
static task_t* find_next_ready(void) {
    if (g_current == NULL || g_task_list == NULL) return NULL;

    task_t* next = g_current->next;
    task_t* start = next;

    do {
        if (next->state == TASK_STATE_READY) {
            return next;
        }
        next = next->next;
    } while (next != start);

    return NULL;  // no READY tasks
}

void yield(void) {
    if (g_current == NULL || g_task_list == NULL) return;

    task_t* cur = g_current;
    task_t* next = find_next_ready();

    if (next == NULL || next == cur) return;

    if (cur->state != TASK_STATE_TERMINATED)
        cur->state = TASK_STATE_READY;
    next->state = TASK_STATE_RUNNING;
    next->time_slice = TIME_SLICE_DEFAULT;
    g_current = next;

    context_switch(&cur->rsp, next->rsp);
}

// Timer tick handler — called from PIT IRQ 0
void scheduler_tick(void) {
    if (g_current == NULL) return;

    // Decrement time slice
    if (g_current->time_slice > 0) {
        g_current->time_slice--;
    }

    // If time slice expired, preempt
    if (g_current->time_slice == 0) {
        task_t* cur = g_current;
        task_t* next = find_next_ready();

        if (next == NULL || next == cur) {
            // No other task ready, or only this task — reload time slice
            cur->time_slice = TIME_SLICE_DEFAULT;
            return;
        }

        // Preempt
        cur->state = TASK_STATE_READY;
        cur->time_slice = TIME_SLICE_DEFAULT;
        next->state = TASK_STATE_RUNNING;
        next->time_slice = TIME_SLICE_DEFAULT;
        g_current = next;

        context_switch(&cur->rsp, next->rsp);
    }
}

void scheduler_run(void) {
    if (g_task_list == NULL) {
        serial_write("SCHED: no tasks to run\n");
        return;
    }

    // Find first READY task
    task_t* start = g_task_list->next;
    task_t* t = start;
    do {
        if (t->state == TASK_STATE_READY) {
            break;
        }
        t = t->next;
    } while (t != start);

    if (t->state != TASK_STATE_READY) {
        serial_write("SCHED: no READY tasks\n");
        return;
    }

    g_current = t;
    t->state = TASK_STATE_RUNNING;
    t->time_slice = TIME_SLICE_DEFAULT;

    serial_write("SCHED: starting task '");
    serial_write(t->name);
    serial_write("'\n");

    // First context switch — switch to the task's stack
    // This will never return (until all tasks terminate)
    UINT64 dummy_rsp;
    context_switch(&dummy_rsp, t->rsp);

    // We only return here when ALL tasks are terminated
    serial_write("SCHED: all tasks finished\n");
    while (1) ASM ("hlt");
}

void task_exit(void) {
    if (g_current == NULL) return;

    serial_write("SCHED: task '");
    serial_write(g_current->name);
    serial_write("' exited\n");

    g_current->state = TASK_STATE_TERMINATED;

    // Yield to next task — we won't come back
    yield();

    // Should never reach here
    while (1) ASM ("hlt");
}

task_t* scheduler_current(void) {
    return g_current;
}