#include "task.h"
#include "timer.h"
#include "memory.h"
#include "panic.h"
#include "cpu.h"
#include "cpu/irq.h"
#include "cpu/mmu.h"
#include "pit8253.h"
#include "port.h"

#include <stdatomic.h>
#include <stdbool.h>

#include <lib/elf.h>
#include <lib/kstdio.h>

#include "arch_thread.h"

extern uint64_t *get_tss_rsp0(void);

static void lock_scheduler(void);
static void unlock_scheduler(void);
static void lock_preempt(void);
static void unlock_preempt(void);

static void update_time(void);

static struct kc_thread *create_thread(void *(*thread_f)(void *), void *p);
static void destroy_thread(struct kc_thread *thread);
static void set_thread(struct kc_thread *thread);
static void block_thread(enum kc_thread_status reason);
static void unblock_thread(struct kc_thread *thread);
static void sleep_thread(uint64_t nanoseconds);
static void sleep_thread_until(uint64_t nanoseconds);

static int sleeping_thread_callback(uint64_t nanoseconds);

static struct kc_thread *ready_thread_pop(void);
static void ready_thread_push(struct kc_thread *thread);
static void ready_thread_push_back(struct kc_thread *thread);

const struct timer_source * timesource;

/* static threads that are always present
 * TODO: make this a per-CPU thing at some point
 */
static struct kc_thread *idle_thread;

// thread lists for the scheduler to manipulate
static struct kc_thread *current_thread;
static struct kc_thread *first_ready_thread;
static struct kc_thread *last_ready_thread;
static struct kc_thread *sleeping_threads;

static volatile atomic_uint_fast64_t preempt_switch_count = 0;
static volatile atomic_bool preempt_switch_flag = false;

static void *idle_thread_entry(void *)
{
    kprintf("idle thread started\n");

	__asm__ ("sti");

    
    while (true)
    {
        __asm__ ("hlt;");
    }

	return nullptr;
}

#define PS2INPUT 0x60
#define PS2STATUS 0x64
#define PS2STATUS_IBF 0x1

static void keyboard_input_loop()
{
    bool has_bytes = true;
    while (has_bytes)
    {
        has_bytes = inb(PS2STATUS) & PS2STATUS_IBF;
        if (has_bytes)
        {
            kprintf("read byte from keyboard input: %hhd\n", inb(PS2INPUT));
        }
        else
        {
            kprintf("no bytes in keyboard input\n");
        }
    }
}

static void *sleepy_thread_entry(void *)
{
    int thread_slept_count = 0;

    while (true)
    {
        sleep_thread(1000000000);
        thread_slept_count++;
        keyboard_input_loop();
    }

    (void)thread_slept_count;
	return nullptr;
}

static uint64_t scheduler_flags;

noreturn void task_init(void)
{
	struct kc_thread boot_thread;
	boot_thread.status = RUNNING;
	current_thread = &boot_thread;

    timesource = &pit8253_timer_source;
    timesource->append_callback(sleeping_thread_callback);
    timesource->start();
    kprintf(
            "starting task management, timesource delta %luns\n",
            timesource->nanoseconds_delta());

	idle_thread = create_thread(idle_thread_entry, nullptr);
	struct kc_thread *sleepy_thread = create_thread(sleepy_thread_entry, nullptr);
	ready_thread_push(sleepy_thread);
	__asm__ volatile ("movq %%rsp, %0" : "=g"(boot_thread.kernel_stack_pointer) :);
	ready_thread_push(idle_thread);
	block_thread(BLOCKED);
    // shouldn't ever get here
	kprintf("the idle thread exited somehow. we're back in the boot thread. this is not good.\n");
    PANIC(DEAD_END);
}

void task_schedule(void)
{
    if (atomic_load(&preempt_switch_count))
    {
        preempt_switch_flag = true;
        return;
    }

    if (first_ready_thread)
    {
        struct kc_thread *this_thread = ready_thread_pop();

        if (this_thread == idle_thread)
        {
            if (first_ready_thread)
            {
                this_thread->status = READY;
                this_thread = ready_thread_pop();
                ready_thread_push(idle_thread);
            }
            else if (current_thread->status == RUNNING)
            {
                return;
            }
            else
            {
                // NULL statement????? idon't like
            }
        }
        set_thread(this_thread);
    }
}


static void lock_scheduler(void)
{
    scheduler_flags = irq_lock();
}

static void lock_preempt(void)
{
    scheduler_flags = irq_lock();
    preempt_switch_count++;
}

static void unlock_scheduler(void)
{
    irq_unlock(scheduler_flags);
}

