#include "node.h"
#include "constants.h"
#include "editor.h"
#include "helper_functions.h"
#include "highlighters.h"
#include "lsp_handler.h"
#include "snake.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

Node root;

void init_root_node() {
  root.type = ROOT;
  root.parent = NULL;
  root.Root.windows = NULL;
  root.Root.count = 0;
  root.Root.focused = NULL;
}

Node *create_window(uint8_t x, uint8_t y, uint8_t width, uint8_t height,
                    _Bool border, _Bool direction) {
  Node *window = malloc(sizeof(Node));
  if (!window)
    die("malloc failed");
  window->type = WINDOW;
  window->parent = &root;
  window->Window.x = x;
  window->Window.y = y;
  window->Window.width = width;
  window->Window.height = height;
  window->Window.border = border;
  Node *root_block = calloc(1, sizeof(Node));
  if (!root_block) {
    free(window);
    die("calloc failed");
  }
  root_block->type = BLOCK;
  root_block->parent = window;
  root_block->Block.children = NULL;
  root_block->Block.count = 0;
  root_block->Block.calc_x = x + (window->Window.border ? 1 : 0);
  root_block->Block.calc_y = y + (window->Window.border ? 1 : 0);
  root_block->Block.calc_rows = height - (window->Window.border ? 2 : 0);
  root_block->Block.calc_cols = width - (window->Window.border ? 2 : 0);
  root_block->Block.direction = direction;
  window->Window.root_block_node = root_block;
  window->Window.visible = 1;
  window->Window.border = 1;
  Node **tmp =
      realloc(root.Root.windows, sizeof(Node *) * (root.Root.count + 1));
  if (!tmp) {
    die("realloc failed");
  }
  root.Root.windows = tmp;
  root.Root.windows[root.Root.count] = window;
  root.Root.count++;
  return window;
}

Node *create_buffer(Node *parent_block, const Color *fg, const Color *bg,
                    uint8_t size, BufferData *data) {
  if (parent_block->type != BLOCK)
    die("parent is not a block");
  Node *buffer = calloc(1, sizeof(Node));
  if (!buffer)
    die("calloc failed");
  buffer->type = BUFFER;
  buffer->parent = parent_block;
  buffer->Buffer.data = data;
  buffer->Buffer.fg = (Color *)fg;
  buffer->Buffer.bg = (Color *)bg;
  buffer->Buffer.size = size;
  Node **tmp = realloc(parent_block->Block.children,
                       sizeof(Node *) * (parent_block->Block.count + 1));
  if (!tmp)
    die("realloc failed");
  parent_block->Block.children = tmp;
  parent_block->Block.children[parent_block->Block.count] = buffer;
  parent_block->Block.count++;
  return buffer;
}

void build_line_highlight_map(Line *lh, ColHighlight *map, Color *default_fg,
                              Color *default_bg) {
  for (size_t i = 0; i < lh->cluster_count; i++) {
    map[i].fg = default_fg;
    map[i].bg = default_bg;
    map[i].flags = 0;
  }
  if (!(!lh || lh->highlight_count == 0)) {
    for (size_t hi = 0; hi < lh->highlight_count; hi++) {
      Highlight *h = &lh->highlights[hi];
      size_t start = h->start_col;
      size_t end = h->end_col;
      if (start >= lh->cluster_count)
        continue;
      if (end > lh->cluster_count)
        end = lh->cluster_count;
      for (size_t i = start; i < end; i++) {
        if (h->fg)
          map[i].fg = h->fg;
        if (h->bg)
          map[i].bg = h->bg;
        map[i].flags |= h->flags;
      }
    }
  }
}

void clear_hl(BufferData *e) {
  for (size_t i = 0; i < e->line_count; i++) {
    free(e->lines[i]->highlights);
    e->lines[i]->highlight_count = 0;
    e->lines[i]->highlights = NULL;
  }
}

static void apply_highlighting(Node *buffer, BufferData *e) {
  clear_hl(e);
  if (buffer->Buffer.hl_func_basic)
    buffer->Buffer.hl_func_basic(buffer);
  if (buffer->Buffer.hl_func_syntax)
    buffer->Buffer.hl_func_syntax(buffer);
}

static void fill_blank_line(Node *buffer, int y, int screen_row, int max_w) {
  for (int col = 0; col < max_w; col++)
    insert_screen(y + screen_row, buffer->Buffer.calc_x + col, " ",
                  buffer->Buffer.fg, buffer->Buffer.bg, 0);
}

