#include "paging.h"
#include "window.h"
#include "drive.h"
#include "panic.h"
#include "task.h"
#include "util.h"
#include "pmap.h"
#include "elf.h"
#include "log.h"
#include "vbe.h"

#define MAX_WINDOWS 64
#define ACTION_BUFFER_PAGES 1
#define AB_END(buf) (buf + (ACTION_BUFFER_PAGES * 4096) / sizeof(struct window_action))

#define GSC(x, y) (x / 4 + y / 4 >= 256 ? 255 : x / 4 + y / 4)
#define BACKGROUND(x, y) ((struct pixel){.r = GSC(x / 2, y / 3), .g = GSC(x, y), .b = GSC(x, y / 2)})

#define BORDER_COLOR ((struct pixel){.r = 0x11, .g = 0x11, .b = 0x11})

#define MOVE_BY 5

static struct window {
  const volatile struct pixel *pixel_buffer_pma;
  uint16_t width;
  uint16_t height;

  uint16_t xpos;
  uint16_t ypos;

  struct window_action *action_buffer;
  struct window_action *next_action_read;
  struct window_action *next_action_write;

  struct window *above;
  struct window *below;

  struct task_state *from_task;
} windows[MAX_WINDOWS];

static struct window *bottom_window = 0;
static struct window *top_window = 0;

static struct pixel *buffer;
static uint32_t buffer_pages;
static uint32_t pix_clear_mask;

static void load_mode_params() {
  if (buffer)
    free_pages(buffer, buffer_pages);

  const uint32_t fb_len = VBE_MODE_INFO->width * VBE_MODE_INFO->height;
  buffer_pages = (fb_len * sizeof(struct pixel) - 1) / 4096 + 1;
  buffer = allocate_kernel_pages(buffer_pages);

  #define MASK(offset, length) ((~((1 << offset) - 1)) - (offset + length == 32 ? 0 : (~((1 << (offset + length)) - 1))))
  pix_clear_mask = ~(
    MASK(VBE_MODE_INFO->red_off,   VBE_MODE_INFO->red_len)   |
    MASK(VBE_MODE_INFO->green_off, VBE_MODE_INFO->green_len) |
    MASK(VBE_MODE_INFO->blue_off,  VBE_MODE_INFO->blue_len)  |
    MASK(VBE_MODE_INFO->alpha_off, VBE_MODE_INFO->alpha_len)
  );
}

static void blit() {
  void *row = VBE_MODE_INFO->frame_buf;
  const struct pixel *from = buffer;
  for (uint16_t y = 0; y < VBE_MODE_INFO->height; ++y) {
    uint32_t *out_pix = row;
    for (uint16_t x = 0; x < VBE_MODE_INFO->width; ++x) {
      *out_pix &= pix_clear_mask;
      *out_pix |= (from->r >> (8 - VBE_MODE_INFO->red_len))   << VBE_MODE_INFO->red_off;
      *out_pix |= (from->g >> (8 - VBE_MODE_INFO->green_len)) << VBE_MODE_INFO->green_off;
      *out_pix |= (from->b >> (8 - VBE_MODE_INFO->blue_len))  << VBE_MODE_INFO->blue_off;
      out_pix = (uint32_t *)((void *)out_pix + VBE_MODE_INFO->bpp / 8);
      ++from;
    }
    row += VBE_MODE_INFO->pitch;
  }
}

static inline void draw_hz_line(uint16_t y, uint16_t xs, uint16_t xm, struct pixel color) {
  if (y >= VBE_MODE_INFO->height)
    return;
  if (xs >= VBE_MODE_INFO->width)
    xs = 0;
  if (xm > VBE_MODE_INFO->width)
    xm = VBE_MODE_INFO->width;
  for (uint16_t x = xs; x < xm; ++x)
    buffer[y * VBE_MODE_INFO->width + x] = color;
}

static inline void draw_vt_line(uint16_t x, uint16_t ys, uint16_t ym, struct pixel color) {
  if (x >= VBE_MODE_INFO->width)
    return;
  if (ys >= VBE_MODE_INFO->height)
    ys = 0;
  if (ym > VBE_MODE_INFO->height)
    ym = VBE_MODE_INFO->height;
  for (uint16_t y = ys; y < ym; ++y)
    buffer[y * VBE_MODE_INFO->width + x] = color;
}

