#include "compiler.h"
#include "../../vm/common.h"
#include "../../vm/libc.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

NamesTable *names_table_init() {
  NamesTable *table = malloc(sizeof(NamesTable));
  table->names = malloc(16 * sizeof(char *));
  table->count = 0;
  table->capacity = 16;
  return table;
}

FunctionTable *function_table_init() {
  FunctionTable *table = malloc(sizeof(FunctionTable));
  table->symbols = malloc(16 * sizeof(FunctionDef));
  table->count = 0;
  table->capacity = 16;
  return table;
}

ArrayTable *array_table_init() {
  ArrayTable *table = malloc(sizeof(ArrayTable));
  table->symbols = malloc(16 * sizeof(ArrayDef));
  table->count = 0;
  table->capacity = 16;
  return table;
}

PlexTable *plex_table_init() {
  PlexTable *table = malloc(sizeof(PlexTable));
  table->symbols = malloc(16 * sizeof(PlexDef));
  table->count = 0;
  table->capacity = 16;
  return table;
}

PlexFieldsTable *plex_fields_table_init() {
  PlexFieldsTable *table = malloc(sizeof(PlexFieldsTable));
  table->plex_refs = malloc(64 * sizeof(u32));
  table->fields = malloc(64 * sizeof(ValueType));
  table->count = 0;
  table->capacity = 64;
  return table;
}

u32 names_table_add(NamesTable *table, const char *name) {
  for (u32 i = 0; i < table->count; i++) {
    if (strcmp(table->names[i], name) == 0) {
      return (u32)i;
    }
  }

  if (table->count >= table->capacity) {
    table->capacity *= 2;
    table->names = realloc(table->names, table->capacity * sizeof(char *));
  }

  table->names[table->count] = malloc(strlen(name) + 1);
  strcpy(table->names[table->count], name);
  u32 index = (u32)table->count;
  table->count++;
  return index;
}

u32 function_table_add(FunctionTable *table, FunctionDef def) {
  if (table->count >= table->capacity) {
    table->capacity *= 2;
    table->symbols =
        realloc(table->symbols, table->capacity * sizeof(FunctionDef));
  }

  table->symbols[table->count] = def;
  u32 index = (u32)table->count;
  table->count++;
  return index;
}

u32 array_table_add(ArrayTable *table, ArrayDef def) {
  if (table->count >= table->capacity) {
    table->capacity *= 2;
    table->symbols = realloc(table->symbols, table->capacity * sizeof(ArrayDef));
  }

  table->symbols[table->count] = def;
  u32 index = (u32)table->count;
  table->count++;
  return index;
}

u32 plex_add(PlexTable *plex_table, u32 name, u32 size, u32 field_start,
             u32 field_count) {
  if (plex_table->count >= plex_table->capacity) {
    plex_table->capacity *= 2;
    plex_table->symbols =
        realloc(plex_table->symbols, plex_table->capacity * sizeof(PlexDef));
  }

  plex_table->symbols[plex_table->count].name = name;
  plex_table->symbols[plex_table->count].size = size;
  plex_table->symbols[plex_table->count].field_ref_start = field_start;
  plex_table->symbols[plex_table->count].field_count = field_count;

  u32 index = (u32)plex_table->count;
  plex_table->count++;
  return index;
}

u32 plex_fields_add(PlexFieldsTable *fields_table, u32 plex_ref,
                    ValueType field) {
  if (fields_table->count + 1 > fields_table->capacity) {
    u32 new_capacity = fields_table->capacity * 2;
    if (new_capacity < fields_table->count + 1) {
      new_capacity = fields_table->count + 1;
    }
    fields_table->plex_refs =
        realloc(fields_table->plex_refs, new_capacity * sizeof(u32));
    fields_table->fields =
        realloc(fields_table->fields, new_capacity * sizeof(ValueType));
    fields_table->capacity = new_capacity;
  }

  u32 start_index = fields_table->count;
  fields_table->plex_refs[start_index] = plex_ref;
  fields_table->fields[start_index] = field;
  fields_table->count++;
  return start_index;
}

int plex_get_field_index_by_name(PlexTable *plex_table,
                                 PlexFieldsTable *fields_table,
                                 NamesTable *names_table, u32 plex_index,
                                 const char *field_name) {
  if (plex_index >= plex_table->count)
    return -1;

  PlexDef *plex_def = &plex_table->symbols[plex_index];
  u32 field_start = plex_def->field_ref_start;
  u32 field_count = plex_def->field_count;

  for (u32 i = 0; i < field_count; i++) {
    u32 field_table_index = field_start + i;
    ValueType *field = &fields_table->fields[field_table_index];

    if (field->name < names_table->count) {
      if (strcmp(names_table->names[field->name], field_name) == 0) {
        return (int)i; // Return field index within the plex
      }
    }
  }
  return -1; // Not found
}