Line *merge_virtual_line(Line *physical, Line *virtual) {
  if (!physical)
    return NULL;
  if (!virtual || virtual->cluster_count == 0)
    return clone_line(physical);
  size_t total_clusters = (physical->cluster_count > virtual->cluster_count)
                              ? physical->cluster_count
                              : virtual->cluster_count;
  Line *merged = calloc(1, sizeof(Line));
  merged->cluster_count = total_clusters;
  merged->clusters = calloc(total_clusters, sizeof(Cluster));
  for (size_t i = 0; i < total_clusters; i++) {
    Cluster c = {0};
    if (i < virtual->cluster_count && virtual->clusters[i].utf8 &&
        virtual->clusters[i].utf8[0] != '\0') {
      c.utf8 = strdup(virtual->clusters[i].utf8);
      c.width = virtual->clusters[i].width;
    } else if (i < physical->cluster_count) {
      c.utf8 = strdup(physical->clusters[i].utf8);
      c.width = physical->clusters[i].width;
    }
    merged->clusters[i] = c;
  }
  size_t total_hl = merged->highlight_count +
                    (virtual ? virtual->highlight_count : 0) +
                    physical->highlight_count;
  merged->highlight_count = 0;
  if (total_hl > 0) {
    merged->highlights = malloc(total_hl * sizeof(Highlight));
    if (!merged->highlights)
      die("malloc failed");
    for (size_t i = 0; i < physical->highlight_count; i++)
      merged->highlights[merged->highlight_count++] = physical->highlights[i];
    if (virtual) {
      for (size_t i = 0; i < virtual->highlight_count; i++)
        merged->highlights[merged->highlight_count++] = virtual->highlights[i];
    }
  }
  return merged;
}

static int draw_line_clusters(Node *buffer, Line *line, size_t line_idx,
                              size_t cursor_row, size_t cursor_col,
                              size_t scroll_col, int screen_row, int max_w,
                              int *cursor_screen_row, int *cursor_screen_col) {
  int screen_col = 0;
  int logical_col = 0;
  ColHighlight *hl_map = malloc(line->cluster_count * sizeof(ColHighlight));
  build_line_highlight_map(line, hl_map, buffer->Buffer.fg, buffer->Buffer.bg);
  for (size_t i = 0; i < line->cluster_count && screen_col < max_w; i++) {
    char *cluster = line->clusters[i].utf8;
    int w = line->clusters[i].width;
    if (logical_col + w <= (int)scroll_col) {
      logical_col += w;
      continue;
    }
    Color *fg = hl_map[i].fg;
    Color *bg = hl_map[i].bg;
    uint8_t flags = hl_map[i].flags;
    if (logical_col < (int)scroll_col && logical_col + w > (int)scroll_col) {
      for (int j = 0; j < (logical_col + w) - (int)scroll_col; j++)
        insert_screen(buffer->Buffer.calc_y + screen_row,
                      buffer->Buffer.calc_x + screen_col++, " ",
                      buffer->Buffer.fg, buffer->Buffer.bg, 0);
      logical_col += w - (int)scroll_col;
      continue;
    }
    if (logical_col >= (int)(scroll_col + max_w))
      break;
    int visible_start =
        (scroll_col > (size_t)logical_col) ? (scroll_col - logical_col) : 0;
    int visible_width = w - visible_start;
    if (visible_width > max_w - screen_col)
      visible_width = max_w - screen_col;
    if (visible_width > 0) {
      insert_screen(buffer->Buffer.calc_y + screen_row,
                    buffer->Buffer.calc_x + screen_col, cluster, fg, bg, flags);
      for (int j = 1; j < visible_width; j++)
        insert_screen(buffer->Buffer.calc_y + screen_row,
                      buffer->Buffer.calc_x + screen_col + j, "\u200B", fg, bg,
                      flags);
      if (line_idx == cursor_row && (size_t)i == cursor_col) {
        *cursor_screen_row = buffer->Buffer.calc_y + screen_row;
        *cursor_screen_col = buffer->Buffer.calc_x + screen_col;
      }
      screen_col += visible_width;
    }
    logical_col += w;
  }
  if (line->cluster_count == 0 && line_idx == cursor_row) {
    *cursor_screen_row = buffer->Buffer.calc_y + screen_row;
    *cursor_screen_col = buffer->Buffer.calc_x;
  }
  if (line->cluster_count == cursor_col && line_idx == cursor_row) {
    *cursor_screen_row = buffer->Buffer.calc_y + screen_row;
    *cursor_screen_col = buffer->Buffer.calc_x + screen_col;
  }
  free(hl_map);
  return screen_col;
}

