#include <asm/interrupt.h>
#include <asm/system.h>
#include <errno.h>
#include <kernel/con.h>
#include <kernel/kernel.h>
#include <kernel/memory.h>
#include <kernel/sched.h>
#include <signal.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>

/*
 * WARNING!
 * don't change this without also updating
 * the number of TSS descriptors in the GDT
 * (located in boot.s)
 */
#define NRTASKS 64

#define USRSTART ((void*) 0x40000000)

#define pages(n) \
  ((n / 4096) + 1)

#define tables(n) \
  ((n / 1024) + 1)

#define getpage(name, p) ({\
  void **_next; \
  _next = map_page(p); \
  name = p; \
  p = *_next; \
  })

extern char _usrbin_start;
extern char _usrbin_size;

struct task_struct *ctask;
uint32_t ctaskn;

static struct task_struct tasks[NRTASKS];

static pid_t nextpid;

static inline void empty_table(void *t) {
  int i;

  for(i = 0; i < PGENT; i++)
    ((uint32_t*) t)[i] = 0;
}

struct task_struct *create_proc(void *bin, size_t len) {
  int i, j;
  int c;
  size_t totpages, npages, ntables;
  void *p, *page, *tab, *map;
  void *pdir, *stackt, *stackp, *kstackp;
  struct task_struct *task = NULL;

  npages = pages(len);
  ntables = tables(pages(len));
  totpages = npages + ntables + 4;

  /* find an unused task structure */
  for(i = 0; i < NRTASKS; i++) {
    if(!tasks[i].pid) {
      task = &tasks[i];
      /* update the GDT */
      set_tss(i, &tasks[i].tss);
      break;
    }
  }
  if(task == NULL)
    return NULL;

  /* find an unused PID */
  for(i = 0; i < NRTASKS; i++) {
    if(tasks[i].pid == nextpid) {
      /* TODO: handle PID exhaustion */
      nextpid = (nextpid + 1) % 0x8000;
      if(nextpid == 0)
        nextpid++;
      i = 0;
    }
  }

  /* here we allocate pages to setup the process'
   * address space. in addition to the pages that
   * will hold the binary image, the following
   * pages are allocated:
   *  - page directory (1 page)
   *  - page tables (1 per every 1024 pages)
   *  - stack page (1 page)
   *  - kernel stack page (1 page)
   *  - stack table (1 page)
   */
  p = alloc_physical_pages(totpages);
  if(p == NULL)
    return NULL;

  getpage(pdir, p);
  getpage(stackt, p);
  getpage(stackp, p);
  getpage(kstackp, p);

  /* populate the page directory */
  c = npages;
  map = map_page(pdir);
  empty_table(map);
  for(i = 0; i < ntables; i++) {
    /* add the table entry to the directory */
    getpage(tab, p);
    map = map_page(pdir);
    ((uint32_t*) map)[i + (PGENT / 4)] = (uint32_t) tab | 0x007;

    /* populate the page table */
    map = map_page(tab);
    for(j = 0; j < PGENT; j++) {
      ((uint32_t*) map)[j] = 0;
      if(c <= 0)
        continue;

      getpage(page, p);
      map = map_page(tab);
      ((uint32_t*) map)[j] = (uint32_t) page | 0x007;
      c--;
    }
  }

  /* populate the stack table */
  map = map_page(stackt);
  empty_table(map);
  ((uint32_t*) map)[PGENT - 3] = (uint32_t) stackp | 0x007;
  /*
   * NOTE: we leave a single 4KB gap between the kernel
   * and user's stacks so we can detect overflows of the
   * kernel's stack.
   */
  ((uint32_t*) map)[PGENT - 1] = (uint32_t) kstackp | 0x003;

  /* add the stack page table to the directory */
  map = map_page(pdir);
  ((uint32_t*) map)[PGENT - 1] = (uint32_t) stackt | 0x007;
  /* add the kernel's page table to the directory */
  ((uint32_t*) map)[0] = (uint32_t) ktab | 0x003;

  __asm__ volatile ("mov %%eax, %%cr3" :: "a" (pdir));

  /* load in the user binary */
  memcpy(USRSTART, bin, len);

  /* initialize the task struct */
  memset(task, 0, sizeof(struct task_struct));
  task->pid = nextpid;
  task->state = TSTATE_RUNNING;

