/*
 * Copyright (C) 2023 Johnny Richard
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
#include "ast_pretty_printer.h"
#include "ast.h"
#include "string_view.h"

#include <assert.h>
#include <stdarg.h>
#include <stdio.h>

const char *operation_kinds[AST_BINOP_N] = {
  [AST_BINOP_ADITION] = "+",  [AST_BINOP_SUBTRACTION] = "-", [AST_BINOP_MULTIPLICATION] = "*",
  [AST_BINOP_DIVISION] = "/", [AST_BINOP_EQUAL] = "==",      [AST_BINOP_NOT_EQUAL] = "!=",
  [AST_BINOP_AND] = "&&",     [AST_BINOP_OR] = "||",         [AST_BINOP_GT] = ">",
  [AST_BINOP_LT] = "<",       [AST_BINOP_LT_EQUAL] = "<=",   [AST_BINOP_GT_EQUAL] = ">=",
};

void
ast_pretty_printer_init(ast_pretty_printer_t *printer, FILE *stream);

inline static void
ast_pretty_printer_add_indentation(ast_pretty_printer_t *printer);

inline static void
ast_pretty_printer_rm_indentation(ast_pretty_printer_t *printer);

inline static void
ast_pretty_printer_set_lst_children(ast_pretty_printer_t *printer);

static void
ast_pretty_printer_print_indentation(ast_pretty_printer_t *printer);

static int
ast_pretty_printer_printf(ast_pretty_printer_t *printer, const char *fmt, ...);

void
ast_pretty_printer_init(ast_pretty_printer_t *printer, FILE *stream)
{
  assert(printer);
  assert(stream);

  printer->indentation_fmt = 0ULL;
  printer->indentation_lst_children = 0ULL;
  printer->indentation_level = 0;
  printer->stream = stream;
}

void
ast_pretty_printer_print_ast(ast_pretty_printer_t *printer, ast_node_t *ast)
{
  assert(ast);

  switch (ast->kind) {
    case AST_NAMESPACE:
      ast_pretty_printer_printf(printer, "Namespace\n");
      ast_pretty_printer_add_indentation(printer);
      {
        for (size_t i = 0; i < ast->data.ns.nodes->size; ++i) {
          if (i + 1 >= ast->data.ns.nodes->size) {
            ast_pretty_printer_set_lst_children(printer);
          }
          ast_pretty_printer_print_ast(printer, vector_at(ast->data.ns.nodes, i));
        }
        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    case AST_IF_STMT: {
      ast_if_stmt_t if_stmt = ast->data.if_stmt;
      ast_pretty_printer_printf(printer, "IfStmt\n");
      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_printf(printer, "condition:\n");
        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, if_stmt.condition);

          ast_pretty_printer_rm_indentation(printer);
        }
        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_printf(printer, "body:\n");
        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, if_stmt.body);

          ast_pretty_printer_rm_indentation(printer);
        }
        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_BINARY_OPERATION: {
      ast_binary_operation_t binop = ast->data.binary_operation;
      ast_pretty_printer_printf(printer, "BinaryOperation operation='%s'\n", operation_kinds[binop.kind]);

      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_printf(printer, "left:\n");

        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, binop.left);

          ast_pretty_printer_rm_indentation(printer);
        }

        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_printf(printer, "right:\n");

        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, binop.right);

          ast_pretty_printer_rm_indentation(printer);
        }

        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_FUNCTION_DECLARATION: {
      ast_function_declaration_t function = ast->data.function;
      ast_pretty_printer_printf(printer, "FunctionDecl name='" SVFMT "'\n", SVARG(&function.identifier.name));

      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_printf(printer, "body:\n");

        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, ast->data.function.body);
          ast_pretty_printer_rm_indentation(printer);
        }

        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_BLOCK: {
      ast_pretty_printer_printf(printer, "Block\n");
      ast_pretty_printer_add_indentation(printer);
      {
        for (size_t i = 0; i < ast->data.block.body->size; ++i) {
          if (i + 1 >= ast->data.block.body->size) {
            ast_pretty_printer_set_lst_children(printer);
          }
          ast_pretty_printer_print_ast(printer, vector_at(ast->data.block.body, i));
        }
        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_LITERAL: {
      ast_literal_t literal = ast->data.literal;

      switch (literal.kind) {
        case AST_LITERAL_INTEGER: {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_printf(printer, "Literal type=i32 value='%d'\n", literal.value.integer);
          break;
        }
        case AST_LITERAL_BOOL:
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_printf(
            printer, "Literal type=bool value='%s'\n", literal.value.boolean ? "true" : "false");
          break;
      }
      break;
    }
    case AST_RETURN_STMT: {
      ast_return_stmt_t return_stmt = ast->data.return_stmt;
      ast_pretty_printer_printf(printer, "ReturnStmt\n");

      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_print_ast(printer, return_stmt.argument);

        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_VARIABLE_DECLARATION: {
      ast_variable_declaration_t var_decl = ast->data.variable_declaration;
      ast_pretty_printer_printf(printer, "VariableDecl name='" SVFMT "'\n", SVARG(&var_decl.identifier.name));

      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_print_ast(printer, var_decl.value);

        ast_pretty_printer_rm_indentation(printer);
      }
      break;
    }
    case AST_VARIABLE_ASSIGNMENT: {
      ast_variable_assignment_t var_assign = ast->data.variable_assignment;
      ast_pretty_printer_printf(printer, "VariableAssignment name='" SVFMT "'\n", SVARG(&var_assign.identifier->name));

      ast_pretty_printer_add_indentation(printer);
      {
        ast_pretty_printer_set_lst_children(printer);
        ast_pretty_printer_printf(printer, "expression:\n");

        ast_pretty_printer_add_indentation(printer);
        {
          ast_pretty_printer_set_lst_children(printer);
          ast_pretty_printer_print_ast(printer, var_assign.expression);

          ast_pretty_printer_rm_indentation(printer);
        }
        ast_pretty_printer_rm_indentation(printer);
      }

      break;
    }
    case AST_VARIABLE: {
      ast_variable_t var = ast->data.variable;
      ast_pretty_printer_printf(printer, "Variable name='" SVFMT "'\n", SVARG(&var.identifier->name));
      break;
    }
    case AST_UNKOWN_NODE: {
      fprintf(printer->stream, "AST_UNKOWN_NODE\n");
      break;
    }
  }
}

static int
ast_pretty_printer_printf(ast_pretty_printer_t *printer, const char *fmt, ...)
{
  int ret = 0;
  va_list args;

  va_start(args, fmt);
  ast_pretty_printer_print_indentation(printer);
  ret = vprintf(fmt, args);
  va_end(args);

  return ret;
}

inline static void
ast_pretty_printer_add_indentation(ast_pretty_printer_t *printer)
{
  printer->indentation_level++;
  printer->indentation_fmt |= 1ULL << printer->indentation_level;
}

inline static void
ast_pretty_printer_rm_indentation(ast_pretty_printer_t *printer)
{
  printer->indentation_level--;
}

inline static void
ast_pretty_printer_set_lst_children(ast_pretty_printer_t *printer)
{
  printer->indentation_lst_children |= 1ULL << printer->indentation_level;
}

static void
ast_pretty_printer_print_indentation(ast_pretty_printer_t *printer)
{
  for (size_t i = 0; i <= printer->indentation_level; ++i) {
    if ((printer->indentation_fmt >> i) & 1ULL) {
      if (i + 1 > printer->indentation_level && (printer->indentation_lst_children & (1ULL << i))) {
        printer->indentation_lst_children ^= (1ULL << i);
        printer->indentation_fmt ^= (1ULL << i);
        fprintf(printer->stream, "└─ ");
      } else if (i + 1 > printer->indentation_level) {
        fprintf(printer->stream, "├─ ");
      } else {
        fprintf(printer->stream, "│  ");
      }
    } else {
      fprintf(printer->stream, "  ");
    }
  }
}