#include "compiler/fzl_token.h"
#include <compiler/fzl_parser.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

char plast_err[256];

fzl_parser_t *fzl_mkparser(fzl_reader_ll_t *first)
{
  if (!first)
  {
    snprintf(plast_err, sizeof(plast_err), "token is null");
    return NULL;
  }

  fzl_parser_t *parser = malloc(sizeof(fzl_parser_t));
  if (!parser)
    return NULL;

  parser->root   = NULL;
  parser->cursor = first->next;
  if (!parser->cursor)
    return NULL;

  return parser;
}

// --- Parse zone ---

// Make a new node
fzl_node_t *fzl_mknode(fzl_nodetype_t type, fzl_token_t *token)
{
  fzl_node_t *node = malloc(sizeof(fzl_node_t));

  if (!node)
    return NULL;
  node->type   = type;
  node->token  = token;
  node->parent = NULL;
  node->next   = NULL;

  return node;
}

// Get token at cursor
static fzl_token_t *peek(fzl_parser_t *parser)
{
  if (!parser || !parser->cursor)
    return NULL;
  return parser->cursor->token;
}

// Get token at (cursor + 1)
static fzl_token_t *next(fzl_parser_t *parser)
{
  if (!parser || !parser->cursor)
    return NULL;

  fzl_reader_ll_t *next = parser->cursor->next;
  if (!next)
    return NULL;

  return next->token;
}

// Consume token at cursor and move to the next
static fzl_token_t *consume(fzl_parser_t *parser)
{
  if (!parser || !parser->cursor)
    return NULL;
  fzl_token_t *tok = parser->cursor->token;
  parser->cursor   = parser->cursor->next;
  return tok;
}

// Check if current token have the same type as specified
static bool is(fzl_parser_t *parser, fzl_toktype_t type)
{
  if (!parser)
    return false;

  fzl_token_t *tok = peek(parser);
  if (!tok)
    return false;

  return tok->type == type;
}

// Convert arithmetic operator token type to node type
static fzl_nodetype_t math_op_to_type(fzl_toktype_t type)
{
  switch (type)
  {
    case FZL_TOKTYPE_PLUS:
      return FZL_NODETYPE_ADD;
    case FZL_TOKTYPE_MINUS:
      return FZL_NODETYPE_SUB;
    case FZL_TOKTYPE_STAR:
      return FZL_NODETYPE_MUL;
    case FZL_TOKTYPE_SLASH:
      return FZL_NODETYPE_DIV;
    case FZL_TOKTYPE_CARET:
      return FZL_NODETYPE_POW;
    case FZL_TOKTYPE_PERCENT:
      return FZL_NODETYPE_MOD;

    default:
      return 0;
  }
}

static fzl_nodetype_t unary_op_to_type(fzl_toktype_t type)
{
  switch (type)
  {
    case FZL_TOKTYPE_PLUS:
      return FZL_NODETYPE_UNARY_PLUS;
    case FZL_TOKTYPE_MINUS:
      return FZL_NODETYPE_UNARY_MINUS;
    case FZL_TOKTYPE_NOT:
      return FZL_NODETYPE_NOT;

    default:
      return 0;
  }
}

fzl_node_t *parse_expression(fzl_parser_t *parser);

