#include "task.h"
#include "mmu.h"
#include "pmm.h"
#include "string.h"
#include "uart.h"
#include "switch.h"
#include "exec.h"

task_t tasks[MAX_TASKS];
int task_count = 0;
task_t* current_task[MAX_CPUS] = { 0 };

void task_create(void (*fn)(void)) {
    if (task_count >= MAX_TASKS) return;

    task_t* t = &tasks[task_count];
    t->entry = fn;
    t->state = TASK_READY;
    t->id = task_count;
    t->last_core = 0;

    t->pgdir = pagedir_create();

    uint64_t stack_virt = 0x40000000;
    void* stack_phys = pmm_alloc_page();
    if (!stack_phys) return;

    map_page(t->pgdir, (void*)stack_virt, stack_phys, MMU_FLAG_RW | MMU_FLAG_USER);

    memset(stack_phys, 0, 4096);
    ((uint64_t*)stack_phys)[511] = (uint64_t)fn; // x30/LR

    t->sp = stack_virt + 4096 - 16 * 8;

    task_count++;
}

int task_create_user_elf(uint64_t entry, pagedir_t* pgdir) {
    if (task_count >= MAX_TASKS) return -1;

    task_t* t = &tasks[task_count];
    t->entry = (void (*)(void))entry;
    t->state = TASK_READY;
    t->id = task_count;
    t->last_core = 0;
    t->pgdir = pgdir;

    uint64_t stack_virt = 0x40000000;
    void* stack_phys = pmm_alloc_page();
    if (!stack_phys) return -1;

    map_page(pgdir, (void*)stack_virt, stack_phys, MMU_FLAG_RW | MMU_FLAG_USER);
    memset(stack_phys, 0, 4096);
    ((uint64_t*)stack_phys)[511] = entry;

    t->sp = stack_virt + 4096 - 16 * 8;

    task_count++;
    return t->id;
}

int task_kill(int id) {
    if (id < 0 || id >= task_count) return -1;
    tasks[id].state = TASK_TERMINATED;
    return 0;
}

void scheduler_yield(void) {
    int cid = 0; // mono-core

    task_t* current = current_task[cid];

    for (int i = 1; i <= task_count; i++) {
        int next_id = (current ? current->id : -1) + i;
        if (next_id >= task_count) next_id -= task_count;

        if (tasks[next_id].state == TASK_READY) {
            task_t* next = &tasks[next_id];

            if (current && current->state == TASK_RUNNING)
                current->state = TASK_READY;

            next->state = TASK_RUNNING;
            current_task[cid] = next;

            if (current)
                context_switch(&current->sp, next->sp);
            else
                context_switch(NULL, next->sp);

            while (1) {} // shouldn't return
        }
    }

    uart_puts("[scheduler_yield] No READY task\n");
    while (1);
}
