#include "codegen.h"
#include "../mem_arena/arena.h"
#include <stdio.h>
#include <string.h>

#define ARR_LEN(ARR) sizeof(ARR) / sizeof(ARR[0])
#define MEM_ARENA_LENGTH 10 * 1024 * 1024
#define LONG_BUFFER_LENGTH 8192
#define MEDIUM_BUFFER_LENGTH 4096
#define SHORT_BUFFER_LENGTH 1024

enum vec_funcs {
  VEC_FUNC_ADD,
  VEC_FUNC_SUB,
  VEC_FUNC_MUL,
  VEC_FUNC_DOT,
  VEC_FUNC_MAGNITUDE,
  VEC_FUNC_PRINT,

  COUNT_VEC_FUNCS,
};

void write_vec_func_implementation(arena_t *arena, func_data_t *func,
                                   const struct_def_t *s, char *body,
                                   const char *operator, const char * separator,
                                   const char *end);

int main(int argc, char *argv[]) {
  const char *output_file = "vec.h";
  FILE *temp = fopen(output_file, "w");
  fclose(temp);

  FILE *fp = fopen(output_file, "a");

  arena_t main_mem = {0};
  arena_init(&main_mem, MEM_ARENA_LENGTH);

  arena_t temp_mem = {0};
  arena_init(&temp_mem, MEM_ARENA_LENGTH);

  const char *header_name = "VEC_H";

  write_header_guard_open(&temp_mem, fp, header_name);

  include_header_t headers[] = {
      {.type = INCLUDE_TYPE_USER, .header = "../aliases.h"},
      {.type = INCLUDE_TYPE_SYSTEM, .header = "stdio.h"},
      {.type = INCLUDE_TYPE_SYSTEM, .header = "math.h"},
  };

  for (u64 i = 0; i < ARR_LEN(headers); ++i) {
    write_include(&temp_mem, fp, &(headers[i]));
  }

  fwrite("\n", 1, 1, fp);

  type_info_t base_types[] = {
      (type_info_t){.name = "i32", .short_name = 'i', .printf_spec = 'd'},
      (type_info_t){.name = "f32", .short_name = 'f', .printf_spec = 'f'},
  };

  u64 dimensions[] = {2, 3};

  const char *member_names[] = {"x", "y", "z"};

  const char *arg_names[] = {"v1", "v2"};

  func_data_t funcs[COUNT_VEC_FUNCS] = {
      [VEC_FUNC_ADD] =
          (func_data_t){
              .base_name = "vec_add",
              .arg_count = 2,
              .arg_names = arg_names,
          },
      [VEC_FUNC_SUB] =
          (func_data_t){
              .base_name = "vec_sub",
              .arg_count = 2,
              .arg_names = arg_names,
          },
      [VEC_FUNC_MUL] =
          (func_data_t){
              .base_name = "vec_mul",
              .arg_count = 2,
              .arg_names = arg_names,
          },
      [VEC_FUNC_DOT] =
          (func_data_t){
              .base_name = "vec_dot",
              .arg_count = 2,
              .arg_names = arg_names,
          },
      [VEC_FUNC_MAGNITUDE] =
          (func_data_t){
              .base_name = "vec_magnitude",
              .arg_count = 1,
              .arg_names = arg_names,
          },
      [VEC_FUNC_PRINT] =
          (func_data_t){
              .base_name = "vec_print",
              .arg_count = 1,
              .arg_names = arg_names,
          },
  };

  for (u64 i = 0; i < COUNT_VEC_FUNCS; ++i) {
    write_macro_func(&temp_mem, fp, &(funcs[i]));
  }

  data_type_t void_data_type = (data_type_t){
      .type = DATA_TYPE_NORMAL,
      .type_info = (type_info_t){.name = "void"},
  };

  for (u64 i = 0; i < ARR_LEN(base_types); ++i) {
    for (u64 j = 0; j < ARR_LEN(dimensions); ++j) {
      type_info_t base_type = base_types[i];
      data_type_t base_data_type =
          (data_type_t){.type = DATA_TYPE_NORMAL, .type_info = base_type};
      u64 mem_count = dimensions[j];

      char *typename = (char *)arena_alloc(&main_mem, SHORT_BUFFER_LENGTH);
      sprintf(typename, "vec%lu%c_t", mem_count, base_type.short_name);

      identifier_t members[mem_count];
      memset(members, 0, sizeof(identifier_t) * mem_count);

      for (u64 k = 0; k < mem_count; ++k) {
        members[k] = (identifier_t){
            .name = member_names[k],
            .type = base_data_type,
        };
      }

      struct_def_t s = (struct_def_t){
          .name = typename,
          .mem_count = mem_count,
          .members = members,
      };

      data_type_t vec_data_type =
          (data_type_t){.type = DATA_TYPE_STRUCT, .struct_def = s};

      write_struct(&temp_mem, fp, &s);

      funcs[VEC_FUNC_ADD].func.ret_type = vec_data_type;
      funcs[VEC_FUNC_SUB].func.ret_type = vec_data_type;
      funcs[VEC_FUNC_MUL].func.ret_type = vec_data_type;
      funcs[VEC_FUNC_DOT].func.ret_type = base_data_type;
      funcs[VEC_FUNC_MAGNITUDE].func.ret_type = base_data_type;
      funcs[VEC_FUNC_PRINT].func.ret_type = void_data_type;

      for (u64 k = 0; k < COUNT_VEC_FUNCS; ++k) {
        func_data_t *func = &(funcs[k]);

        char *name = (char *)arena_alloc(&main_mem, SHORT_BUFFER_LENGTH);
        sprintf(name, "%s_%s", func->base_name, typename);

        u64 alloc_size = sizeof(identifier_t) * func->arg_count;
        identifier_t *args = (identifier_t *)arena_alloc(&main_mem, alloc_size);

        for (u64 l = 0; l < func->arg_count; ++l) {
          args[l] =
              (identifier_t){.type = vec_data_type, .name = func->arg_names[l]};
        }

        func->func.name = name;
        func->func.args = args;

        write_func_declaration(&temp_mem, fp, func);
      }

      const char *vec_impl_begin = "#ifdef VEC_IMPLEMENTATION\n\n";
      fwrite(vec_impl_begin, strlen(vec_impl_begin), 1, fp);

      const char *operators[VEC_FUNC_DOT] = {
          [VEC_FUNC_ADD] = "+",
          [VEC_FUNC_SUB] = "-",
          [VEC_FUNC_MUL] = "*",
      };

      for (u64 k = 0; k < VEC_FUNC_DOT; ++k) {
        func_data_t *func = &(funcs[k]);

        char *body = (char *)arena_alloc(&main_mem, MEDIUM_BUFFER_LENGTH);
        sprintf(body, "\treturn (%s) {", typename);

        write_vec_func_implementation(&main_mem, func, &s, body, operators[k],
                                      ", ", "};\n");

        write_func_defintion(&temp_mem, fp, func);
      }

      func_data_t *vec_dot_func = &(funcs[VEC_FUNC_DOT]);
      char *vec_dot_body = (char *)arena_alloc(&main_mem, MEDIUM_BUFFER_LENGTH);
      strcat(vec_dot_body, "\treturn ");

      write_vec_func_implementation(&main_mem, vec_dot_func, &s, vec_dot_body,
                                    operators[VEC_FUNC_MUL], " +", " 0;\n");

      write_func_defintion(&temp_mem, fp, vec_dot_func);

      func_data_t *vec_magnitude_func = &(funcs[VEC_FUNC_MAGNITUDE]);
      char *vec_magnitude_body =
          (char *)arena_alloc(&main_mem, MEDIUM_BUFFER_LENGTH);
      sprintf(vec_magnitude_body, "\t%s dotproduct = %s(%s, %s);\n",
              base_type.name, vec_dot_func->func.name,
              vec_magnitude_func->arg_names[0],
              vec_magnitude_func->arg_names[0]);
      strcat(vec_magnitude_body, "\treturn (");
      strcat(vec_magnitude_body, base_type.name);
      strcat(vec_magnitude_body, ")sqrtf(dotproduct);\n");

      vec_magnitude_func->func.body = vec_magnitude_body;

      write_func_defintion(&temp_mem, fp, vec_magnitude_func);

      func_data_t *vec_print_func = &(funcs[VEC_FUNC_PRINT]);
      char *vec_print_body =
          (char *)arena_alloc(&main_mem, MEDIUM_BUFFER_LENGTH);
      strcat(vec_print_body, "\tprintf(\"{%");

      for (u64 k = 0; k < s.mem_count; ++k) {
        if (k > 0) {
          strcat(vec_print_body, ", %");
        }

        strcat(vec_print_body, &(base_type.printf_spec));
      }

      strcat(vec_print_body, "}\\n\",");

      for (u64 k = 0; k < s.mem_count; ++k) {
        if (k > 0) {
          strcat(vec_print_body, ",");
        }

        for (u64 l = 0; l < vec_print_func->arg_count; ++l) {
          char *arg = (char *)arena_alloc(&main_mem, SHORT_BUFFER_LENGTH);
          sprintf(arg, " %s.%s", vec_print_func->arg_names[l],
                  s.members[k].name);

          strcat(vec_print_body, arg);
        }
      }

      strcat(vec_print_body, ");\n");

      vec_print_func->func.body = vec_print_body;

      write_func_defintion(&temp_mem, fp, vec_print_func);

      const char *vec_impl_end = "#endif // VEC_IMPLEMENTATION\n\n";
      fwrite(vec_impl_end, strlen(vec_impl_end), 1, fp);
    }

    arena_clear(&main_mem);
  }

  write_header_guard_close(&temp_mem, fp, header_name);

  arena_free(&temp_mem);
  arena_free(&main_mem);

  fclose(fp);

  return 0;
}