// Parse primary
fzl_node_t *parse_primary(fzl_parser_t *parser)
{
  if (!parser || !peek(parser))
    return NULL;

  if (is(parser, FZL_TOKTYPE_STR_LITERAL))
  {
    fzl_token_t *tok = consume(parser);
    return fzl_mknode(FZL_NODETYPE_LITERAL, tok);
  }
  else if (is(parser, FZL_TOKTYPE_NUM_LITERAL))
  {
    fzl_token_t *tok = consume(parser);
    return fzl_mknode(FZL_NODETYPE_LITERAL, tok);
  }
  else if (is(parser, FZL_TOKTYPE_BOOL_LITERAL))
  {
    fzl_token_t *tok = consume(parser);
    return fzl_mknode(FZL_NODETYPE_LITERAL, tok);
  }
  else if (is(parser, FZL_TOKTYPE_LPAREN))
  {
    consume(parser);
    fzl_node_t *node = parse_expression(parser);
    if (!is(parser, FZL_TOKTYPE_RPAREN))
    {
      snprintf(plast_err, sizeof(plast_err), "expected ')' at %ld:%ld",
               parser->cursor->token->ln, parser->cursor->token->cn);
      return NULL;
    }
    consume(parser);
    return node;
  }
  else if (is(parser, FZL_TOKTYPE_NAME))
  {
    fzl_token_t *name = consume(parser);
    return fzl_mknode(FZL_NODETYPE_IDENTIFIER, name);
  }

  snprintf(plast_err, sizeof(plast_err), "failed to parse primary at %ld:%ld",
           peek(parser)->ln, peek(parser)->cn);
  return NULL;
}