static void pad_remaining_cols(Node *buffer, int y, int screen_row,
                               int screen_col, int max_w) {
  for (; screen_col < max_w; screen_col++)
    insert_screen(y + screen_row, buffer->Buffer.calc_x + screen_col, " ",
                  buffer->Buffer.fg, buffer->Buffer.bg, 0);
}

void draw_buffer(Node *buffer) {
  if (!buffer || buffer->type != BUFFER)
    die("node is not a buffer");
  BufferData *e = buffer->Buffer.data;
  if (!e || e->line_count == 0)
    return;
  apply_highlighting(buffer, e);
  size_t total_lines = e->line_count;
  size_t scroll_row = e->scroll_row;
  VirtLine **virt_lines;
  Line **lines;
  Coords coords = filter_visible_lines(e, &virt_lines, &lines, &total_lines);
  size_t scroll_col = e->scroll_col;
  size_t cursor_row = coords.row;
  size_t cursor_col = coords.col;
  int cursor_screen_row = -1;
  int cursor_screen_col = -1;
  int alt_screen_row = 0;
  for (size_t screen_row = 0; screen_row < (size_t)buffer->Buffer.calc_rows;
       screen_row++) {
    size_t line_idx = scroll_row + alt_screen_row;
    alt_screen_row++;
    int orig_idx = -1;
    if (line_idx < total_lines) {
      orig_idx = 0;
      for (; orig_idx < (int)e->line_count; orig_idx++) {
        if (e->lines[orig_idx] == lines[line_idx])
          break;
      }
    }
    int gutter_width = 1;
    if (e->line_count > 0) {
      int max_idx = e->line_count;
      gutter_width = 0;
      while (max_idx > 0) {
        max_idx /= 10;
        gutter_width++;
      }
    }
    gutter_width++;
    char buf[12];
    if (orig_idx == -1)
      buf[0] = '\0';
    else
      snprintf(buf, sizeof(buf), "%d", orig_idx + 1);
    int len = strlen(buf);
    const Color *color = colors + DARK;
    if (line_idx == cursor_row)
      color = colors + WHITE;
    for (int j = 0; j < gutter_width; j++) {
      char tmp[2];
      if (j < gutter_width - len)
        tmp[0] = ' ';
      else
        tmp[0] = buf[j - (gutter_width - len)];
      tmp[1] = '\0';
      insert_screen(buffer->Buffer.calc_y + screen_row,
                    buffer->Buffer.calc_x + j - gutter_width - 1, tmp, color,
                    buffer->Buffer.bg, 0);
    }
    char tmp[2];
    tmp[0] = ' ';
    tmp[1] = '\0';
    insert_screen(buffer->Buffer.calc_y + screen_row, buffer->Buffer.calc_x - 1,
                  tmp, buffer->Buffer.fg, buffer->Buffer.bg, 0);
    if (line_idx >= total_lines) {
      fill_blank_line(buffer, buffer->Buffer.calc_y, screen_row,
                      buffer->Buffer.calc_cols);
      continue;
    }
    Line *line = lines[line_idx];
    Line *virtual = NULL;
    if (virt_lines[line_idx] && virt_lines[line_idx]->line_count > 0)
      virtual = virt_lines[line_idx]->lines[0];
    Line *line_f = merge_virtual_line(line, virtual);
    int screen_col =
        draw_line_clusters(buffer, line_f, line_idx, cursor_row, cursor_col,
                           scroll_col, screen_row, buffer->Buffer.calc_cols,
                           &cursor_screen_row, &cursor_screen_col);
    for (size_t i = 0; i < line_f->cluster_count; i++)
      free(line_f->clusters[i].utf8);
    free(line_f->highlights);
    free(line_f->clusters);
    free(line_f);
    pad_remaining_cols(buffer, buffer->Buffer.calc_y, screen_row, screen_col,
                       buffer->Buffer.calc_cols);
    if (virt_lines[line_idx] && virt_lines[line_idx]->line_count > 1) {
      for (size_t i = 1; i < virt_lines[line_idx]->line_count; i++) {
        virtual = virt_lines[line_idx]->lines[i];
        screen_row++;
        draw_line_clusters(buffer, virtual, line_idx, cursor_row, cursor_col,
                           scroll_col, screen_row, buffer->Buffer.calc_cols,
                           &cursor_screen_row, &cursor_screen_col);
        pad_remaining_cols(buffer, buffer->Buffer.calc_y, screen_row, 0,
                           buffer->Buffer.calc_cols);
      }
    }
  }
  free(lines);
  for (size_t i = 0; i < total_lines; i++) {
    VirtLine *vl = virt_lines[i];
    if (!vl)
      continue;
    if (vl->is_tmp) {
      for (size_t j = 0; j < vl->line_count; j++) {
        Line *line = vl->lines[j];
        if (!line)
          continue;
        for (size_t k = 0; k < line->cluster_count; k++)
          free(line->clusters[k].utf8);
        free(line->clusters);
        free(line->highlights);
        free(line);
      }
      free(vl->lines);
      free(vl);
    }
  }
  free(virt_lines);
  if (cursor_screen_row >= 0 && cursor_screen_col >= 0 &&
      root.Root.focused == buffer)
    printf("\x1b[%d;%dH", cursor_screen_row + 1, cursor_screen_col + 1);
}