ValueType *plex_get_field(PlexTable *plex_table, PlexFieldsTable *fields_table,
                          u32 plex_index, u32 field_in_plex_index) {
  if (plex_index >= plex_table->count)
    return nil;

  PlexDef *plex_def = &plex_table->symbols[plex_index];
  if (field_in_plex_index >= plex_def->field_count)
    return nil;

  u32 field_table_index = plex_def->field_ref_start + field_in_plex_index;
  return &fields_table->fields[field_table_index];
}

ValueType *plex_get_field_by_name(PlexTable *plex_table,
                                  PlexFieldsTable *fields_table,
                                  NamesTable *names_table, u32 plex_index,
                                  const char *field_name) {
  int field_index = plex_get_field_index_by_name(
      plex_table, fields_table, names_table, plex_index, field_name);
  if (field_index == -1)
    return nil;

  return plex_get_field(plex_table, fields_table, plex_index, (u32)field_index);
}

Symbol *global(VM *vm) {
  Symbol s;
  ValueType t;

  s.ref.global = vm->mp;

  Token token_type = nextToken();
  Token array_or_eq = nextToken();
  if (array_or_eq.type == TOKEN_LBRACKET) {
    Token rb = nextToken();
    if (rb.type != TOKEN_RBRACKET)
      return nil;

    Token eq = nextToken();
    if (eq.type != TOKEN_EQ)
      return nil;

    t.type = ARRAY;
    ValueType array_type;

    switch (token_type.type) {
    case TOKEN_TYPE_I8:
      array_type.type = I8;
      break;
    case TOKEN_TYPE_I16:
      array_type.type = I16;
      break;
    case TOKEN_TYPE_INT:
      array_type.type = I32;
      break;
    case TOKEN_TYPE_U8:
      array_type.type = U8;
      break;
    case TOKEN_TYPE_U16:
      array_type.type = U16;
      break;
    case TOKEN_TYPE_NAT:
      array_type.type = U32;
      break;
    case TOKEN_TYPE_REAL:
      array_type.type = F32;
      break;
    case TOKEN_TYPE_STR:
      array_type.type = STR;
      break;
    case TOKEN_IDENTIFIER:
      break;
    default:
      return nil;
    }

  } else {
    // its not an array, so should be =
    if (array_or_eq.type != TOKEN_EQ)
      return nil;

    switch (token_type.type) {
    case TOKEN_TYPE_I8:
      t.type = I8;
      break;
    case TOKEN_TYPE_I16:
      t.type = I16;
      break;
    case TOKEN_TYPE_INT:
      t.type = I32;
      break;
    case TOKEN_TYPE_U8:
      t.type = U8;
      break;
    case TOKEN_TYPE_U16:
      t.type = U16;
      break;
    case TOKEN_TYPE_NAT:
      t.type = U32;
      break;
    case TOKEN_TYPE_REAL:
      t.type = F32;
      break;
    case TOKEN_TYPE_STR:
      t.type = STR;
      break;
    case TOKEN_IDENTIFIER:
      break;
    default:
      return nil;
    }
  }

  s.type = t;

  Token value = nextToken();

  return nil;
}

typedef struct {
  Token current;
  Token previous;
  bool hadError;
  bool panicMode;
} Parser;

typedef enum {
  PREC_NONE,
  PREC_ASSIGNMENT, /* = */
  PREC_OR,         /* or */
  PREC_AND,        /* and */
  PREC_EQUALITY,   /* == != */
  PREC_COMPARISON, /* < > <= >= */
  PREC_TERM,       /* + - */
  PREC_FACTOR,     /* * / */
  PREC_UNARY,      /* not */
  PREC_CALL,       /* . () */
  PREC_PRIMARY
} Precedence;

typedef void (*ParseFn)(char *program);

typedef struct {
  ParseFn prefix;
  ParseFn infix;
  Precedence precedence;
} ParseRule;

typedef struct {
  SymbolTable table;
  Symbol current;
  Symbol last;
  i8 rp; // Next free register
} Compiler;

Parser parser;

const char *internalErrorMsg =
    "FLAGRANT COMPILER ERROR\n\nCompiler over.\nBug = Very Yes.";

bool isType(TokenType type) {
  return type == TOKEN_TYPE_INT || type == TOKEN_TYPE_NAT ||
         type == TOKEN_TYPE_REAL || type == TOKEN_TYPE_STR ||
         type == TOKEN_TYPE_BOOL;
}