static void paint_and_above(const struct window *w) {
  switch_to_kernel_cr3();
  for (const struct window *i = w; i; i = i->above) {
    const uint16_t ys = i->ypos < VBE_MODE_INFO->height ? 0 : -i->ypos;
    const uint16_t xs = i->xpos < VBE_MODE_INFO->width  ? 0 : -i->xpos;
    const uint16_t ym = ((i->ypos + i->height) & 0xffff) <= VBE_MODE_INFO->height ? i->height : VBE_MODE_INFO->height - i->ypos;
    const uint16_t xm = ((i->xpos + i->width)  & 0xffff) <= VBE_MODE_INFO->width  ? i->width  : VBE_MODE_INFO->width  - i->xpos;

  //logf(LOG_INFO, "y: %n .. %n", ys, ym - 1);
  //logf(LOG_INFO, "x: %n .. %n", xs, xm - 1);

    for (uint16_t y = ys; y < ym; ++y)
      for (uint16_t x = xs; x < xm; ++x)
        buffer[(y + i->ypos) * VBE_MODE_INFO->width + x + i->xpos] = i->pixel_buffer_pma[y * i->width + x];

    draw_hz_line(i->ypos             - 2, i->xpos - 2, i->xpos + i->width + 2, BORDER_COLOR);
    draw_hz_line(i->ypos             - 1, i->xpos - 2, i->xpos + i->width + 2, BORDER_COLOR);
    draw_hz_line(i->ypos + i->height,     i->xpos - 2, i->xpos + i->width + 2, BORDER_COLOR);
    draw_hz_line(i->ypos + i->height + 1, i->xpos - 2, i->xpos + i->width + 2, BORDER_COLOR);
    draw_vt_line(i->xpos             - 2, i->ypos,     i->ypos + i->height,    BORDER_COLOR);
    draw_vt_line(i->xpos             - 1, i->ypos,     i->ypos + i->height,    BORDER_COLOR);
    draw_vt_line(i->xpos + i->width,      i->ypos,     i->ypos + i->height,    BORDER_COLOR);
    draw_vt_line(i->xpos + i->width  + 1, i->ypos,     i->ypos + i->height,    BORDER_COLOR);
  }

  blit();
  switch_to_task_cr3();
}

static void paint_bg() {
  for (uint16_t y = 0; y < VBE_MODE_INFO->height; ++y)
    for (uint16_t x = 0; x < VBE_MODE_INFO->width; ++x)
      buffer[y * VBE_MODE_INFO->width + x] = BACKGROUND(x, y);
}

static void paint_all() {
  paint_bg();
  paint_and_above(bottom_window);
}

struct window *new_window(uint16_t width, uint16_t height, const struct pixel *pixel_buffer) {
  if (!pixel_buffer) {
    logf(LOG_WARN, "Refusing to create window with null pixel buffer for task %s.", active_task->name);
    return 0;
  }

  struct window *w;
  for (uint8_t i = 0; i < MAX_WINDOWS; ++i)
    if (!windows[i].pixel_buffer_pma) {
      w = &windows[i];
      goto got_window;
    }
  return 0;
got_window:

  w->pixel_buffer_pma = vma_to_pma(active_task->page_directory, pixel_buffer);
  w->width = width;
  w->height = height;
  w->xpos = (VBE_MODE_INFO->width  / 2) - (width  / 2);
  w->ypos = (VBE_MODE_INFO->height / 2) - (height / 2);

  struct window_action *const ab = allocate_kernel_pages(ACTION_BUFFER_PAGES);
  w->action_buffer = ab;
  w->next_action_read = ab;
  w->next_action_write = ab;

  if (top_window)
    top_window->above = w;
  else
    bottom_window = w;
  w->above = 0;
  w->below = top_window;
  top_window = w;

  w->from_task = active_task;

  paint_and_above(w);
  return w;
}

static void del_no_paint(struct window *w) {
  if (w == top_window)
    top_window = w->below;
  if (w == bottom_window)
    bottom_window = w->above;
  if (w->below)
    w->below->above = w->above;
  if (w->above)
    w->above->below = w->below;

  free_pages(w->action_buffer, ACTION_BUFFER_PAGES);
  w->pixel_buffer_pma = 0;
}

void del_window(struct window *w) {
  del_no_paint(w);
  paint_all();
}

void delete_any_windows_from(struct task_state *tstate) {
  bool need_to_paint = false;
  for (struct window *w = windows; w < windows + MAX_WINDOWS; ++w)
    if (w->pixel_buffer_pma && (w->from_task == tstate)) {
      del_no_paint(w);
      need_to_paint = true;
    }
  if (need_to_paint)
    paint_all();
}

void resize_window(struct window *w, uint16_t width, uint16_t height, const struct pixel *pixel_buffer) {
  const bool smaller = (width < w->width) || (height < w->height);

  w->width = width;
  w->height = height;
  reassign_pixel_buffer(w, pixel_buffer);

  if (smaller)
    paint_all();
  else
    paint_and_above(w);
}

void reassign_pixel_buffer(struct window *w, const struct pixel *pixel_buffer) {
  w->pixel_buffer_pma = vma_to_pma(active_task->page_directory, pixel_buffer);
}

void push_window_paint(const struct window *w) {
  paint_and_above(w);
}

struct window_action next_window_action(struct window *w) {
  if (w->next_action_write == w->next_action_read)
    return (struct window_action){.action_type = NONE};
  const struct window_action *const action = w->next_action_read;
  if (++(w->next_action_read) >= AB_END(w->action_buffer))
    w->next_action_read = w->action_buffer;
  return *action;
}

