#include "all.h"
#include <stdlib.h>

static inline lex_tok_e peek(parser_t *p) {
  return iv_get(&p->lx->tokens, p->pos);
}

static inline lex_tok_e advance(parser_t *p) {
  return iv_get(&p->lx->tokens, p->pos++);
}

static ast_node_t *new_node(ast_type_e type, int val) {
  ast_node_t *n = malloc(sizeof(ast_node_t));
  n->type = type;
  n->val = val;
  n->next = NULL;
  n->children = (node_vec){0};
  return n;
}

static ast_node_t *parse_block(parser_t *p);

static ast_node_t *parse_single(parser_t *p) {
  lex_tok_e t = peek(p);

  switch (t) {

  case LEX_PINC: {
    int count = 0;
    while (peek(p) == LEX_PINC) {
      advance(p);
      count++;
    }
    return new_node(AST_INC, count);
  }

  case LEX_PDEC: {
    int count = 0;
    while (peek(p) == LEX_PDEC) {
      advance(p);
      count++;
    }
    return new_node(AST_DEC, count);
  }

  case LEX_VINC: {
    int count = 0;
    while (peek(p) == LEX_VINC) {
      advance(p);
      count++;
    }
    return new_node(AST_PTR_INC, count);
  }

  case LEX_VDEC: {
    int count = 0;
    while (peek(p) == LEX_VDEC) {
      advance(p);
      count++;
    }
    return new_node(AST_PTR_DEC, count);
  }

  case LEX_OUT:
    advance(p);
    return new_node(AST_OUT, 1);

  case LEX_IN:
    advance(p);
    return new_node(AST_IN, 1);

  case LEX_LB:
    advance(p); // consume '['
    return parse_block(p);

  case LEX_RB:
    return NULL;

  case LEX_EOF:
    return NULL;

  default:
    advance(p);
    return NULL;
  }
}

static ast_node_t *parse_block(parser_t *p) {
  ast_node_t *loop = new_node(AST_LOOP, 1);

  while (1) {
    lex_tok_e t = peek(p);

    if (t == LEX_EOF) {
      break;
    }

    if (t == LEX_RB) {
      advance(p); // consume ']'
      break;
    }

    ast_node_t *child = parse_single(p);
    if (child) {
      node_vec_push(&loop->children, child);
    }
  }

  return loop;
}

parser_t parser_init(lexer_t *lx) { return (parser_t){.lx = lx, .pos = 0}; }

ast_node_t *parse_program(parser_t *p) {
  ast_node_t *head = NULL;
  ast_node_t *tail = NULL;

  while (peek(p) != LEX_EOF) {

    ast_node_t *node = parse_single(p);
    if (!node)
      continue;

    if (!head) {
      head = tail = node;
    } else {
      tail->next = node;
      tail = node;
    }
  }

  return head;
}

void ast_free(ast_node_t *node) {
  if (!node)
    return;

  if (node->next)
    ast_free(node->next);

  for (size_t i = 0; i < node->children.len; i++) {
    ast_free(node->children._data[i]);
  }

  node_vec_free(&node->children);

  free(node);
}