#include <thread.h>
#include <logger.h>
#include <utils.h>
#include <kmalloc.h>
#include <mman.h>
#include <queue.h>
#include <interrupt.h>
#include <string.h>
#include <timer.h>

queue_t *global_run_queue = 0x0;
queue_t *global_wait_queue = 0x0;
queue_t *global_gc_queue = 0x0;

thread_t *get_current(void)
{
  uint64_t tpidr_el1;
  R_SYSREG(tpidr_el1, tpidr_el1);
  return (thread_t *)tpidr_el1;
}

static inline
void _thread_wrapper()
{
  enable_interrupt();

  thread_t *cur = get_current();

  LOG("Running thread wrapper"); DEBUG_THREAD(cur->func);
  cur->func(cur->param);

  cur->status = THREAD_STATUS_ZOMBIE;
  
  schedule();
}

static inline
thread_t *_make_thread(thread_func_t func, uint64_t param)
{
  thread_t *ret = (thread_t *)kmalloc(sizeof(thread_t));
  
  *ret = (thread_t){
    .func = func,
    .param = param,

    .status = THREAD_STATUS_RUNABLE,
    .retcode = 0,

    .stack = allocate_page(256),
    .stack_size = 256 << 12,

    .regs = kmalloc(sizeof(cpu_state_t)),

    .process = 0x0,
  };
  memset(ret->regs, 0, sizeof(cpu_state_t));
  ret->regs->sp = (uint64_t)ret->stack + ret->stack_size - 64;
  ret->regs->lr = (uint64_t)_thread_wrapper;
  ret->regs->fp = ret->regs->sp;

  return ret;
}

void run_thread(thread_func_t func, uint64_t param)
{
  // Runs before enable interrupt
  QUEUE_PUSH_BACK(global_run_queue, _make_thread(func, param));
}

static inline
int32_t _idle(uint64_t)
{
  DEBUG_THREAD("enter idle thread");

  for (;;) {
    enable_interrupt();
    schedule();
    DEBUG_THREAD(global_gc_queue->size);
    // DEBUG_THREAD("back to idle thread");
  }

  return 0;
}

void thread_init(void)
{
  global_run_queue = make_queue();
  global_wait_queue = make_queue();
  global_gc_queue = make_queue();

  thread_t *th = kmalloc(sizeof(thread_t));
  W_SYSREG(tpidr_el1, th);

  run_thread(_idle, 0x0);
}

void schedule(void)
{
  disable_interrupt();

  thread_t *th = get_current(), *next = 0x0;
  LOG("Run thread"); LOG(global_run_queue->size); DEBUG_THREAD(th);
  LOG(th->stack); DEBUG_THREAD(th->regs->sp);

  while (global_run_queue->size > 0 && !next) {
    next = QUEUE_POP_FRONT(thread_t, global_run_queue);
    if (next->status == THREAD_STATUS_ZOMBIE) {
      QUEUE_PUSH_BACK(global_gc_queue, next);
      next = 0x0;
    } else
      QUEUE_PUSH_BACK(global_run_queue, next);
  }
  
  if (next && th != next)
    context_switch(th->regs, next->regs, next);

  enable_interrupt();
}