/*
 * 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 "gas_assembly_generator.h"

#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static void
gas_assembly_generator_binary_operation(gas_assembly_generator_t *gen, ast_binary_operation_t *binary_operation);

static void
gas_assembly_generator_compile_function(gas_assembly_generator_t *gen, ast_function_declaration_t *func);

static void
gas_assembly_generator_compile_function_call(gas_assembly_generator_t *gen, ast_function_call_t *func_call);

static void
gas_assembly_generator_compile_literal(gas_assembly_generator_t *gen, ast_literal_t *literal);

static void
gas_assembly_generator_compile_return_stmt(gas_assembly_generator_t *gen, ast_return_stmt_t *return_stmt);

static void
gas_assembly_generator_compile_variable_declaration(gas_assembly_generator_t *gen,
                                                    ast_variable_declaration_t *variable_declaration);
static void
gas_assembly_generator_compile_variable_assignment(gas_assembly_generator_t *gen,
                                                   ast_variable_assignment_t *variable_assignment);

static void
gas_assembly_generator_compile_block(gas_assembly_generator_t *gen, ast_block_t *block);

static void
gas_assembly_generator_compile_if_stmt(gas_assembly_generator_t *gen, ast_if_stmt_t *if_stmt);

static void
gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable);

static void
gas_assembly_generator_set_latest_evaluation_as_literal(gas_assembly_generator_t *gen, int64_t value)
{
  gen->latest_evaluation =
    (evaluation_result_t){ .kind = EVALUATION_RESULT_AS_LITERAL_INTEGER, .data = { .literal_int = value } };
}

static void
gas_assembly_generator_set_latest_evaluation_as_literal_boolean(gas_assembly_generator_t *gen, bool value)
{
  gen->latest_evaluation =
    (evaluation_result_t){ .kind = EVALUATION_RESULT_AS_LITERAL_BOOL, .data = { .literal_bool = value } };
}

static void
gas_assembly_generator_set_latest_evaluation_to_rax(gas_assembly_generator_t *gen)
{
  gen->latest_evaluation = (evaluation_result_t){ .kind = EVALUATION_RESULT_ON_RAX };
}

static void
gas_assembly_generator_set_latest_evaluation_on_stack(gas_assembly_generator_t *gen, int stack_offset)
{
  gen->latest_evaluation =
    (evaluation_result_t){ .kind = EVALUATION_RESULT_ON_STACK, .data = { .stack_offset = stack_offset } };
}

static void
gas_assembly_generator_compile_condition(gas_assembly_generator_t *gen,
                                         ast_node_t *condition,
                                         uint64_t true_label_index,
                                         bool jump_when_true,
                                         uint64_t false_label_index,
                                         bool jump_when_false);

typedef struct ref_entry_t
{
  ast_identifier_t *id;
  int stack_offset;
} ref_entry_t;

// TODO: We need a destroy function for ref_entry.
static ref_entry_t *
ref_entry_new(void)
{
  ref_entry_t *entry = (ref_entry_t *)malloc(sizeof(ref_entry_t));

  if (entry == NULL) {
    fprintf(stderr, "ref_entry_new: Out of memory: %s\n", strerror(errno));
    exit(EXIT_FAILURE);
  }

  return entry;
}

static int
ref_find_variable_reference_stack_offset(vector_t *refs, ast_identifier_t *identifier)
{
  for (int i = refs->size - 1; i >= 0; --i) {
    ref_entry_t *entry = (ref_entry_t *)vector_at(refs, i);
    if (entry->id == identifier) {
      return entry->stack_offset;
    }
  }
  assert(false);
}

static void
ref_set_variable_reference_stack_offset(vector_t *refs, ast_identifier_t *identifier, int stack_offset)
{
  ref_entry_t *ref_entry = ref_entry_new();

  *ref_entry = (ref_entry_t){
    .id = identifier,
    .stack_offset = stack_offset,
  };

  vector_push_back(refs, ref_entry);
}

void
gas_assembly_generator_init(gas_assembly_generator_t *gen, FILE *stream)
{
  assert(gen && stream);
  gen->stream = stream;
  gen->refs = vector_new();
  gen->stack_offset = 0;
  gen->counter_label_index = 1;
  gen->return_label_index = 0;
}

void
gas_assembly_generator_compile(gas_assembly_generator_t *gen, ast_node_t *ast)
{
  gen->latest_evaluation.kind = EVALUATION_RESULT_VOID;

  switch (ast->kind) {
    case AST_NAMESPACE:
      gas_assembly_generator_compile_ns(gen, &ast->data.ns);
      break;
    case AST_BINARY_OPERATION:
      gas_assembly_generator_binary_operation(gen, &ast->data.binary_operation);
      break;
    case AST_FUNCTION_DECLARATION:
      gas_assembly_generator_compile_function(gen, &ast->data.function);
      break;
    case AST_FUNCTION_CALL:
      gas_assembly_generator_compile_function_call(gen, &ast->data.function_call);
      break;
    case AST_LITERAL:
      gas_assembly_generator_compile_literal(gen, &ast->data.literal);
      break;
    case AST_RETURN_STMT:
      gas_assembly_generator_compile_return_stmt(gen, &ast->data.return_stmt);
      break;
    case AST_VARIABLE_DECLARATION:
      gas_assembly_generator_compile_variable_declaration(gen, &ast->data.variable_declaration);
      break;
    case AST_VARIABLE_ASSIGNMENT:
      gas_assembly_generator_compile_variable_assignment(gen, &ast->data.variable_assignment);
      break;
    case AST_VARIABLE:
      gas_assembly_generator_compile_variable(gen, &ast->data.variable);
      break;
    case AST_BLOCK:
      gas_assembly_generator_compile_block(gen, &ast->data.block);
      break;
    case AST_IF_STMT:
      gas_assembly_generator_compile_if_stmt(gen, &ast->data.if_stmt);
      break;
    case AST_UNKOWN_NODE:
      assert(false && "unreachable");
  }
}

void
gas_assembly_generator_compile_ns(gas_assembly_generator_t *gen, ast_namespace_t *ns)
{
  for (size_t i = 0; i < ns->nodes->size; i++) {
    gas_assembly_generator_compile(gen, vector_at(ns->nodes, i));
  }
}

static void
gas_assembly_generator_compile_function(gas_assembly_generator_t *gen, ast_function_declaration_t *func)
{
  assert(func);

  uint64_t previous_index = gen->return_label_index;
  uint64_t return_label_index = gen->return_label_index = gen->counter_label_index++;

  fprintf(gen->stream, SVFMT ":\n", SVARG(&func->identifier.name));
  fprintf(gen->stream, "  push %%rbp\n");
  fprintf(gen->stream, "  mov %%rsp, %%rbp\n");

  gas_assembly_generator_compile(gen, func->body);

  fprintf(gen->stream, ".L%ld:\n", return_label_index);
  fprintf(gen->stream, "  pop %%rbp\n");
  fprintf(gen->stream, "  ret\n");

  gen->return_label_index = previous_index;
}

static void
gas_assembly_generator_compile_function_call(gas_assembly_generator_t *gen, ast_function_call_t *func_call)
{
  fprintf(gen->stream, "  call " SVFMT "\n", SVARG(&func_call->identifier->name));
  gas_assembly_generator_set_latest_evaluation_to_rax(gen);
}

void
gas_assembly_generator_compile_linux_main(gas_assembly_generator_t *gen, ast_node_t *ns)
{
  if (ast_node_ns_get_function_node_by_name(ns, "main") == NULL) {
    fprintf(stderr, "[ERROR]: no main function has been defined!\n");
    exit(EXIT_FAILURE);
  }

  fprintf(gen->stream, ".global _start\n");
  fprintf(gen->stream, ".text\n");

  gas_assembly_generator_compile(gen, ns);

  fprintf(gen->stream, "_start:\n");
  fprintf(gen->stream, "  call main\n");

  switch (gen->latest_evaluation.kind) {
    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
      fprintf(gen->stream, "  mov $%ld, %%rdi\n", gen->latest_evaluation.data.literal_int);
      break;
    case EVALUATION_RESULT_ON_RAX:
      fprintf(gen->stream, "  mov %%rax, %%rdi\n");
      break;
    case EVALUATION_RESULT_ON_STACK:
      fprintf(gen->stream, "  mov %d(%%rbp), %%rdi\n", gen->latest_evaluation.data.stack_offset);
      break;
    case EVALUATION_RESULT_VOID:
      fprintf(gen->stream, "  xor %%rdi, %%rdi\n");
      break;
    case EVALUATION_RESULT_AS_LITERAL_BOOL:
      assert(false && "Unexpected");
      break;
  }

  fprintf(gen->stream, "  mov $60, %%rax\n");
  fprintf(gen->stream, "  syscall\n");
}

static void
gas_assembly_generator_compile_return_stmt(gas_assembly_generator_t *gen, ast_return_stmt_t *return_stmt)
{
  assert(gen && return_stmt);
  assert(gen->return_label_index != 0 && "There are no return label index set");

  gas_assembly_generator_compile(gen, return_stmt->argument);

  switch (gen->latest_evaluation.kind) {
    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
      fprintf(gen->stream, "  mov $%ld, %%rax\n", gen->latest_evaluation.data.literal_int);
      gas_assembly_generator_set_latest_evaluation_to_rax(gen);
      break;
    case EVALUATION_RESULT_ON_RAX:
      gas_assembly_generator_set_latest_evaluation_to_rax(gen);
      break;
    case EVALUATION_RESULT_ON_STACK:
      fprintf(gen->stream, "  mov %d(%%rbp), %%rax\n", gen->latest_evaluation.data.stack_offset);
      gas_assembly_generator_set_latest_evaluation_to_rax(gen);
      break;
    case EVALUATION_RESULT_VOID:
      break;
    case EVALUATION_RESULT_AS_LITERAL_BOOL:
      assert(false && "Unexpected");
      break;
  }

  fprintf(gen->stream, "  jmp .L%ld\n", gen->return_label_index);
}

static void
gas_assembly_generator_compile_variable_assign_value(gas_assembly_generator_t *gen,
                                                     ast_node_t *expression,
                                                     int stack_offset)
{
  if (expression->result_type == TYPE_BOOL) {
    if (expression->kind == AST_LITERAL) {
      assert(expression->data.literal.kind == AST_LITERAL_BOOL);

      fprintf(gen->stream, "  movq $%d, %d(%%rbp)\n", expression->data.literal.value.boolean, stack_offset);
      return;
    }

    uint64_t true_label_index = gen->counter_label_index++;
    uint64_t false_label_index = gen->counter_label_index++;
    uint64_t after_assign_label_index = gen->counter_label_index++;

    gas_assembly_generator_compile_condition(gen, expression, true_label_index, false, false_label_index, true);

    fprintf(gen->stream, ".L_%ld:\n", true_label_index);
    fprintf(gen->stream, "  movq $1, %d(%%rbp)\n", stack_offset);
    fprintf(gen->stream, "  jmp .L_%ld\n", after_assign_label_index);

    fprintf(gen->stream, ".L_%ld:\n", false_label_index);
    fprintf(gen->stream, "  movq $0, %d(%%rbp)\n", stack_offset);

    fprintf(gen->stream, ".L_%ld:\n", after_assign_label_index);
    return;
  }

  gas_assembly_generator_compile(gen, expression);

  switch (gen->latest_evaluation.kind) {
    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
      fprintf(gen->stream, "  movq $%ld, %d(%%rbp)\n", gen->latest_evaluation.data.literal_int, stack_offset);
      break;
    case EVALUATION_RESULT_ON_RAX:
      fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", stack_offset);
      break;
    case EVALUATION_RESULT_ON_STACK:
      fprintf(gen->stream, "  movq %d(%%rbp), %%rax\n", gen->latest_evaluation.data.stack_offset);
      fprintf(gen->stream, "  movq %%rax, %d(%%rbp)\n", stack_offset);
      break;
    case EVALUATION_RESULT_AS_LITERAL_BOOL:
    case EVALUATION_RESULT_VOID:
      assert(false && "Unexpected void result for variable declaration");
      // FIXME: store the latest node at latest evaluation and print_ast
  }
}

static void
gas_assembly_generator_compile_variable_declaration(gas_assembly_generator_t *gen,
                                                    ast_variable_declaration_t *variable_declaration)
{
  gen->stack_offset -= 8;
  ref_set_variable_reference_stack_offset(gen->refs, &variable_declaration->identifier, gen->stack_offset);
  gas_assembly_generator_compile_variable_assign_value(gen, variable_declaration->value, gen->stack_offset);
}

static void
gas_assembly_generator_compile_variable_assignment(gas_assembly_generator_t *gen,
                                                   ast_variable_assignment_t *variable_assignment)
{
  int stack_offset = ref_find_variable_reference_stack_offset(gen->refs, variable_assignment->identifier);

  gas_assembly_generator_compile_variable_assign_value(gen, variable_assignment->expression, stack_offset);
}

static void
gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable)
{
  int stack_offset = ref_find_variable_reference_stack_offset(gen->refs, variable->identifier);

  gas_assembly_generator_set_latest_evaluation_on_stack(gen, stack_offset);
}

static void
gas_assembly_generator_compile_literal(gas_assembly_generator_t *gen, ast_literal_t *literal)
{
  switch (literal->kind) {
    case AST_LITERAL_INTEGER:
      gas_assembly_generator_set_latest_evaluation_as_literal(gen, literal->value.integer);
      return;
    case AST_LITERAL_BOOL:
      gas_assembly_generator_set_latest_evaluation_as_literal_boolean(gen, literal->value.boolean);
      return;
    default:
      assert(false && "no code generation strategy.");
  }
}

static void
gas_assembly_generator_compile_block(gas_assembly_generator_t *gen, ast_block_t *block)
{
  for (size_t i = 0; i < block->body->size; i++) {
    gas_assembly_generator_compile(gen, vector_at(block->body, i));
  }
}

static void
gas_assembly_generator_binary_operation(gas_assembly_generator_t *gen, ast_binary_operation_t *binary_operation)
{
  gas_assembly_generator_compile(gen, binary_operation->right);
  evaluation_result_t right_evaluation = gen->latest_evaluation;

  if (right_evaluation.kind == EVALUATION_RESULT_ON_RAX) {
    gen->stack_offset -= 8;
    fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", gen->stack_offset);
  }

  gas_assembly_generator_compile(gen, binary_operation->left);
  evaluation_result_t left_evaluation = gen->latest_evaluation;

  switch (right_evaluation.kind) {
    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
      fprintf(gen->stream, "  mov $%ld, %%rcx\n", right_evaluation.data.literal_int);
      break;
    case EVALUATION_RESULT_ON_RAX:
      fprintf(gen->stream, "  mov %d(%%rbp), %%rcx\n", gen->stack_offset);
      gen->stack_offset += 8;
      break;
    case EVALUATION_RESULT_ON_STACK:
      fprintf(gen->stream, "  movq %d(%%rbp), %%rcx\n", right_evaluation.data.stack_offset);
      break;
    case EVALUATION_RESULT_AS_LITERAL_BOOL:
      fprintf(gen->stream, "  mov $%d, %%rcx\n", right_evaluation.data.literal_bool);
      break;
    case EVALUATION_RESULT_VOID:
      assert(false && "Unexpected void result for binary operation");
      // FIXME: store the latest node at latest evaluation and print_ast
  }

  gas_assembly_generator_set_latest_evaluation_to_rax(gen);

  switch (left_evaluation.kind) {
    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
      fprintf(gen->stream, "  mov $%ld, %%rax\n", left_evaluation.data.literal_int);
      break;
    case EVALUATION_RESULT_ON_STACK:
      fprintf(gen->stream, "  movq %d(%%rbp), %%rax\n", left_evaluation.data.stack_offset);
      break;
    case EVALUATION_RESULT_ON_RAX:
      break;
    case EVALUATION_RESULT_AS_LITERAL_BOOL:
      fprintf(gen->stream, "  mov $%d, %%rax\n", left_evaluation.data.literal_bool);
      break;
    case EVALUATION_RESULT_VOID:
      assert(false && "Unexpected void result for binary operation");
  }

  if (binary_operation->kind == AST_BINOP_ADITION) {
    fprintf(gen->stream, "  add %%rcx, %%rax\n");
    return;
  }

  if (binary_operation->kind == AST_BINOP_SUBTRACTION) {
    fprintf(gen->stream, "  sub %%rcx, %%rax\n");
    return;
  }

  if (binary_operation->kind == AST_BINOP_MULTIPLICATION) {
    fprintf(gen->stream, "  mul %%rcx\n");
    return;
  }

  if (binary_operation->kind == AST_BINOP_DIVISION) {
    fprintf(gen->stream, "  xor %%rdx, %%rdx\n");
    fprintf(gen->stream, "  div %%rcx\n");
    return;
  }

  if (binary_operation->kind == AST_BINOP_LT || binary_operation->kind == AST_BINOP_GT ||
      binary_operation->kind == AST_BINOP_EQUAL || binary_operation->kind == AST_BINOP_NOT_EQUAL ||
      binary_operation->kind == AST_BINOP_GT || binary_operation->kind == AST_BINOP_LT ||
      binary_operation->kind == AST_BINOP_LT_EQUAL || binary_operation->kind == AST_BINOP_GT_EQUAL) {
    fprintf(gen->stream, "  cmp %%rcx, %%rax\n");
    return;
  }

  assert(false && "No strategy defined for binary operation");
}

static char *jump_map[AST_BINOP_N] = {
  [AST_BINOP_EQUAL] = "je",     [AST_BINOP_NOT_EQUAL] = "jne", [AST_BINOP_LT] = "jl",
  [AST_BINOP_LT_EQUAL] = "jle", [AST_BINOP_GT] = "jg",         [AST_BINOP_GT_EQUAL] = "jge",
};

static char *inverse_jump_map[AST_BINOP_N] = {
  [AST_BINOP_EQUAL] = "jne",   [AST_BINOP_NOT_EQUAL] = "je", [AST_BINOP_LT] = "jge",
  [AST_BINOP_LT_EQUAL] = "jg", [AST_BINOP_GT] = "jle",       [AST_BINOP_GT_EQUAL] = "jl",
};

static void
gas_assembly_generator_compile_condition(gas_assembly_generator_t *gen,
                                         ast_node_t *condition,
                                         uint64_t true_label_index,
                                         bool jump_when_true,
                                         uint64_t false_label_index,
                                         bool jump_when_false)
{
  assert(condition->kind == AST_BINARY_OPERATION || condition->kind == AST_LITERAL);

  if (condition->kind == AST_LITERAL) {
    assert(condition->data.literal.kind == AST_LITERAL_BOOL);

    if (jump_when_false && !condition->data.literal.value.boolean) {
      fprintf(gen->stream, "  jmp .L_%ld\n", false_label_index);
      return;
    }

    if (jump_when_true && condition->data.literal.value.boolean) {
      fprintf(gen->stream, "  jmp .L_%ld\n", true_label_index);
    }

    return;
  }

  if (condition->data.binary_operation.kind == AST_BINOP_OR) {
    uint64_t or_false_branch_label_index = gen->counter_label_index++;

    gas_assembly_generator_compile_condition(
      gen, condition->data.binary_operation.left, true_label_index, true, or_false_branch_label_index, false);

    fprintf(gen->stream, ".L_%ld:\n", or_false_branch_label_index);

    gas_assembly_generator_compile_condition(
      gen, condition->data.binary_operation.right, true_label_index, false, false_label_index, true);
    return;
  }

  if (condition->data.binary_operation.kind == AST_BINOP_AND) {
    uint64_t and_true_branch_label_index = gen->counter_label_index++;

    gas_assembly_generator_compile_condition(
      gen, condition->data.binary_operation.left, and_true_branch_label_index, false, false_label_index, true);

    fprintf(gen->stream, ".L_%ld:\n", and_true_branch_label_index);

    gas_assembly_generator_compile_condition(
      gen, condition->data.binary_operation.right, true_label_index, false, false_label_index, true);
    return;
  }

  gas_assembly_generator_compile(gen, condition);

  if (jump_when_false) {
    char *jumper = inverse_jump_map[condition->data.binary_operation.kind];

    assert(jumper != NULL);

    fprintf(gen->stream, "  %s .L_%ld\n", jumper, false_label_index);
    return;
  }

  if (jump_when_true) {
    char *jumper = jump_map[condition->data.binary_operation.kind];

    assert(jumper != NULL);

    fprintf(gen->stream, "  %s .L_%ld\n", jumper, true_label_index);
  }
}

static void
gas_assembly_generator_compile_if_stmt(gas_assembly_generator_t *gen, ast_if_stmt_t *if_stmt)
{
  switch (if_stmt->condition->kind) {
    case AST_BINARY_OPERATION:
    case AST_LITERAL: {
      uint64_t true_label_index = gen->counter_label_index++;
      uint64_t false_label_index = gen->counter_label_index++;

      gas_assembly_generator_compile_condition(
        gen, if_stmt->condition, true_label_index, false, false_label_index, true);

      fprintf(gen->stream, ".L_%ld:\n", true_label_index);
      gas_assembly_generator_compile(gen, if_stmt->body);

      fprintf(gen->stream, ".L_%ld:\n", false_label_index);
      break;
    }
    case AST_VARIABLE: {
      uint64_t false_label_index = gen->counter_label_index++;
      int stack_offset =
        ref_find_variable_reference_stack_offset(gen->refs, if_stmt->condition->data.variable.identifier);

      fprintf(gen->stream, "  movq %d(%%rbp), %%rax\n", stack_offset);
      fprintf(gen->stream, "  cmpq $1, %%rax\n");
      fprintf(gen->stream, "  jne .L_%ld\n", false_label_index);

      gas_assembly_generator_compile(gen, if_stmt->body);

      fprintf(gen->stream, ".L_%ld:\n", false_label_index);
      break;
    }
    default:
      assert(false);
  }
}