Node *create_bubble(int row, int col, int width, int height, char *text) {
  Node *node = calloc(1, sizeof(Node));
  node->type = BUBBLE;
  node->Bubble.row = row;
  node->Bubble.col = col;
  node->Bubble.width = width;
  node->Bubble.height = height;
  node->Bubble.lines = calloc(height, sizeof(Line));
  char *copy = strdup(text);
  if (!copy)
    return node;
  size_t line_idx = 0;
  char *line = strtok(copy, "\n");
  while (line && line_idx < (size_t)height) {
    size_t len = strlen(line);
    if (len > (size_t)width)
      len = width;
    Line *dst = &node->Bubble.lines[line_idx];
    *dst = split_graphemes(line);
    line_idx++;
    line = strtok(NULL, "\n");
  }
  node->Bubble.line_count = line_idx;
  free(copy);
  return node;
}

void draw_bubble(Node *node) {
  if (!node || node->type != BUBBLE)
    return;
  Line *lines = node->Bubble.lines;
  insert_screen(node->Bubble.row, node->Bubble.col, "╭", colors + BLUE, colors,
                0);
  for (size_t j = 1; j < node->Bubble.width; j++) {
    insert_screen(node->Bubble.row, node->Bubble.col + j, "─", colors + BLUE,
                  colors, 0);
    insert_screen(node->Bubble.row + node->Bubble.height, node->Bubble.col + j,
                  "─", colors + BLUE, colors, 0);
  }
  insert_screen(node->Bubble.row, node->Bubble.col + node->Bubble.width, "╮",
                colors + BLUE, colors, 0);
  insert_screen(node->Bubble.row + node->Bubble.height, node->Bubble.col, "╰",
                colors + BLUE, colors, 0);
  for (size_t i = 1; i < node->Bubble.height; i++) {
    insert_screen(node->Bubble.row + i, node->Bubble.col, "│", colors + BLUE,
                  colors, 0);
    insert_screen(node->Bubble.row + i, node->Bubble.col + node->Bubble.width,
                  "│", colors + BLUE, colors, 0);
  }
  insert_screen(node->Bubble.row + node->Bubble.height,
                node->Bubble.col + node->Bubble.width, "╯", colors + BLUE,
                colors, 0);
  for (size_t i = 0; i < node->Bubble.height - 1; i++) {
    Line line = lines[i];
    for (size_t j = 0; j < node->Bubble.width - 1; j++) {
      if (j >= line.cluster_count) {
        insert_screen(node->Bubble.row + i + 1, node->Bubble.col + j + 1, " ",
                      colors + BLUE, colors, 0);
        continue;
      }
      Cluster c = line.clusters[j];
      insert_screen(node->Bubble.row + i + 1, node->Bubble.col + j + 1, c.utf8,
                    colors + WHITE, colors, 0);
    }
  }
}

Node *create_block(Node *parent_block, _Bool direction, uint8_t size) {
  if (parent_block->type != BLOCK)
    die("parent is not a block");
  Node *block = calloc(1, sizeof(Node));
  if (!block)
    die("calloc failed");
  block->type = BLOCK;
  block->parent = parent_block;
  block->Block.direction = direction;
  block->Block.size = size;
  Node **tmp = realloc(parent_block->Block.children,
                       sizeof(Node *) * (parent_block->Block.count + 1));
  if (!tmp)
    die("realloc failed");
  parent_block->Block.children = tmp;
  parent_block->Block.children[parent_block->Block.count] = block;
  parent_block->Block.count++;
  return block;
}

