#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"

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);

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;

struct thread_queue
{
	struct kc_thread
		*first,
		*last;
};

static void thread_queue_push(struct thread_queue *queue, struct kc_thread *thread);
static void thread_queue_push_back(struct thread_queue *queue, struct kc_thread *thread);

static void thread_queue_push(struct thread_queue *queue, struct kc_thread *thread)
{
	thread->next = queue->first;
	queue->first = thread;

	if (!queue->last) 
	{
		thread_queue_push_back(queue, thread);
	}
}

static void thread_queue_push_back(struct thread_queue *queue, struct kc_thread *thread)
{
	if (!queue->first)
	{
		thread_queue_push(queue, thread);
	}
	else if (queue->last)
	{
		queue->last->next = thread;
	}

	queue->last = thread;
}

struct kc_thread *thread_queue_pop(struct thread_queue *queue)
{
	if (!queue->first)
	{
		return nullptr;
	}

	struct kc_thread *thread = queue->first;
	queue->first = thread->next;
	thread->next = nullptr;
	if (!queue->first)
	{
		queue->last = nullptr;
	}

	return thread;
}

bool thread_queue_any(struct thread_queue queue)
{
	return queue.first && queue.last;
}

static struct thread_queue ready_queue;
static struct thread_queue sleep_queue;

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;
}

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

    while (true)
    {
        sleep_thread(sleep_nanoseconds);
        thread_slept_count++;
		kprintf("thread slept %d times\n", thread_slept_count);
    }

    (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, (void *)3000000000);
	struct kc_thread *sleepy_thread2 = create_thread(sleepy_thread_entry, (void *)5000000000);
	thread_queue_push(&ready_queue, sleepy_thread);
	thread_queue_push(&ready_queue, sleepy_thread2);
	__asm__ volatile ("movq %%rsp, %0" : "=g"(boot_thread.kernel_stack_pointer) :);
	thread_queue_push(&ready_queue, 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);
}

static inline bool check_preempt_count()
{
	if (atomic_load(&preempt_switch_count) > 0)
	{
		return true;
	}

	return false;
}

void task_schedule(void)
{
	if(check_preempt_count())
	{
		preempt_switch_flag = true;
		return;
	}

	struct kc_thread *t = thread_queue_pop(&ready_queue);

	if (t == idle_thread)
	{
		thread_queue_push_back(&ready_queue, t);
		t = thread_queue_pop(&ready_queue);
	}

	if (t && current_thread->status != RUNNING)
	{
		set_thread(t);
	}
}

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) > 0)
    {
        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 (thread == current_thread)
	{
		return;
	}

    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;
        thread_queue_push(&ready_queue, 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 (!thread_queue_any(ready_queue) || (current_thread == idle_thread))
    {
        unlock_preempt();
        set_thread(thread);
    }
    else
    {
        thread_queue_push_back(&ready_queue, 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;
    }

    current_thread->sleep_expiration = nanoseconds;
	
	thread_queue_push_back(&sleep_queue, current_thread);

    unlock_preempt();

    block_thread(SLEEPING);
}

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

	struct thread_queue queue = sleep_queue;

	sleep_queue = (struct thread_queue){ nullptr, nullptr };

	{
		struct kc_thread *t;
		while((t = thread_queue_pop(&queue)) != nullptr)
		{
			if (t->sleep_expiration <= timestamp)
			{
				t->sleep_expiration = 0;
				unblock_thread(t);
			}
			else
			{
				thread_queue_push(&sleep_queue, t);
			}
		}
	}

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

    return 0;
}

typedef void (locked_scheduler_action)(void *);

static void with_locked_scheduler(
		locked_scheduler_action *action,
		void *action_parameter)
{
	lock_scheduler();
	action(action_parameter);
	unlock_scheduler();
}

static void terminate_action(void *result)
{
	union kc_thread_result *result_value = result;
	current_thread->result = *result_value;
	current_thread->kernel_stack_pointer = current_thread->kernel_stack_head;
	// TODO: wake threads waiting
	// TODO: write waitng thread code
	block_thread(TERMINATED);
}

void kct_terminate(union kc_thread_result result)
{
	with_locked_scheduler(terminate_action, &result);
}