//
// Created by rick on 22-02-21.
//

#include "task.h"
#include <cpu/cpu.h>
#include <stdbool.h>
#include <libc/libc.h>
#include <mem/malloc.h>
#include <mem/pmm.h>

#include <attributes.h>
#include <libk/libk.h>

#define stack_end(task) ((task)->stack + ((task)->stack_page_count * PAGE_SIZE))

#define TASK_STATE_UNKNOWN      (0)
#define TASK_STATE_RUNNABLE     (1 << 0)
#define TASK_STATE_RUNNING      (1 << 1)
#define TASK_STATE_WAIT_IRQ     (1 << 2)
#define TASK_STATE_WAIT_SIGNAL  (1 << 3)

#define TASK_STATE_STOPPED      (1 << 6)
#define TASK_STATE_ERROR        (1 << 7)


typedef struct task {
    bool present: 1;
    uint8_t state;
    uint16_t wait_irq;

    uint32_t tid;

    struct task *next;

    void *stack;
    uint32_t stack_page_count;

    task_registers_t *task_registers;
} task_t;

typedef struct {
    task_registers_t task_registers;
    isr_registers_t isr_registers;
    task_entrypoint entrypoint;
    void *entry_data;
    char alignment[4];
} packed task_stack_start;

volatile uint32_t task_locked = 0;

task_t *idle_task = NULL;
task_t *first_task = NULL;
task_t *last_task = NULL;
task_t *current_task = NULL;

uint32_t last_tid = 0;

extern switch_task(task_registers_t
**, task_registers_t*);

extern cdecl noreturn void __task_entry_point();

extern noreturn void __task_entry_point_inner();

// explicit cdecl calling convention
void cdecl noreturn task_entry_point(task_entrypoint entrypoint, void *entry_data) {
    entrypoint(entry_data);
//    task_end_self();
    while (true); // halt
}

void task_ensure_enabled() {
    if (current_task == NULL) {
        k_panics("Tasking should be enabled\n");
    }
}

void task_lock_acquire() {
    if (__sync_add_and_fetch(&task_locked, 1) == UINT32_MAX) {
        k_panics("To many task locks");
    };
}

void task_lock_free() {
    if (__sync_sub_and_fetch(&task_locked,  1) == UINT32_MAX) {
        k_panics("To many task free's");
    }
}

uint32_t task_get_current_tid() {
    return current_task->tid;
}

void task_signal(uint32_t tid) {
    task_t *t = first_task;
    while (t != NULL) {
        if (t->tid == tid) {
            if (t->state != TASK_STATE_WAIT_SIGNAL) {
                // todo
            }
            t->state = TASK_STATE_RUNNABLE;
            break;
        }
        t = t->next;
    }
}

void task_suspend() {
    current_task->state = TASK_STATE_WAIT_SIGNAL;
    task_switch_next();
}

void noreturn task_idle(void *data) {
    while (true) __asm__("hlt");
}

void task_notify_irq(uint8_t irq_no) {
    uint16_t irq__bit = 1 << irq_no;
    task_t *t = first_task;
    while (t != NULL) {
        if (t->state == TASK_STATE_WAIT_IRQ && (t->wait_irq & irq__bit) != 0) {
            t->wait_irq = 0;
            t->state = TASK_STATE_RUNNABLE;
        }
        t = t->next;
    }
}

void task_wait_irq(uint16_t irq_bits) {
    current_task->wait_irq = irq_bits;
    current_task->state = TASK_STATE_WAIT_IRQ;
    task_switch_next();
}

void task_switch_next_inner(task_t *next_task) {
    if (task_locked != 0) {
        return; // don't switch while the task is locked
    }
    if (next_task == current_task) {
        return;
    }
    task_t *previous_task = current_task;
    current_task = next_task;
    if (previous_task->state == TASK_STATE_RUNNING) {
        previous_task->state = TASK_STATE_RUNNABLE;
    }
    current_task->state = TASK_STATE_RUNNING;
    // switch task
    switch_task(previous_task == NULL ? NULL : &previous_task->task_registers, current_task->task_registers);
}

void task_start_first() {
    if (task_locked > 0) {
        k_panics("Tasking locked before start\n");
    }
    task_switch_next_inner(first_task);
}

task_t *task_first_runnable() {
    task_t *next_task = current_task;
    bool looped = false;
    if (next_task == idle_task || next_task->next == NULL) {
        looped = true;
        next_task = first_task;
    } else {
        next_task = current_task->next;
    }
    do {
        if (!looped && next_task == NULL) {
            looped = true;
            next_task = first_task;
        }
        if (next_task == NULL) {
            break;
        }
        if (next_task->state == TASK_STATE_RUNNABLE) {
            return next_task;
        }
        next_task = next_task->next;
    } while (next_task != current_task);
    if (current_task->state == TASK_STATE_RUNNING)
        return current_task;
    return idle_task;
}

void task_switch_next() {
    if (current_task == NULL) {
        return;
    }
    task_switch_next_inner(task_first_runnable());
}

task_t *task_create(task_entrypoint entrypoint, void *entry_data) {
    task_t *new_task = malloc(sizeof(task_t));
    memset((uint8_t *) new_task, 0, sizeof(task_t));
    new_task->tid = last_tid++;
    new_task->state = TASK_STATE_RUNNABLE;

    new_task->stack = pmm_get_pages(16);
    new_task->stack_page_count = 16;
    // todo check for null

    // prepare stack
    // cdecl format
    task_stack_start *stack = (task_stack_start *) (stack_end(new_task) - sizeof(task_stack_start));
    memset((uint8_t *) stack, 0, sizeof(task_stack_start));
    // setup isr return frame
    stack->isr_registers.cs = 0x08; // todo proper gdt setup
    stack->isr_registers.ds = 0x10; // todo
//    stack->isr_registers.ss = 0x10; // todo

    stack->isr_registers.eip = (uint32_t) __task_entry_point_inner;
    stack->isr_registers.eflags = 0x200;
    // setup task switch frame
    stack->task_registers.eip = (uint32_t) __task_entry_point;
    stack->task_registers.ebp = (uint32_t) stack_end(new_task);
    // setup entrypoint
    stack->entrypoint = entrypoint;
    stack->entry_data = entry_data;

    new_task->task_registers = &stack->task_registers;

    new_task->present = true;
    return new_task;

}

void task_init() {
    idle_task = task_create(task_idle, NULL);
}

uint32_t task_spawn(task_entrypoint entrypoint, void *entry_data) {
    task_t *new_task = task_create(entrypoint, entry_data);
    if (first_task == NULL) {
        // first task
        first_task = new_task;
    } else {
        last_task->next = new_task;
    }
    last_task = new_task;
    return new_task->tid;
}

void task_end(uint32_t tid) {
    task_t *t = first_task;
    while (t != NULL) {
        if (t->tid == tid) {
            t->state = TASK_STATE_STOPPED;
            break;
        }
        t = t->next;
    }
}