void write_header_guard_open(arena_t *arena, FILE *fp, const char *header_id) {
  char header_guard[LONG_BUFFER_LENGTH] = {0};

  sprintf(header_guard, "#ifndef %s\n#define %s\n\n", header_id, header_id);

  u64 length = strlen(header_guard);

  fwrite(header_guard, length, 1, fp);
}

void write_header_guard_close(arena_t *arena, FILE *fp, const char *header_id) {
  char *header_guard = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!header_guard) {
    return;
  }

  sprintf(header_guard, "#endif // %s", header_id);

  u64 length = strlen(header_guard);

  fwrite(header_guard, length, 1, fp);

  arena_clear(arena);
}

void write_include(arena_t *arena, FILE *fp, const include_header_t *header) {
  u64 length = 0;

  char *include_line = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!include_line) {
    return;
  }

  switch (header->type) {
  case INCLUDE_TYPE_SYSTEM:
    sprintf(include_line, "#include <%s>\n", header->header);

    length = strlen(include_line);

    fwrite(include_line, length, 1, fp);

    break;
  case INCLUDE_TYPE_USER:
    sprintf(include_line, "#include \"%s\"\n", header->header);

    length = strlen(include_line);

    fwrite(include_line, length, 1, fp);

    break;
  default:
    break;
  }

  arena_clear(arena);
}

