#include "scheduler.h"
#include <stdlib.h>

void task_a(void) {
  while (true) printk("A");
}

void task_b(void) {
  while (true) printk("B");
}

static uint8_t stack_a[4096] = {0};
static uint8_t stack_b[4096] = {0};
constexpr uint16_t max_tasks = 2;
static cpu_state_t* tasks[max_tasks];

Scheduler::Scheduler(uint16_t cs) : InterruptHandler(0x00) {
  add_task(stack_a, cs, &task_a);
  add_task(stack_b, cs, task_b);
}

bool Scheduler::add_task(uint8_t* stack, uint16_t cs, void (*entry)()) {
  if (last_task_id == uint16_t_max) return false;

  uint16_t id = last_task_id++;
  ++num_tasks;

  cpu_state_t* cpu_state = reinterpret_cast<cpu_state_t*>(stack + 4096 - sizeof(cpu_state_t));
  cpu_state->eip = reinterpret_cast<uint32_t>(entry);
  cpu_state->cs = cs;
  cpu_state->eflags = 0x202;

  tasks[id] = cpu_state;
  printk("adding task: cs ", uhex{cpu_state->cs}, " stack ", uhex{reinterpret_cast<uint32_t>(tasks[id])}, " entry ",
         uhex{cpu_state->eip}, '\n');
  return true;
}

cpu_state_t* Scheduler::trigger(cpu_state_t* cpu) {
  tasks[current_task] = cpu;

  if (++current_task >= num_tasks) current_task %= num_tasks;

  // printk("swapped to task ", current_task, '\n');
  return tasks[current_task];
}