#include "highlighters.h"
#include "../libs/tree-sitter-c/bindings/c/tree_sitter/tree-sitter-c.h"
#include "../libs/tree-sitter-ruby/bindings/c/tree-sitter-ruby.h"
#include "../libs/tree-sitter/lib/include/tree_sitter/api.h"
#include "constants.h"
#include "editor.h"
#include "helper_functions.h"
#include "highlighters.h"
#include "main.h"
#include "node.h"
#include "ruby_ts_table.h"
#include "snake.h"
#include <assert.h>
#include <ctype.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <termios.h>
#include <tree_sitter/api.h>
#include <unistd.h>

void tree_sitter_highlighter(Node *buffer) {
  BufferData *e = buffer->Buffer.data;
  TSParser *parser = e->parser;
  TSTree *tree = e->tree;
  if (!parser) {
    parser = ts_parser_new();
    if (!parser)
      die("ts_parser_new failed");
    switch (e->lang) {
    case RUBY:
      ts_parser_set_language(parser, tree_sitter_ruby());
      break;
    case C_LANG:
      ts_parser_set_language(parser, tree_sitter_c());
      break;
    default:
      die("unknown language");
    }
    e->parser = parser;
    e->tree = NULL;
  }
  char *code = lines_to_string(e);
  TSTree *old_tree = tree;
  tree = ts_parser_parse_string(parser, old_tree, code, strlen(code));
  TSNode t_root = ts_tree_root_node(tree);
  if (old_tree && old_tree != tree) {
    ts_tree_delete(old_tree);
  }
  if (!tree)
    die("ts_parser_parse_string failed");
  e->tree = tree;
  Fold *prev_folds = e->folds;
  int prev_fold_count = e->fold_count;
  e->folds = NULL;
  e->fold_count = 0;
  collect_tree_highlights(e, t_root, code, prev_folds, prev_fold_count);
  free(prev_folds);
  free(code);
}

char *ts_node_text(char *src, TSNode node) {
  size_t start = ts_node_start_byte(node);
  size_t end = ts_node_end_byte(node);
  size_t len = end - start;
  char *text = malloc(len + 1);
  memcpy(text, src + start, len);
  text[len] = '\0';
  return text;
}

void collect_tree_highlights(BufferData *e, TSNode node, char *src,
                             Fold *prev_folds, int prev_fold_count) {
  TSTreeCursor cursor = ts_tree_cursor_new(node);
  if (ts_tree_cursor_goto_first_child(&cursor)) {
    do {
      TSNode current = ts_tree_cursor_current_node(&cursor);
      TSPoint start = ts_node_start_point(current);
      TSPoint end = ts_node_end_point(current);
      if (start.row != end.row) {
        e->folds = realloc(e->folds, (e->fold_count + 1) * sizeof(Fold));
        e->folds[e->fold_count].start = start.row;
        e->folds[e->fold_count].end = end.row;
        char *text = ts_node_text(src, current);
        uint32_t hash = fnv1a_32(text, strlen(text));
        e->folds[e->fold_count].hash = hash;
        bool folded = false;
        if (prev_folds) {
          for (size_t i = 0; i < prev_fold_count; i++) {
            if (prev_folds[i].hash == hash) {
              folded = prev_folds[i].folded;
              break;
            }
          }
        }
        e->folds[e->fold_count].folded = folded;
        free(text);
        e->fold_count++;
      }
      if (ts_node_is_named(current)) {
        highlight_node(e, current);
        collect_tree_highlights(e, current, src, prev_folds, prev_fold_count);
      }
    } while (ts_tree_cursor_goto_next_sibling(&cursor));
  }
  ts_tree_cursor_delete(&cursor);
}

void highlight_node(BufferData *e, TSNode current) {
  TSPoint start = ts_node_start_point(current);
  TSPoint end = ts_node_end_point(current);
  TSSymbol symbol = ts_node_symbol(current);
  if (start.row >= e->line_count)
    return;
  const Color *fg = NULL;
  uint8_t flags = 0;
  symbol_colorify(symbol, &fg, &flags);
  if (!fg && !flags)
    return;
  for (size_t row = start.row; row <= end.row && row < e->line_count; row++) {
    Line *line = e->lines[row];
    size_t col_start = (row == start.row) ? start.column : 0;
    size_t col_end =
        (row == end.row) ? end.column : e->lines[row]->cluster_count;
    append_highlight_byte(line, col_start, col_end, (Color *)fg, NULL, flags);
  }
}

void append_highlight_byte(Line *line, size_t start_col, size_t end_col,
                           Color *fg, Color *bg, uint8_t flags) {
  size_t start = byte_to_grapheme(line, start_col);
  size_t end = byte_to_grapheme(line, end_col);
  line->highlight_count++;
  line->highlights =
      realloc(line->highlights, line->highlight_count * sizeof(Highlight));
  Highlight *h = &line->highlights[line->highlight_count - 1];
  h->fg = fg;
  h->bg = bg;
  h->flags = flags;
  h->start_col = start;
  h->end_col = end;
}

static int is_word_char(char c) {
  return isalnum((unsigned char)c) || c == '_';
}