// Parse postfix expressions (parens, index access, dots)
fzl_node_t *parse_postfix(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_primary(parser);
  fzl_node_t *right = NULL;

  if (!left)
  {
    snprintf(plast_err, sizeof(plast_err), "failed to parse postfix at %ld:%ld",
             peek(parser)->ln, peek(parser)->cn);
    return NULL;
  }

  while (is(parser, FZL_TOKTYPE_LPAREN))
  {
    consume(parser);

    size_t args_capacity = 2;
    size_t args_count    = 0;
    fzl_node_t **args    = malloc(args_capacity * sizeof(fzl_node_t *));

    while (!is(parser, FZL_TOKTYPE_RPAREN))
    {
      fzl_node_t *arg = parse_expression(parser);
      if (!arg)
      {
        free(args);
        return NULL;
      }

      if (args_capacity <= args_count)
      {
        args_capacity *= 2;
        args = realloc(args, args_capacity * sizeof(fzl_node_t));
      }

      args[args_count] = arg;
      args_count++;

      if (is(parser, FZL_TOKTYPE_COMMA))
        consume(parser);
    }

    fzl_token_t *rparen = consume(parser);
    if (!rparen || rparen->type != FZL_TOKTYPE_RPAREN)
    {
      free(args);
      snprintf(plast_err, sizeof(plast_err),
               "expected ')' after argument list at %ld:%ld", peek(parser)->ln,
               peek(parser)->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(FZL_NODETYPE_CALL, left->token);
    node->call.callee     = left;
    node->call.args       = args;
    node->call.args_count = args_count;

    left = node;
  }

  return left;
}

// Parse unary operations
fzl_node_t *parse_unary(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  if (!is(parser, FZL_TOKTYPE_MINUS) && !is(parser, FZL_TOKTYPE_PLUS) &&
      !is(parser, FZL_TOKTYPE_NOT))
    return parse_postfix(parser);

  fzl_token_t *op     = consume(parser);
  fzl_node_t *operand = parse_unary(parser);

  if (!operand)
    return NULL;

  fzl_node_t *result       = fzl_mknode(unary_op_to_type(op->type), op);
  result->unary_op.operand = operand;

  return result;
}

fzl_node_t *parse_power(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_unary(parser);
  fzl_node_t *right = NULL;

  if (!left)
  {
    snprintf(plast_err, sizeof(plast_err), "failed to parse power at %ld:%ld",
             peek(parser)->ln, peek(parser)->cn);
    return NULL;
  }

  if (is(parser, FZL_TOKTYPE_CARET))
  {
    fzl_token_t *caret = consume(parser);

    right = parse_unary(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err),
               "failed to parse right side of power at %ld:%ld",
               peek(parser)->ln, peek(parser)->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(FZL_NODETYPE_POW, caret);
    node->binary_op.left  = left;
    node->binary_op.right = right;

    left = node;
  }

  return left;
}

// Parse FACTOR: multiplication and division
fzl_node_t *parse_factor(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_power(parser);
  fzl_node_t *right = NULL;

  if (!left)
  {
    if (!plast_err[0])
      snprintf(plast_err, sizeof(plast_err), "invalid expression at %ld:%ld",
               parser->cursor->token->ln, parser->cursor->token->cn);
    return NULL;
  }

  while (is(parser, FZL_TOKTYPE_STAR) || is(parser, FZL_TOKTYPE_SLASH) ||
         is(parser, FZL_TOKTYPE_PERCENT))
  {
    fzl_token_t *op = consume(parser);

    right = parse_factor(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err),
               "unfinished factor (mul/div) expression at %ld:%ld",
               left->token->ln, left->token->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(math_op_to_type(op->type), op);
    node->binary_op.left  = left;
    node->binary_op.right = right;
    left                  = node;
  }

  return left;
}

// Parse TERM: addition and subtraction
fzl_node_t *parse_term(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_factor(parser);
  fzl_node_t *right = NULL;

  if (!left)
  {
    snprintf(plast_err, sizeof(plast_err), "invalid expression at %ld:%ld",
             parser->cursor->token->ln, parser->cursor->token->cn);
    return NULL;
  }

  while (is(parser, FZL_TOKTYPE_PLUS) || is(parser, FZL_TOKTYPE_MINUS))
  {
    fzl_token_t *op = consume(parser);

    right = parse_factor(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err),
               "unfinished term (add/sub) expression at %ld:%ld",
               left->token->ln, left->token->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(math_op_to_type(op->type), op);
    node->binary_op.left  = left;
    node->binary_op.right = right;
    left                  = node;
  }

  return left;
}

// Parse conditional operations
fzl_node_t *parse_cond(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_term(parser);
  fzl_node_t *right = NULL;

  if (!left)
    return NULL;

  while (is(parser, FZL_TOKTYPE_EQEQ) || is(parser, FZL_TOKTYPE_NEQ) ||
         is(parser, FZL_TOKTYPE_LANGLE) || is(parser, FZL_TOKTYPE_LEQ) ||
         is(parser, FZL_TOKTYPE_RANGLE) || is(parser, FZL_TOKTYPE_REQ))
  {
    fzl_token_t *op = consume(parser);
    right           = parse_term(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err), "unfinished expression at %ld:%ld",
               op->ln, op->cn);
      return NULL;
    }

    fzl_node_t *result;
    switch (op->type)
    {
      case FZL_TOKTYPE_EQEQ:
        result = fzl_mknode(FZL_NODETYPE_EQUAL, op);
        break;
      case FZL_TOKTYPE_NEQ:
        result = fzl_mknode(FZL_NODETYPE_NEQUAL, op);
        break;
      case FZL_TOKTYPE_LANGLE:
        result = fzl_mknode(FZL_NODETYPE_LESS, op);
        break;
      case FZL_TOKTYPE_LEQ:
        result = fzl_mknode(FZL_NODETYPE_LESS_EQUAL, op);
        break;
      case FZL_TOKTYPE_RANGLE:
        result = fzl_mknode(FZL_NODETYPE_GREATER, op);
        break;
      case FZL_TOKTYPE_REQ:
        result = fzl_mknode(FZL_NODETYPE_GREATER_EQUAL, op);
        break;
      default:
        return NULL;
    }

    result->binary_op.left  = left;
    result->binary_op.right = right;
    left                    = result;
  }

  return left;
}

// Parse conditional &&
fzl_node_t *parse_and(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_cond(parser);
  fzl_node_t *right = NULL;

  if (!left)
    return NULL;

  while (is(parser, FZL_TOKTYPE_DAMPERSAND))
  {
    fzl_token_t * or = consume(parser);
    right            = parse_cond(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err), "unfinished expression at %ld:%ld",
               or->ln, or->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(FZL_NODETYPE_AND, or);
    node->binary_op.left  = left;
    node->binary_op.right = right;
    left                  = node;
  }

  return left;
}

