#include <stdlib.h>
#include <string.h>

#include <drivers/irq.h>

#include <scheduler.h>

#define STACK_SIZE 4096

task_t* task_list = NULL;
task_t* current_task = NULL;
uint32_t next_task_id = 1;

static void schedule(void);

/* scheduler init */
void scheduler_init(void)
{
  set_irq_handler(0, (irq_func_t) scheduler_tick);
}

static void schedule(void)
{
  if (!task_list)
  {
    return;
  }

  /* wake sleeping tasks */
  task_t* t = task_list;

  extern uint64_t pit_ticks;

  do
  {
    if (t->state == TASK_SLEEPING && pit_ticks >= t->wakeup_tick)
    {
      t->state = TASK_READY;
    }

    t = t->next;
  }
  while (t != task_list);

  /* find next ready task */
  task_t* start = current_task;

  do
  {
    current_task = current_task->next;

    if (current_task->state == TASK_READY)
    {
      return;
    }
  }
  while (current_task != start);
}


/* tick handler */
void scheduler_tick(registers_t* regs)
{
  if (!task_list)
  {
    printf("!task_list\n");
    
    return;
  }

  if (!current_task)
  {
    printf("!current_task\n");
    current_task = task_list;
  }

  /* save current task state */
  *(current_task->regs) = *regs;

  /* pick next */
  schedule();

  /* load next task state */
  *regs = *(current_task->regs);
}

/* task creation */
task_t* create_task(void (*entry)())
{
  task_t* task = malloc(sizeof(task_t));
  
  if (!task)
  {
    return NULL;
  }

  memset(task, 0, sizeof(task_t));

  /* allocate stack */
  uint8_t* stack = malloc(STACK_SIZE);
  
  if (!stack)
  {
    free(task);
    return NULL;
  }

  printf("Creating task %i @ %p\n", next_task_id, entry);

  /*
    align stack to 16 bytes for safety (SSE friendly).
    stack grows downward.
  */
  uintptr_t stack_top = (uintptr_t)(stack + STACK_SIZE);
  stack_top &= ~0xF;

  registers_t* regs = (registers_t*) (stack_top - sizeof(registers_t));
  memset(regs, 0, sizeof(registers_t));

  /* initialize register frame */
  regs->eip = (uint32_t) entry;
  regs->cs = 0x08;     /* Kernel code segment */
  regs->eflags = 0x202; /* IF=1 */

  task->id = next_task_id++;
  task->regs = regs;
  task->stack_base = (uint32_t*) stack;
  task->state = TASK_READY;

  /* insert into circular list */
  if (!task_list)
  {
    task_list = task;
    task->next = task;
  }
  else
  {
    task_t* last = task_list;
    while (last->next != task_list)
    {
      last = last->next;
    }

    last->next = task;
    task->next = task_list;
  }

  printf("Task %i @ %p created successfully\n", next_task_id - 1, entry);

  return task;
}

/* task destruction */
void destroy_task(task_t* task)
{
  free(task->stack_base);
  free(task);
}