/*
 * 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_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_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_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 } };
}

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 ref_entry_t *
find_ref_entry(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;
    }
  }
  return NULL;
}

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;
}

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_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_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:
    case AST_UNKOWN_NODE:
      assert(false && "unreachable");
  }
}

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

  if (!string_view_eq(func->identifier.name, string_view_from_str("main"))) {
    fprintf(stderr, "[ERROR]: no main function has been defined!\n");
    exit(EXIT_FAILURE);
  }

  fprintf(gen->stream, ".global _start\n");
  fprintf(gen->stream, ".text\n");
  fprintf(gen->stream, "_start:\n");
  fprintf(gen->stream, "  push %%rbp\n");
  fprintf(gen->stream, "  mov %%rsp, %%rbp\n");

  gas_assembly_generator_compile(gen, func->body);

  fprintf(gen->stream, "  pop %%rbp\n");
}

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

  gas_assembly_generator_compile(gen, return_stmt->argument);

  if (gen->latest_evaluation.kind == EVALUATION_RESULT_AS_LITERAL_INTEGER) {
    fprintf(gen->stream, "  mov $%ld, %%rdi\n", gen->latest_evaluation.data.literal_int);
  } else {
    fprintf(gen->stream, "  mov %%rax, %%rdi\n");
  }
  fprintf(gen->stream, "  mov $60, %%rax\n");
  fprintf(gen->stream, "  syscall\n");
}

static void
gas_assembly_generator_compile_variable_declaration(gas_assembly_generator_t *gen,
                                                    ast_variable_declaration_t *variable_declaration)
{
  gen->stack_offset -= 8;
  gas_assembly_generator_compile(gen, variable_declaration->value);

  ref_entry_t *entry = ref_entry_new();
  *entry = (ref_entry_t){ .id = &variable_declaration->identifier, .stack_offset = gen->stack_offset };
  vector_push_back(gen->refs, entry);

  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, entry->stack_offset);
      break;
    case EVALUATION_RESULT_ON_RAX:
      fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", entry->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", entry->stack_offset);
      break;
    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_assignment(gas_assembly_generator_t *gen,
                                                   ast_variable_assignment_t *variable_assignment)
{
  ref_entry_t *entry = find_ref_entry(gen->refs, variable_assignment->identifier);
  assert(entry && "reference not found");

  gas_assembly_generator_compile(gen, variable_assignment->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, entry->stack_offset);
      break;
    case EVALUATION_RESULT_ON_RAX:
      fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", entry->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", entry->stack_offset);
      break;
    case EVALUATION_RESULT_VOID:
      assert(false && "Unexpected void result for variable declaration");
      // FIXME: store the latest node at latest evaluation and print_ast
  }
  // Today we dont't support chain-assignment like this:
  //
  //   a = b = 1;
  //
  // If we start supporting it, we need to change the latest evaluation kind to
  // stack.
  gen->latest_evaluation.kind = EVALUATION_RESULT_VOID;
}

static void
gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable)
{
  ref_entry_t *entry = find_ref_entry(gen->refs, variable->identifier);
  assert(entry && "reference not found");
  gas_assembly_generator_set_latest_evaluation_on_stack(gen, entry->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;
    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_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_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;
  }

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