// Parse conditional ||
fzl_node_t *parse_or(fzl_parser_t *parser)
{
  if (!parser)
    return NULL;

  fzl_node_t *left  = parse_and(parser);
  fzl_node_t *right = NULL;

  if (!left)
    return NULL;

  while (is(parser, FZL_TOKTYPE_DVBAR))
  {
    fzl_token_t * or = consume(parser);
    right            = parse_and(parser);
    if (!right)
    {
      snprintf(plast_err, sizeof(plast_err), "unfinished expression at %ld:%ld",
               or->ln, or->cn);
      return NULL;
    }

    fzl_node_t *node      = fzl_mknode(FZL_NODETYPE_OR, or);
    node->binary_op.left  = left;
    node->binary_op.right = right;
    left                  = node;
  }

  return left;
}

// Expression parsing
fzl_node_t *parse_expression(fzl_parser_t *parser) { return parse_and(parser); }

// Statement parsing
fzl_node_t *parse_statement(fzl_parser_t *parser)
{
  if (is(parser, FZL_TOKTYPE_TYPE))
  {
    fzl_token_t *type = peek(parser);
    if (type->type == FZL_TOKTYPE_TYPE)
      consume(parser);

    if (type->type == FZL_TOKTYPE_TYPE || next(parser)->type == FZL_TOKTYPE_EQ)
    {

      fzl_node_t *identifier = parse_primary(parser);
      if (!identifier || identifier->type != FZL_NODETYPE_IDENTIFIER)
      {
        snprintf(plast_err, sizeof(plast_err), "expected identifier at %ld:%ld",
                 type->ln, type->cn);
        return NULL;
      }

      fzl_node_t *left     = fzl_mknode(FZL_NODETYPE_DECLARATION, type);
      left->var.type       = type;
      left->var.identifier = identifier;

      if (!is(parser, FZL_TOKTYPE_EQ))
        return left;

      fzl_node_t *right;
      while (is(parser, FZL_TOKTYPE_EQ))
      {
        fzl_token_t *asn = consume(parser);
        right            = parse_expression(parser);
        if (!right)
          return NULL;

        fzl_node_t *node      = fzl_mknode(FZL_NODETYPE_ASSIGN, asn);
        node->binary_op.left  = left;
        node->binary_op.right = right;

        left = node;
      }

      return left;
    }
  }

  if (is(parser, FZL_TOKTYPE_NAME) && next(parser)->type == FZL_TOKTYPE_EQ)
  {
    fzl_token_t *name = consume(parser);
    fzl_node_t *left  = fzl_mknode(FZL_NODETYPE_IDENTIFIER, name);

    while (is(parser, FZL_TOKTYPE_EQ))
    {
      fzl_token_t *eq = consume(parser);

      fzl_node_t *right = parse_expression(parser);
      if (!right)
        return NULL;

      fzl_node_t *assn      = fzl_mknode(FZL_NODETYPE_ASSIGN, eq);
      assn->binary_op.left  = left;
      assn->binary_op.right = right;

      left = assn;
    }

    return left;
  }

  if (is(parser, FZL_TOKTYPE_IF))
  {
    fzl_token_t *if_tok = consume(parser);
    fzl_token_t *lparen = consume(parser);
    if (lparen->type != FZL_TOKTYPE_LPAREN)
    {
      snprintf(plast_err, sizeof(plast_err),
               "expected '(' before condition at %ld:%ld", lparen->ln,
               lparen->cn);
      return NULL;
    }

    fzl_node_t *condition = parse_expression(parser);
    if (!condition)
      return NULL;

    fzl_token_t *rparen = consume(parser);
    if (rparen->type != FZL_TOKTYPE_RPAREN)
    {
      snprintf(plast_err, sizeof(plast_err),
               "expected ')' after condition at %ld:%ld", rparen->ln,
               rparen->cn);
      return NULL;
    }

    fzl_node_t *else_block = NULL;
    fzl_node_t *true_block = parse_statement(parser);
    if (!true_block)
      return NULL;

    if (is(parser, FZL_TOKTYPE_ELSE))
    {
      consume(parser); // else
      else_block = parse_statement(parser);
      if (!else_block)
        return NULL;
    }

    fzl_node_t *stmt_node      = fzl_mknode(FZL_NODETYPE_CONDITIONAL, if_tok);
    stmt_node->cond.condition  = condition;
    stmt_node->cond.true_block = true_block;
    stmt_node->cond.else_block = else_block;

    return stmt_node;
  }

  if (is(parser, FZL_TOKTYPE_LBRACE))
  {
    fzl_token_t *block_tok = consume(parser);

    size_t remainder   = 2;
    size_t count       = 0;
    fzl_node_t **stmts = malloc(remainder * sizeof(fzl_node_t *));

    while (peek(parser) && !is(parser, FZL_TOKTYPE_RBRACE))
    {
      fzl_node_t *node = parse_statement(parser);
      if (!node)
        return NULL;

      if (remainder <= count)
      {
        remainder *= 2;
        stmts = realloc(stmts, remainder);
      }

      stmts[count++] = node;
    }
    fzl_token_t *rbrace = consume(parser);
    if (!rbrace || rbrace->type != FZL_TOKTYPE_RBRACE)
    {
      snprintf(plast_err, sizeof(plast_err),
               "missing '}' after block at %ld:%ld", rbrace->ln, rbrace->cn);
      return NULL;
    }

    fzl_node_t *node       = fzl_mknode(FZL_NODETYPE_BLOCK, block_tok);
    node->list.entries     = stmts;
    node->list.entry_count = count;
    return node;
  }

  return parse_expression(parser);
}