static void send_action(struct window *w, struct window_action packet) {
  struct window_action *next_next = w->next_action_write + 1;
  if (next_next >= AB_END(w->action_buffer))
    next_next = w->action_buffer;
  if (next_next != w->next_action_read) {
    *(w->next_action_write) = packet;
    w->next_action_write = next_next;
    unwait(w->from_task, (struct wait){.mode = WINDOW_ACTION});
  }
}

#define RUN_COMMAND_FILE "sys/winspace.rc"

enum wm_action {
  WM_SHUFFLE_UP,
  WM_SHUFFLE_DOWN,
  WM_MOVE_LEFT,
  WM_MOVE_RIGHT,
  WM_MOVE_UP,
  WM_MOVE_DOWN,
  WM_RUN_COMMAND,

  N_WM_ACTIONS
};

static struct key_packet keybinds[] = {
  {.key_id = KEY_PAGE_DOWN,   .modifiers = WINS},
  {.key_id = KEY_PAGE_UP,     .modifiers = WINS},
  {.key_id = KEY_LEFT_ARROW,  .modifiers = WINS},
  {.key_id = KEY_RIGHT_ARROW, .modifiers = WINS},
  {.key_id = KEY_UP_ARROW,    .modifiers = WINS},
  {.key_id = KEY_DOWN_ARROW,  .modifiers = WINS},
  {.key_id = KEY_SPACE,       .modifiers = WINS}
};

static inline bool fuzzy_key_match(struct key_packet t, struct key_packet a) {
  if (t.key_id != a.key_id)
    return false;

  if (((t.modifiers & SHIFTS) == SHIFTS) && (a.modifiers & SHIFTS))
    a.modifiers |= SHIFTS;
  if (((t.modifiers & CTRLS)  == CTRLS)  && (a.modifiers & CTRLS))
    a.modifiers |= CTRLS;
  if (((t.modifiers & ALTS)   == ALTS)   && (a.modifiers & ALTS))
    a.modifiers |= ALTS;
  if (((t.modifiers & WINS)   == WINS)   && (a.modifiers & WINS))
    a.modifiers |= WINS;

  return a.modifiers == t.modifiers;
}

static char *run_command;
static char *run_command_pass = "";

void init_win() {
  buffer = 0;
  load_mode_params();

  paint_bg();
  blit();

  const file_id_t fid = drives->get_file(drives, RUN_COMMAND_FILE);
  if (!fid)
    PANIC("Couldn't open " RUN_COMMAND_FILE ".");

  const uint32_t len = drives->get_file_length(drives, fid);
  run_command = allocate_kernel_pages(len / 4096 + 1);
  if (!run_command)
    PANIC("Couldn't make buffer for " RUN_COMMAND_FILE " contents.");

  fmcpy(run_command, drives, fid, 0, len);
  run_command[len] = '\0';

  drives->free_file(drives, fid);

  for (char *c = run_command; *c; ++c)
    if (*c == ' ') {
      *c = '\0';
      run_command_pass = c + 1;
      break;
    }
}

void on_action(struct window_action packet) {
//logf(LOG_INFO, "Window action, top window = 0x%d from %s.", top_window, top_window->from_task->name);

  if (packet.action_type == NOT_READY)
    return;

  if (packet.action_type == KEY_DOWN)
    for (uint8_t i = 0; i < N_WM_ACTIONS; ++i)
      if (fuzzy_key_match(keybinds[i], packet.as_key)) {
        switch_to_kernel_cr3();
        struct window *old_top, *old_bottom;
        switch (i) {
        case WM_SHUFFLE_UP:
          if (!top_window || !top_window->below)
            break;
          old_top = top_window;
          old_bottom = bottom_window;
          top_window = old_top->below;
          top_window->above = 0;
          old_top->below = 0;
          old_top->above = old_bottom;
          old_bottom->below = old_top;
          bottom_window = old_top;
          paint_and_above(bottom_window->above);
          break;
        case WM_SHUFFLE_DOWN:
          if (!top_window || !top_window->below)
            break;
          old_top = top_window;
          old_bottom = bottom_window;
          bottom_window = old_bottom->above;
          bottom_window->below = 0;
          old_bottom->above = 0;
          old_bottom->below = old_top;
          old_top->above = old_bottom;
          top_window = old_bottom;
          paint_and_above(top_window);
          break;
        case WM_MOVE_LEFT:
          if (top_window) {
            top_window->xpos -= MOVE_BY;
            paint_all();
          }
          break;
        case WM_MOVE_RIGHT:
          if (top_window) {
            top_window->xpos += MOVE_BY;
            paint_all();
          }
          break;
        case WM_MOVE_UP:
          if (top_window) {
            top_window->ypos -= MOVE_BY;
            paint_all();
          }
          break;
        case WM_MOVE_DOWN:
          if (top_window) {
            top_window->ypos += MOVE_BY;
            paint_all();
          }
          break;
        case WM_RUN_COMMAND:
          if (!try_elf_run(drives, run_command, run_command_pass, 0))
            PANIC("Couldn't run program listed in " RUN_COMMAND_FILE ".");
        }
        switch_to_task_cr3();
        return;
      }

  if (top_window)
    send_action(top_window, packet);
}