void write_data_type(char *dst, const data_type_t *data_type) {
  switch (data_type->type) {
  case DATA_TYPE_NORMAL:
    strcat(dst, data_type->type_info.name);

    break;
  case DATA_TYPE_STRUCT:
    strcat(dst, data_type->struct_def.name);

    break;
  default:
    break;
  }
}

void write_identifier(char *dst, const identifier_t *identifier) {
  write_data_type(dst, &(identifier->type));

  strcat(dst, " ");
  strcat(dst, identifier->name);
}

void write_struct(arena_t *arena, FILE *fp, const struct_def_t *struct_def) {
  char *struct_output = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!struct_output) {
    return;
  }

  strcat(struct_output, "typedef struct {");

  for (u64 i = 0; i < struct_def->mem_count; ++i) {
    strcat(struct_output, "\n\t");

    const identifier_t *member = &(struct_def->members[i]);

    write_identifier(struct_output, member);

    strcat(struct_output, ";");
  }

  strcat(struct_output, "\n} ");
  strcat(struct_output, struct_def->name);
  strcat(struct_output, ";\n\n");

  u64 length = strlen(struct_output);

  fwrite(struct_output, length, 1, fp);

  arena_clear(arena);
}

void write_func_prototype(char *dst, const func_data_t *func) {
  strcat(dst, "INTERNAL ");

  write_data_type(dst, &(func->func.ret_type));

  strcat(dst, " ");
  strcat(dst, func->func.name);
  strcat(dst, "(");

  for (u64 i = 0; i < func->arg_count; ++i) {
    if (i > 0) {
      strcat(dst, ", ");
    }

    const identifier_t *arg = &(func->func.args[i]);

    write_identifier(dst, arg);
  }

  strcat(dst, ")");
}