// Internal parsing entry point
fzl_node_t *parse_program(fzl_parser_t *parser)
{
  fzl_node_t *first = NULL;
  fzl_node_t *last  = NULL;
  fzl_node_t *node  = NULL;
  while (peek(parser) != NULL)
  {
    node = parse_statement(parser);
    if (!node)
      return NULL;

    if (!last)
    {
      first = node;
      last  = node;
    }
    else
    {
      last->next = node;
      last       = node;
    }
  }

  if (!first)
  {
    snprintf(plast_err, sizeof(plast_err), "nothing to parse >:(");
    return NULL;
  }

  return first;
}

// Parse tokens and build AST
int fzl_parse(fzl_parser_t *parser)
{
  fzl_node_t *ast = parse_program(parser);
  if (!ast)
    return -1;

  parser->root = ast;
  return 0;
}

// --- Parse zone end ---

char *fzl_parser_geterror() { return plast_err; }

void fzl_freenode(fzl_node_t *node)
{
  if (!node)
    return;
  free(node);
}

void fzl_freeast(fzl_node_t *root)
{
  if (!root)
    return;

  while (root)
  {
    if (root->type == FZL_NODETYPE_BLOCK)
    {
      if (root->list.entries)
      {
        for (size_t i = 0; i < root->list.entry_count; i++)
          if (root->list.entries[i])
            fzl_freeast(root->list.entries[i]);

        free(root->list.entries);
      }
    }
    else if (root->type == FZL_NODETYPE_CALL)
    {
      if (root->call.args_count)
      {
        for (size_t i = 0; i < root->call.args_count; i++)
          if (root->call.args[i])
            fzl_freeast(root->call.args[i]);

        free(root->call.args);
      }
    }
    else if (root->type == FZL_NODETYPE_ADD || root->type == FZL_NODETYPE_SUB ||
             root->type == FZL_NODETYPE_MUL || root->type == FZL_NODETYPE_DIV ||
             root->type == FZL_NODETYPE_POW || root->type == FZL_NODETYPE_MOD)
    {
      if (root->binary_op.left)
        fzl_freeast(root->binary_op.left);
      if (root->binary_op.right)
        fzl_freeast(root->binary_op.right);
    }
    else if (root->type == FZL_NODETYPE_UNARY_PLUS ||
             root->type == FZL_NODETYPE_UNARY_MINUS)
    {
      if (root->unary_op.operand)
        fzl_freeast(root->unary_op.operand);
    }

    fzl_freenode(root);
    root = root->next;
  }
}

void fzl_freeparser(fzl_parser_t *parser)
{
  if (!parser)
    return;
  free(parser);
}