void errorAt(Token *token, const char *message) {
  if (parser.panicMode)
    return;
  parser.panicMode = true;
  fprintf(stderr, "[line %d] Error", token->line);

  if (token->type == TOKEN_EOF) {
    fprintf(stderr, " at end");
  } else if (token->type == TOKEN_ERROR) {
  } else {
    fprintf(stderr, " at '%.*s'", token->length, token->start);
  }

  fprintf(stderr, ": %s\n", message);
  parser.hadError = true;
}

void error(const char *message) { errorAt(&parser.previous, message); }

void errorAtCurrent(const char *message) { errorAt(&parser.current, message); }

void advance() {
  parser.previous = parser.current;

  for (;;) {
    parser.current = nextToken();
    if (parser.current.type != TOKEN_ERROR)
      break;

    errorAtCurrent(parser.current.start);
  }
}

void consume(TokenType type, const char *message) {
  if (parser.current.type == type) {
    advance();
    return;
  }

  errorAtCurrent(message);
}

static int allocateRegister(Compiler *c) {
  char buffer[38];
  if (c->rp + 1 > 31) {
    sprintf(buffer, "Out of registers (used %d, max 32)", c->rp + 1);
    error(buffer);
    return -1;
  }

  return c->rp++;
}

static void popRegister(Compiler *c) {
  if (c->rp - 1 > 0) {
    c->rp--;
  }
}

static void freeRegister(Compiler *c, u8 reg) {
  if (reg == c->rp - 1) {
    c->rp--;
  }
}

static void clearRegisters(Compiler *c, u8 reg) { c->rp = 0; }

void emit_byte(VM *vm, u8 byte) { vm->code[vm->cp++] = byte; }

void emit_u32(VM *vm, u32 value) {
  write_u32(vm, code, vm->cp, value);
  vm->cp += 4;
}

void emit_opcode(VM *vm, Opcode op) { emit_byte(vm, op); }

static bool check(TokenType type) { return parser.current.type == type; }

static bool match(TokenType type) {
  if (!check(type))
    return false;
  advance();
  return true;
}

static void expression(Compiler *c, VM *vm) {
  USED(c);
  USED(vm);
}

void number(Compiler *c, VM *vm) {
  emit_opcode(vm, OP_LOAD_IMM);
  int reg = allocateRegister(c);
  if (reg < 0)
    return;
  emit_byte(vm, reg);

  c->last = Symbol{ .type=parser.previous.type };

  switch (parser.previous.type) {
  case TOKEN_INT_LITERAL: {
    char *endptr;
    i32 value = (i32)strtol(parser.previous.start, &endptr, 10);
    emit_u32(vm, value);
    return;
  }
  case TOKEN_UINT_LITERAL: {
    long value = atol(parser.previous.start);
    emit_u32(vm, value);
    return;
  }
  case TOKEN_FLOAT_LITERAL: {
    float value = atof(parser.previous.start);
    fixed_t fvalue = float_to_fixed(value);
    emit_u32(vm, fvalue);
    return;
  }
  default:
    return; // Unreachable.
  }

  errorAtCurrent("Invalid number format");
}

static void unary(Compiler *c, VM *vm) {
  TokenType operatorType = parser.previous.type;

  // Compile the operand.
  expression(c, vm);

  // Emit the operator instruction.
  switch (operatorType) {
  case TOKEN_MINUS: {
    switch (c->last.type) {
    case TOKEN_UINT_LITERAL:
      emit_opcode(vm, OP_NEG_UINT);
    case TOKEN_FLOAT_LITERAL:
      emit_opcode(vm, OP_NEG_FLOAT);
    default:
      emit_opcode(vm, OP_NEG_INT);
    }

    int dest = allocateRegister();
    emit_byte(vm, dest);
    emit_byte(vm, dest);
  }
  default:
    return; // Unreachable.
  }
}

static void emitHalt(Compiler *c, VM *vm) {
  emit_opcode(vm, OP_HALT);
  advance();
  number(c, vm);
}

static void endCompiler(Compiler *c, VM *vm) { emitHalt(c, vm); }

static void grouping(Compiler *c, VM *vm) {
  expression(c, vm);
  consume(TOKEN_RPAREN, "Expect ')' after expression.");
}

bool compile(const char *source, VM *vm) {
  USED(source);
  USED(vm);
  initLexer(source);

  parser.hadError = false;
  parser.panicMode = false;

  Compiler compiler;
  advance();
  expression(&compiler, vm);
  consume(TOKEN_EOF, "Expect end of expression.");
  endCompiler(&compiler, vm);

  return parser.hadError;
}