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

// 汇编函数：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;  // 循环链表头

// 跳板函数：新任务在 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;

    // 分配内核栈
    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;

    // 复制任务名称
    str_copy(task->name, name, TASK_NAME_LEN);

    // 设置首次 context_switch 时的初始栈
    // 栈向下增长。context_switch 会弹出 6 个寄存器然后 ret。
    //
    // 布局（高地址 → 低地址）：
    //   [stack + TASK_STACK_SIZE]     <- 栈顶
    //   返回地址 = task_entry_trampoline（ret 跳转到这里）
    //   rbx = 0
    //   rbp = 0
    //   r12 = 0
    //   r13 = 0
    //   r14 = 0
    //   r15 = 0                       <- RSP 初始指向这里
    //
    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;

    // 插入循环链表
    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;
}

// 在循环链表中查找下一个就绪任务
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;  // 没有就绪任务
}

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);
}

// 定时器 tick 处理 — 由 PIT IRQ 0 调用
void scheduler_tick(void) {
    if (g_current == NULL) return;

    // 递减时间片
    if (g_current->time_slice > 0) {
        g_current->time_slice--;
    }

    // 时间片用完则抢占
    if (g_current->time_slice == 0) {
        task_t* cur = g_current;
        task_t* next = find_next_ready();

        if (next == NULL || next == cur) {
            // 没有其他就绪任务，或仅此一个 — 重置时间片
            cur->time_slice = TIME_SLICE_DEFAULT;
            return;
        }

        // 抢占
        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;
    }

    // 查找第一个就绪任务
    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");

    // 首次上下文切换 — 切换到任务栈
    // 此后不会返回（直到所有任务终止）
    UINT64 dummy_rsp;
    context_switch(&dummy_rsp, t->rsp);

    // 只有所有任务终止后才会返回到这里
    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;

    // 让出 CPU 给下一个任务 — 不会回来
    yield();

    // 不应到达此处
    while (1) ASM ("hlt");
}

task_t* scheduler_current(void) {
    return g_current;
}