  /* initialize the task's state */
  task->tss.cr3 = (uint32_t) pdir;
  task->tss.cs = 0x1B;
  task->tss.ds = 0x23;
  task->tss.es = 0x23;
  task->tss.fs = 0x23;
  task->tss.gs = 0x23;
  task->tss.ss = 0x23;
  task->tss.eip = (uint32_t) USRSTART;
  task->tss.esp = 0xFFFFDFFF;
  task->tss.esp0 = 0xFFFFFFFF;
  task->tss.ss0 = 0x10;
  __asm__ (
      "pushf\n" \
      "pop %%eax" \
      : "=a" (task->tss.eflags) \
      );

  return task;
}

void kill_proc(struct task_struct *task) {
  int i, j;
  uint32_t *x;

  task->pid = 0;

  for(i = 0; i < PGENT; i++) {
    x = map_page((void*) task->tss.cr3);
    if(!x[i])
      continue;
    x = map_page((void*) (x[i] & 0xFFFFF000));
    for(j = 0; j < PGENT; j++) {
      if(x[i] & 1)
        free_physical_page((void*) (x[i] & 0xFFFFF000));
    }
  }
  free_physical_page((void*) task->tss.cr3);

  reschedule();
}

void sched_init(void) {
  int i;
  struct task_struct *task;

  memset(&tasks, 0, sizeof(tasks));
  for(i = 0; i < NRTASKS; i++)
    clear_tss(i);

  nextpid = 1;
  ctask = NULL;

  create_proc(&_usrbin_start, (size_t) &_usrbin_size);
  /* create_proc(&_usrbin_start, (size_t) &_usrbin_size); */

  switch_to(0, &tasks[0]);
}

void reschedule(void) {
  int i;
  uint32_t n;

  if(ctask == NULL)
    return;

  n = ctaskn;
  while(1) {
    /* check alarms and signals */
    for(i = 0; i < NRTASKS; i++) {
      if(tasks[i].alarm && tasks[i].alarm < ticks) {
        tasks[i].alarm = 0;
        tasks[i].signal |= 1 << (SIGALRM - 1);
      }
      if(tasks[i].signal && tasks[i].state == TSTATE_INTERRUPTIBLE)
        tasks[i].state = TSTATE_RUNNING;
    }

    /* check if a process is ready to run */
    for(i = 0; i < NRTASKS; i++) {
      n = (n + 1) % NRTASKS;
      if(tasks[n].pid && tasks[n].state == TSTATE_RUNNING) {
        switch_to(n, &tasks[n]);
        return;
      }
    }

    /* if not, then halt the machine (temporarily re-enabling interrupts) */
    __asm__ volatile (
        "pushf\n"
        "sti\n"
        "hlt\n"
        "popf\n"
        );
  }
}

void sched_tick(void) {
  reschedule();
}

void wake_up(struct task_struct **task) {
  if(task == NULL || *task == NULL)
    return;

  (*task)->state = TSTATE_RUNNING;
  *task = NULL;
}

void sleep_on(struct task_struct **waiting) {
  struct task_struct *t;

  if(waiting == NULL)
    return;

  t = *waiting;
  *waiting = ctask;
  ctask->state = TSTATE_UNINTERRUPTIBLE;
  reschedule();

  *waiting = t;
  if(t)
    t->state = TSTATE_RUNNING;
}

void interruptible_sleep_on(struct task_struct **waiting) {
  struct task_struct *t;

  if(waiting == NULL)
    return;

  t = *waiting;
  *waiting = ctask;
  while(1) {
    ctask->state = TSTATE_INTERRUPTIBLE;
    reschedule();

    if(*waiting == NULL || *waiting == ctask)
      break;

    (*waiting)->state = TSTATE_RUNNING;
  }

  *waiting = NULL;
  if(t)
    t->state = TSTATE_RUNNING;
}

void sighandler_default(uint32_t sig) {
  switch(sig) {
    case SIGKILL:
      kill_proc(ctask);
      break;
    default:
      printk("[kernel] Handling signal 0x%02x for PID 0x%04x\n", sig, ctask->pid);
      break;
  }
}

int sys_kill(pid_t pid, int sig) {
  int i;

  if(sig == 0)
    return 0;

  for(i = 0; i < NRTASKS; i++) {
    if(tasks[i].pid == pid) {
      tasks[i].signal |= (1 << (sig - 1));
      return 0;
    }
  }

  return -ESRCH;
}