void append_highlight(Line *lh, size_t start_col, size_t end_col, Color *fg,
                      Color *bg, uint8_t flags) {
  lh->highlight_count++;
  lh->highlights =
      realloc(lh->highlights, lh->highlight_count * sizeof(Highlight));
  Highlight *h = &lh->highlights[lh->highlight_count - 1];
  h->fg = fg;
  h->bg = bg;
  h->flags = flags;
  h->start_col = start_col;
  h->end_col = end_col;
}

void clear_virt(BufferData *e) {
  for (size_t i = 0; i < e->line_count; i++) {
    VirtLine *vl = e->virtual_lines[i];
    if (!vl)
      continue;
    if (vl->lines) {
      for (size_t j = 0; j < vl->line_count; j++) {
        Line *inner = vl->lines[j];
        if (!inner)
          continue;
        for (size_t k = 0; k < inner->cluster_count; k++) {
          if (inner->clusters && inner->clusters[k].utf8)
            free(inner->clusters[k].utf8);
        }
        if (inner->clusters)
          free(inner->clusters);
        if (inner->highlights)
          free(inner->highlights);
        free(inner);
      }
      free(vl->lines);
      vl->line_count = 0;
    }
  }
}

void word_highlighter(Node *buffer) {
  BufferData *e = buffer->Buffer.data;
  if (buffer->Buffer.data->selecting) {
    int start_row = buffer->Buffer.data->selection_start_row;
    int start_col = buffer->Buffer.data->selection_start_col;
    int end_row = e->cursor_row;
    int end_col = e->cursor_col;
    if (start_row > end_row || (start_row == end_row && start_col > end_col)) {
      int tmp_row = start_row, tmp_col = start_col;
      start_row = end_row;
      start_col = end_col;
      end_row = tmp_row;
      end_col = tmp_col;
    }
    for (int row = start_row; row <= end_row; row++) {
      Line *l = e->lines[row];
      if (!l)
        continue;
      int s = 0;
      int e_col = (int)l->cluster_count;
      if (row == start_row)
        s = start_col;
      if (row == end_row)
        e_col = end_col;
      if (s < 0)
        s = 0;
      if (e_col > (int)l->cluster_count)
        e_col = (int)l->cluster_count;
      if (s < e_col)
        append_highlight(l, s, e_col + 1, NULL, (Color *)colors + DARK, 0b000);
    }
  }
  clear_virt(e);
  for (size_t x = 0; x < e->line_count; x++) {
    Line *l = e->lines[x];
    VirtLine *v_l = e->virtual_lines[x];
    if (!l || !v_l)
      continue;
    if (v_l->line_count == 0) {
      Line *tmp = calloc(1, sizeof(Line));
      if (!tmp)
        die("calloc failed");
      tmp->cluster_count = l->cluster_count;
      tmp->clusters = calloc(l->cluster_count, sizeof(Cluster));
      if (!tmp->clusters)
        die("calloc failed");
      v_l->lines = calloc(1, sizeof(Line *));
      if (!v_l->lines)
        die("calloc failed");
      v_l->lines[0] = tmp;
      v_l->line_count = 1;
    }
    Line *virt = v_l->lines[0];
    if (!virt->clusters) {
      virt->cluster_count = l->cluster_count;
      virt->clusters = calloc(l->cluster_count, sizeof(Cluster));
    }
    for (size_t y = 0; y < l->cluster_count; y++) {
      if (l->clusters[y].utf8[0] != ' ')
        break;
      virt->clusters[y].utf8 = strdup("·");
      virt->clusters[y].width = 1;
      append_highlight(virt, y, y + 1, (Color *)colors + DARK, NULL, 0b010);
    }
  }
  if (!e || e->line_count == 0)
    return;
  Line *line = e->lines[e->cursor_row];
  if (!line || e->cursor_col >= line->cluster_count)
    return;
  if (!is_word_char(line->clusters[e->cursor_col].utf8[0]))
    return;
  size_t start = e->cursor_col;
  size_t end = e->cursor_col;
  while (start > 0 && is_word_char(line->clusters[start - 1].utf8[0]))
    start--;
  while (end < line->cluster_count && is_word_char(line->clusters[end].utf8[0]))
    end++;
  if (start == end)
    return;
  size_t word_len = end - start;
  char **word_clusters = malloc(word_len * sizeof(char *));
  if (!word_clusters)
    return;
  for (size_t i = 0; i < word_len; i++)
    word_clusters[i] = line->clusters[start + i].utf8;
  for (size_t row = 0; row < e->line_count; row++) {
    Line *l = e->lines[row];
    if (!l)
      continue;
    for (size_t i = 0; i + word_len <= l->cluster_count; i++) {
      bool match = true;
      for (size_t j = 0; j < word_len; j++) {
        if (strcmp(l->clusters[i + j].utf8, word_clusters[j]) != 0) {
          match = false;
          break;
        }
      }
      if (match) {
        append_highlight(l, i, i + word_len, NULL, NULL, 0b100);
        i += word_len - 1;
      }
    }
  }
  free(word_clusters);
}