static void unlock_preempt(void)
{
    if (atomic_load(&preempt_switch_count) >= 1)
    {
        preempt_switch_count--;
    }

    if (!atomic_load(&preempt_switch_count) && 
            atomic_load(&preempt_switch_flag))
    {
        preempt_switch_flag = false;
        task_schedule();
    }

    irq_unlock(scheduler_flags);
}

void update_time(void)
{
    static uint64_t last_elapsed = 0;
    if (current_thread)
    {
        uint64_t current_elapsed = timesource->nanoseconds_elapsed();
        uint64_t delta = current_elapsed - last_elapsed;
        last_elapsed = current_elapsed;
        current_thread->time_elapsed += delta;
    }
}

#define KC_THREAD_SIZE 0x3000

static struct kc_thread *create_thread(void *(*thread_f)(void *), void *p)
{
	void *buffer = vm_alloc(KC_THREAD_SIZE, VM_ALLOC_ANONYMOUS|VM_ALLOC_ANY); 
	memset(buffer, 0, KC_THREAD_SIZE);

	uintptr_t stack_head = (uintptr_t)buffer + KC_THREAD_SIZE;
	struct kc_thread *thread = (struct kc_thread *)(stack_head -= sizeof(*thread));

	thread->kernel_stack_head = (void *)stack_head;
	thread->kernel_stack_pointer = (void *)stack_head;
	thread->status = READY;

	arch_create_thread(thread, thread_f, p);

    return thread;
}

static void destroy_thread(struct kc_thread *thread)
{
    (void)thread;
}

static void set_thread(struct kc_thread *thread)
{
    if (atomic_load(&preempt_switch_count))
    {
        preempt_switch_flag = true;
        return;
    }
    update_time();
    struct kc_thread *previous_thread = current_thread;
    current_thread = thread;
    
    if (previous_thread->status == RUNNING)
    {
        previous_thread->status = READY;
        ready_thread_push(previous_thread);
    }

    current_thread->status = RUNNING;
	arch_swap_thread(previous_thread, current_thread);

}

static void block_thread(enum kc_thread_status reason)
{
    lock_scheduler();
    current_thread->status = reason;
    task_schedule();
    unlock_scheduler();
}

static void unblock_thread(struct kc_thread *thread)
{
    lock_scheduler();

    thread->status = READY;

    if (!first_ready_thread || (current_thread == idle_thread))
    {
        unlock_preempt();
        set_thread(thread);
    }
    else
    {
        ready_thread_push_back(thread);
    }

    unlock_scheduler();
}

static void sleep_thread(uint64_t nanoseconds)
{
    sleep_thread_until(timesource->nanoseconds_elapsed() + nanoseconds);
}

static void sleep_thread_until(uint64_t nanoseconds)
{
    lock_preempt();

    if (nanoseconds < timesource->nanoseconds_elapsed())
    {
        unlock_scheduler();
        return;
    }

    //kprintf("time elapsed is now %ld\r\n", timesource->nanoseconds_elapsed());
    //kprintf("thread will now sleep until %ld\r\n", nanoseconds);
    current_thread->sleep_expiration = nanoseconds;
    current_thread->next = sleeping_threads;
    sleeping_threads = current_thread;

    unlock_preempt();

    block_thread(SLEEPING);
}

static struct kc_thread *ready_thread_pop(void)
{
    struct kc_thread *thread = first_ready_thread;

    if (thread)
    {
        first_ready_thread = thread->next;
    }

    if (!first_ready_thread)
    {
        last_ready_thread = NULL;
    }

    return thread;
}

static void ready_thread_push(struct kc_thread *thread)
{
    thread->next = first_ready_thread;
    first_ready_thread = thread;

    if (!last_ready_thread)
    {
        last_ready_thread = first_ready_thread;
    }
}

static void ready_thread_push_back(struct kc_thread *thread)
{
    if (last_ready_thread)
    {
        last_ready_thread->next = thread;
    }
    else
    {
        first_ready_thread = thread;
    }
    last_ready_thread = thread;
}

static int sleeping_thread_callback(uint64_t nanoseconds)
{
    lock_preempt();

    struct kc_thread *sleeping = sleeping_threads;
    sleeping_threads = NULL;

    while (sleeping != NULL)
    {
        struct kc_thread *this_thread = sleeping;
        sleeping = sleeping->next;

        if (this_thread->sleep_expiration <= nanoseconds)
        {
            kprintf("thread awakened: %p\n", this_thread);
            this_thread->sleep_expiration = 0;
            unblock_thread(this_thread);
        }
        else
        {
            this_thread->next = sleeping_threads;
            sleeping_threads = this_thread;
        }
    }

	unlock_preempt();
	lock_scheduler();
	task_schedule();
	unlock_scheduler();

    return 0;
}