#include "tree_sitter/runtime.h"
#include "runtime/alloc.h"
#include "runtime/tree_cursor.h"
#include "runtime/language.h"
#include "runtime/tree.h"

TSTreeCursor ts_tree_cursor_new(const TSTree *tree) {
  TSTreeCursor self;
  ts_tree_cursor_init((TreeCursor *)&self, tree);
  return self;
}

void ts_tree_cursor_init(TreeCursor *self, const TSTree *tree) {
  self->tree = tree;
  array_init(&self->stack);
  array_push(&self->stack, ((TreeCursorEntry) {
    .subtree = tree->root,
    .position = length_zero(),
    .child_index = 0,
    .structural_child_index = 0,
  }));
}

void ts_tree_cursor_delete(TSTreeCursor *_self) {
  TreeCursor *self = (TreeCursor *)_self;
  array_delete(&self->stack);
}

bool ts_tree_cursor_goto_first_child(TSTreeCursor *_self) {
  TreeCursor *self = (TreeCursor *)_self;
  TreeCursorEntry *last_entry = array_back(&self->stack);
  const Subtree *subtree = last_entry->subtree;
  if (subtree->children.size == 0 || subtree->visible_child_count == 0) return false;
  Length position = last_entry->position;

  bool did_descend;
  do {
    did_descend = false;

    const TSSymbol *alias_sequence = ts_language_alias_sequence(
      self->tree->language,
      subtree->alias_sequence_id
    );
    uint32_t structural_child_index = 0;
    for (uint32_t i = 0; i < subtree->children.size; i++) {
      const Subtree *child = subtree->children.contents[i];
      bool visible = child->visible;
      if (!child->extra) {
        visible |= alias_sequence && alias_sequence[structural_child_index];
      }

      if (visible || (child->children.size > 0 && child->visible_child_count > 0)) {
        array_push(&self->stack, ((TreeCursorEntry) {
          .subtree = child,
          .child_index = i,
          .structural_child_index = structural_child_index,
          .position = position,
        }));

        if (visible) {
          return true;
        } else {
          subtree = child;
          did_descend = true;
          break;
        }
      }
      if (!child->extra) structural_child_index++;
      position = length_add(position, ts_subtree_total_size(child));
    }
  } while (did_descend);

  return false;
}

int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *_self, uint32_t goal_byte) {
  TreeCursor *self = (TreeCursor *)_self;
  uint32_t initial_size = self->stack.size;
  TreeCursorEntry *last_entry = array_back(&self->stack);
  const Subtree *subtree = last_entry->subtree;
  Length position = last_entry->position;
  uint32_t visible_child_index = 0;

  bool did_descend;
  do {
    did_descend = false;

    const TSSymbol *alias_sequence = ts_language_alias_sequence(
      self->tree->language,
      subtree->alias_sequence_id
    );
    uint32_t structural_child_index = 0;
    for (uint32_t i = 0; i < subtree->children.size; i++) {
      const Subtree *child = subtree->children.contents[i];
      Length next_position = length_add(position, ts_subtree_total_size(child));
      bool at_goal = next_position.bytes > goal_byte;

      bool visible = child->visible;
      if (!child->extra) {
        visible |= alias_sequence && alias_sequence[structural_child_index];
      }

      uint32_t visible_child_count = child->children.size > 0 ? child->visible_child_count : 0;

      if (at_goal) {
        if (visible || visible_child_count > 0) {
          array_push(&self->stack, ((TreeCursorEntry) {
            .subtree = child,
            .child_index = i,
            .structural_child_index = structural_child_index,
            .position = position,
          }));

          if (visible) {
            return visible_child_index;
          } else {
            subtree = child;
            did_descend = true;
            break;
          }
        }
      } else {
        if (visible) {
          visible_child_index++;
        } else {
          visible_child_index += visible_child_count;
        }
      }

      if (!child->extra) structural_child_index++;
      position = next_position;
    }
  } while (did_descend);

  self->stack.size = initial_size;
  return -1;
}

bool ts_tree_cursor_goto_next_sibling(TSTreeCursor *_self) {
  TreeCursor *self = (TreeCursor *)_self;
  TreeCursorEntry *child_entry = array_back(&self->stack);

  for (unsigned i = self->stack.size - 2; i + 1 > 0; i--) {
    TreeCursorEntry *parent_entry = &self->stack.contents[i];

    const Subtree *parent = parent_entry->subtree;
    uint32_t child_index = child_entry->child_index;
    uint32_t structural_child_index = child_entry->structural_child_index;
    Length position = child_entry->position;
    const Subtree *child = parent->children.contents[child_index];
    const TSSymbol *alias_sequence = ts_language_alias_sequence(
      self->tree->language,
      parent->alias_sequence_id
    );

    while (++child_index < parent->children.size) {
      if (!child->extra) structural_child_index++;
      position = length_add(position, ts_subtree_total_size(child));
      child = parent->children.contents[child_index];

      bool visible = child->visible;
      if (!child->extra) {
        visible |= alias_sequence && alias_sequence[structural_child_index];
      }

      if (visible || (child->children.size > 0 && child->visible_child_count > 0)) {
        self->stack.contents[i + 1] = (TreeCursorEntry) {
          .subtree = child,
          .child_index = child_index,
          .structural_child_index = structural_child_index,
          .position = position,
        };
        self->stack.size = i + 2;

        if (visible) {
          return true;
        } else {
          ts_tree_cursor_goto_first_child(_self);
          return true;
        }
      }
    }

    child_entry = parent_entry;
    if (parent->visible) break;
  }

  return false;
}

bool ts_tree_cursor_goto_parent(TSTreeCursor *_self) {
  TreeCursor *self = (TreeCursor *)_self;
  for (unsigned i = self->stack.size - 2; i + 1 > 0; i--) {
    TreeCursorEntry *entry = &self->stack.contents[i];
    if (entry->subtree->visible) {
      self->stack.size = i + 1;
      return true;
    }
  }
  return false;
}

TSNode ts_tree_cursor_current_node(const TSTreeCursor *_self) {
  const TreeCursor *self = (const TreeCursor *)_self;
  TreeCursorEntry *last_entry = array_back(&self->stack);
  TSSymbol alias_symbol = 0;
  if (self->stack.size > 1) {
    TreeCursorEntry *parent_entry = &self->stack.contents[self->stack.size - 2];
    const TSSymbol *alias_sequence = ts_language_alias_sequence(
      self->tree->language,
      parent_entry->subtree->alias_sequence_id
    );
    if (alias_sequence && !last_entry->subtree->extra) {
      alias_symbol = alias_sequence[last_entry->structural_child_index];
    }
  }
  return ts_node_new(
    self->tree,
    last_entry->subtree,
    length_add(last_entry->position, last_entry->subtree->padding),
    alias_symbol
  );
}