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

#include <attributes.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <myke/cpu/cpu.h>
#include <myke/libk/libk.h>
#include <myke/mem/pmm.h>
#include <myke/libk/syscall.h>
#include <myke/tasks/task.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)

const char *task_state_str(uint8_t state) {
    switch (state) {
        case TASK_STATE_RUNNABLE:
            return "RUNNABLE   ";
        case TASK_STATE_RUNNING:
            return "RUNNING    ";
        case TASK_STATE_WAIT_IRQ:
            return "WAIT_IRQ   ";
        case TASK_STATE_WAIT_SIGNAL:
            return "WAIT_SIGNAL";
        case TASK_STATE_STOPPED:
            return "STOPPED    ";
        case TASK_STATE_ERROR:
            return "ERROR      ";
        case TASK_STATE_UNKNOWN:
        default:
            return "UNKNOWN    ";
    }
}

int errno = 0;

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

    // identity
    pid_t tid;
    char *name;

    // linked list
    struct task *next;

    // persistence/memory
    int errno;
    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];
} att_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;

pid_t last_tid = 0;

pid_t reaper_pid = 0;

extern void switch_task(task_registers_t **, task_registers_t *);

extern att_cdecl att_noreturn void __task_entry_point();

extern att_noreturn void __task_entry_point_inner();

// internal api
void task_free(task_t *task);

// explicit cdecl calling convention
void att_unused att_cdecl att_noreturn task_entry_point(task_entrypoint entrypoint, void *entry_data) {
    entrypoint(entry_data);
    syscall_kill_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");
    }
}

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

void task_signal(pid_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();
}

// internal tasks
void att_noreturn task_idle(void *data) {
    while (true) __asm__("hlt");
}

void att_noreturn task_reaper(void *data) {
    while (true) {
        bool did_reap = false;
        task_lock_acquire();
        task_t *t = first_task;
        task_t *p = first_task;
        task_t *n = NULL;
        while (t != NULL) {
            n = t->next;
            if (t->state == TASK_STATE_STOPPED) {
                did_reap = true;
                if (t == first_task) {
                    if (n == NULL) {
                        k_panics("No more tasks");
                    }
                    first_task = n;
                    task_free(t);
                    t = NULL;
                } else {
                    p->next = n;
                    task_free(t);
                    t = p;
                }
            }
            p = t;
            t = n;
        }
        task_lock_free();
        if (did_reap) {
            syscall_yield_job();
        } else {
            syscall_job_suspend();
        }
    }
}

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 != NULL) {
        previous_task->errno = errno;
        if (previous_task->state == TASK_STATE_RUNNING) {
            previous_task->state = TASK_STATE_RUNNABLE;
        }
    }
    current_task->state = TASK_STATE_RUNNING;
    errno = current_task->errno;
    // 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, char *name) {
    task_t *new_task = malloc(sizeof(task_t));
    memset((uint8_t *) new_task, 0, sizeof(task_t));
    new_task->tid = last_tid++;
    new_task->name = strdup(name);
    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_free(task_t *task) {
    pmm_free_pages(task->stack, 16);
    free(task);
}

void task_init() {
    idle_task = task_create(task_idle, NULL, "idle");
    reaper_pid = task_spawn(task_reaper, NULL, "reaper");
}

pid_t task_spawn(task_entrypoint entrypoint, void *entry_data, char *name) {
    task_t *new_task = task_create(entrypoint, entry_data, name);
    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(pid_t tid) {
    task_t *t = first_task;
    while (t != NULL) {
        if (t->tid == tid) {
            t->state = TASK_STATE_STOPPED;
            task_signal(reaper_pid);
            break;
        }
        t = t->next;
    }
}

void task_print_all() {
    // acquiring task lock as reference to current task may disappear
    // might overrun kprint buffer
    task_lock_acquire();
    printf("tasks:\n");
    task_t *c = first_task;
    while (c != NULL) {
        printf("%d - %s %s\n", c->tid, task_state_str(c->state), c->name);
        c = c->next;
    }
    task_lock_free();
}