void calculate_block(Node *block) {
  if (!block || block->type != BLOCK)
    return;
  uint8_t total_size = (block->Block.direction == HORIZONTAL)
                           ? block->Block.calc_cols
                           : block->Block.calc_rows;
  uint8_t used_space = 0;
  int unsized_count = 0;
  for (int i = 0; i < block->Block.count; i++) {
    Node *child = block->Block.children[i];
    uint8_t size = 0;
    if (child->type == BUFFER)
      size = child->Buffer.size;
    else if (child->type == BLOCK)
      size = child->Block.size;
    if (size && size > 0)
      used_space += size;
    else
      unsized_count++;
  }
  uint8_t remaining_space =
      (total_size > used_space) ? total_size - used_space : 0;
  uint8_t default_size =
      (unsized_count > 0) ? remaining_space / unsized_count : 0;
  uint8_t leftover = (unsized_count > 0) ? remaining_space % unsized_count : 0;
  uint8_t cur_x = block->Block.calc_x;
  uint8_t cur_y = block->Block.calc_y;
  for (int i = 0; i < block->Block.count; i++) {
    Node *child = block->Block.children[i];
    uint8_t child_size = 0;
    if (child->type == BUFFER) {
      child_size = (child->Buffer.size > 0) ? child->Buffer.size : default_size;
      if (child->Buffer.size == 0 && leftover > 0) {
        child_size++;
        leftover--;
      }
      child->Buffer.calc_x = cur_x;
      child->Buffer.calc_y = cur_y;
      if (block->Block.direction == HORIZONTAL) {
        child->Buffer.calc_cols = child_size;
        child->Buffer.calc_rows = block->Block.calc_rows;
        cur_x += child_size;
      } else {
        child->Buffer.calc_cols = block->Block.calc_cols;
        child->Buffer.calc_rows = child_size;
        cur_y += child_size;
      }
      int n = child->Buffer.data->line_count;
      int count = 0;
      do {
        count++;
        n /= 10;
      } while (n != 0);
      child->Buffer.calc_cols -= count + 2;
      child->Buffer.calc_x += count + 2;
    } else if (child->type == BLOCK) {
      child_size = (child->Block.size > 0) ? child->Block.size : default_size;
      if (child->Block.size == 0 && leftover > 0) {
        child_size++;
        leftover--;
      }
      child->Block.calc_x = cur_x;
      child->Block.calc_y = cur_y;
      if (block->Block.direction == HORIZONTAL) {
        child->Block.calc_cols = child_size;
        child->Block.calc_rows = block->Block.calc_rows;
        cur_x += child_size;
      } else {
        child->Block.calc_cols = block->Block.calc_cols;
        child->Block.calc_rows = child_size;
        cur_y += child_size;
      }
      calculate_block(child);
    }
  }
}

static void free_line(Line *line) {
  if (!line)
    return;
  for (size_t j = 0; j < line->cluster_count; j++)
    free(line->clusters[j].utf8);
  free(line->clusters);
  free(line->highlights);
  free(line);
}

static void free_virtual_line(VirtLine *vline) {
  if (!vline)
    return;
  for (size_t i = 0; i < vline->line_count; i++)
    free_line(vline->lines[i]);
  free(vline->lines);
  free(vline);
}

static void free_buffer(BufferData *e) {
  if (!e)
    return;
  if (e->lines) {
    for (size_t i = 0; i < e->line_count; i++)
      free_line(e->lines[i]);
    free(e->lines);
  }
  if (e->virtual_lines) {
    for (size_t i = 0; i < e->line_count; i++)
      free_virtual_line(e->virtual_lines[i]);
    free(e->virtual_lines);
  }
  if (e->tree)
    ts_tree_delete(e->tree);
  if (e->parser)
    ts_parser_delete(e->parser);
  free(e->folds);
  free(e);
}

void free_all_nodes(Node *node) {
  if (!node)
    return;
  switch (node->type) {
  case ROOT:
    for (size_t i = 0; i < node->Root.count; i++)
      free_all_nodes(node->Root.windows[i]);
    free(node->Root.windows);
    break;
  case WINDOW:
    if (node->Window.root_block_node)
      free_all_nodes(node->Window.root_block_node);
    break;
  case BLOCK:
    for (size_t i = 0; i < node->Block.count; i++)
      free_all_nodes(node->Block.children[i]);
    free(node->Block.children);
    break;
  case BUFFER:
    free_buffer(node->Buffer.data);
    free(node->Buffer.folder);
    break;
  }
  if (node != &root)
    free(node);
}
