/*
 * 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(gas_assembly_generator_t *gen, ast_variable_t *variable);

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


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)
{
  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_IDENTIFIER:
      assert(false && "TODO: ast identifier not implemented yet");
      break;
    case AST_VARIABLE_DECLARATION:
      gas_assembly_generator_compile_variable_declaration(gen, &ast->data.variable_declaration);
      break;
    case AST_VARIABLE:
      gas_assembly_generator_compile_variable(gen, &ast->data.variable);
      break;
    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");

  for (size_t i = 0; i < func->body->size; i++) {
    gas_assembly_generator_compile(gen, vector_at(func->body, i));
  }

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

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

  fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", entry->stack_offset);
}

static void
gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable)
{
  ref_entry_t *entry = NULL;
  for (int i = gen->refs->size - 1; i >= 0; --i) {
    ref_entry_t *entry_tmp = (ref_entry_t *)vector_at(gen->refs, i);
    if (entry_tmp->id == variable->identifier) {
      entry = entry_tmp;
      break;
    }
  }
  assert(entry && "reference not found");
  fprintf(gen->stream, "  mov %d(%%rbp), %%rax\n", 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:
      fprintf(gen->stream, "  mov $%d, %%rax\n", literal->value.integer);
      return;
    default:
      assert(false && "no code generation strategy.");
  }
}

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

  gen->stack_offset -= 8;
  fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", gen->stack_offset);

  gas_assembly_generator_compile(gen, binary_operation->left);

  fprintf(gen->stream, "  mov %d(%%rbp), %%rcx\n", gen->stack_offset);
  gen->stack_offset += 8;

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