void write_func_declaration(arena_t *arena, FILE *fp, const func_data_t *func) {
  char *func_declaration = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!func_declaration) {
    return;
  }

  write_func_prototype(func_declaration, func);

  strcat(func_declaration, ";\n\n");

  u64 length = strlen(func_declaration);

  fwrite(func_declaration, length, 1, fp);

  arena_clear(arena);
}

void write_func_defintion(arena_t *arena, FILE *fp, const func_data_t *func) {
  char *func_definition = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!func_definition) {
    return;
  }

  write_func_prototype(func_definition, func);

  strcat(func_definition, " {\n");

  strcat(func_definition, func->func.body);

  strcat(func_definition, "}\n\n");

  u64 length = strlen(func_definition);

  fwrite(func_definition, length, 1, fp);

  arena_clear(arena);
}

void write_macro_func(arena_t *arena, FILE *fp, const func_data_t *func) {
  char *func_definition = (char *)arena_alloc(arena, LONG_BUFFER_LENGTH);
  if (!func_definition) {
    return;
  }

  strcat(func_definition, "#define ");
  strcat(func_definition, func->base_name);
  strcat(func_definition, "(T");

  for (u64 i = 0; i < func->arg_count; ++i) {
    strcat(func_definition, ", ");
    strcat(func_definition, func->arg_names[i]);
  }

  strcat(func_definition, ") ");
  strcat(func_definition, func->base_name);
  strcat(func_definition, "_##T(");

  for (u64 i = 0; i < func->arg_count; ++i) {
    if (i > 0) {
      strcat(func_definition, ", ");
    }

    strcat(func_definition, func->arg_names[i]);
  }

  strcat(func_definition, ")\n\n");

  u64 length = strlen(func_definition);

  fwrite(func_definition, length, 1, fp);

  arena_clear(arena);
}

void write_vec_func_implementation(arena_t *arena, func_data_t *func,
                                   const struct_def_t *s, char *body,
                                   const char *operator, const char * separator,
                                   const char *end) {
  for (u64 l = 0; l < s->mem_count; ++l) {
    for (u64 m = 0; m < func->arg_count; ++m) {
      if (m > 0) {
        strcat(body, " ");
        strcat(body, operator);
      }

      char *arg = (char *)arena_alloc(arena, SHORT_BUFFER_LENGTH);
      sprintf(arg, " %s.%s", func->arg_names[m], s->members[l].name);

      strcat(body, arg);
    }

    strcat(body, separator);
  }

  strcat(body, end);

  func->func.body = body;
}