#include "kernel/kernel.h"
#include <kernel/task.h>
#include <kernel/mem.h>

extern uint32_t boot_page_tab[1024];
extern kernel_ctx_t ctx;

//int task_init(task_t *task, uint32_t *img) {
int task_init(task_t *task, elf_t *elf) {

    if (!elf) return 0;
    uint8_t *s = (uint8_t *)mmap_find_first_free_block();
    /*
       memset(s, 0, 0x1000);
       int idx = 0;
    // Loop section header to find ALLOC sections then allocate them
    for (int i = 0; i < elf->e_shnum; i++) {
    section_t *sec = (section_t *)((uint32_t)elf + elf->e_shoff + i*elf->e_shentsize);
    if (!(sec->sh_flags & 0x2)) continue;
    if (!sec->sh_addr) continue;
    idx += sec->sh_offset;
    memcpy(s + idx, (void *)((uint32_t)elf + sec->sh_offset), sec->sh_offset);
    LOG("%d\n", sec->sh_offset);
    }
    */
    memcpy(s, elf, 0x1000);

    //section_t *text = elf_find_section(elf, ".text");
    //if (!text) return 0;
    task->eip = elf->e_entry;
    // Allocate a page for a stack
    task->esp = (uint32_t)mmap_find_first_free_block() + 0x1000;
    // Allocate a page for stdin
    // LATER: Allow for allocating more on demand
    task->stdin = (uint8_t *)mmap_find_first_free_block();
    task->stdin_len = 0;

    uint32_t *page_dir = mmap_find_first_free_block();
    task->cr3 = (uint32_t)page_dir - KERNEL_VMA;
    uint32_t *page_tab = mmap_find_first_free_block();
    for (int i = 0; i < 1024; i++) {
        // Map kernel to 3GiB
        if (i == 768) page_dir[i] = ((uint32_t)boot_page_tab - KERNEL_VMA) | 7;
        // Map 3.4GiB-4GiB as non-present kernel entries
        else if (i > 768) page_dir[i] = 4;
        // Map remaining entries as non-present u/s
        else page_dir[i] = 6;
    }
    // Map the ELF image to its specified entry
    // This will cause problems if the ELF image is larger than 4 MiB so I need to fix this
    //page_dir[PAGE_DIR_INDEX(elf->e_entry)] = ((uint32_t)page_tab - KERNEL_VMA) | 7;
    page_dir[PAGE_DIR_INDEX(elf->e_entry)] = ((uint32_t)page_tab - KERNEL_VMA) | 7;

    /*
       uint32_t elf_phys = (uint32_t)s - KERNEL_VMA;
       for (int i = 0; i < 1024; i++) {
       page_tab[i] = (elf_phys + 0x1000*i) | 7;
       }
       */
    page_tab[PAGE_TAB_INDEX(elf->e_entry)] = ((uint32_t)s - KERNEL_VMA) | 7;

    return 1;
}

void task_sort() {
    if (!ctx.current_task || !ctx.current_task->next) return; // If the list is empty or has only one node

    task_t* new_head = (task_t *)ctx.current_task->next;
    task_t* temp = new_head;

    while (temp->next) {
        temp = (task_t *)temp->next;
    }

    temp->next = (uint32_t *)ctx.current_task;
    ctx.current_task->next = 0;
    ctx.current_task = new_head;
}

void schedule(void) {
    // If there's only one task running
    if (!ctx.current_task->next) return;
}