#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "common.h"
#include "compiler.h"
#include "debug.h"
#include "object.h"
#include "memory.h"
#include "table.h"
#include "vm.h"

#define READ_BYTE() (*vm.ip++)
#define READ_CONSTANT() (vm.chunk->constants.values[READ_BYTE()])
#define READ_SHORT() (vm.ip += 2, (uint16_t)((vm.ip[-1] << 8) | vm.ip[-2]))
#define READ_3BYTE() (vm.ip += 3, (uint32_t)((vm.ip[-1] << 16)| (vm.ip[-2] << 8) | vm.ip[-3]))
#define READ_STRING() AS_STRING(READ_CONSTANT())
#define READ_STRING_LONG() AS_STRING(readConstantLong())

#define BINARY_OP(valueType, op) \
    do { \
      if (!IS_NUMBER(peek(0)) || !IS_NUMBER(peek(1))) { \
        runtimeError("Operands must be numbers."); \
        return INTERPRET_RUNTIME_ERROR; \
      } \
      \
      double b = AS_NUMBER(pop()); \
      double a = AS_NUMBER(pop()); \
      push(valueType(a op b)); \
    } while (false)

VM vm;

static void resetStack() {
  vm.sp = 0;
  vm.capacity = 0;
  vm.stack = NULL;
}

static void runtimeError(const char* format, ...) {
  va_list args;
  va_start(args, format);
  vfprintf(stderr, format, args);
  va_end(args);
  fputs("\n", stderr);

  size_t instruction = vm.ip - vm.chunk->code;
  fprintf(stderr, "[line %d] in script\n", vm.chunk->lines[instruction]);

  resetStack();
}

void initVM() {
  resetStack();
  vm.objects = NULL;

  initTable(&vm.globals);
  initTable(&vm.strings);
}

void freeVM() {
  freeTable(&vm.globals);
  freeTable(&vm.strings);
  freeObjects();
}

void push(Value value) {
  if (vm.sp == vm.capacity) {
    int oldCapacity = vm.capacity;
    vm.capacity = GROW_CAPACITY(oldCapacity);
    vm.stack = GROW_ARRAY(vm.stack, Value, oldCapacity, vm.capacity);
  }
  vm.stack[vm.sp] = value;
  vm.sp++;
}

Value pop() {
  vm.sp--;
  return vm.stack[vm.sp];
}

Value peek(int distance) {
  return vm.stack[vm.sp - 1 - distance];
}

static bool isFalsey(Value value) {
  return IS_NIL(value) || (IS_BOOL(value) && !AS_BOOL(value));
}

static void concatenate() {
  ObjString* b = AS_STRING(pop());
  ObjString* a = AS_STRING(pop());

  int length = a->length + b->length;
  char* chars = ALLOCATE(char, length + 1);
  memcpy(chars, a->chars, a->length);
  memcpy(chars + a->length, b->chars, b->length);
  chars[length] = '\0';

  ObjString* result = takeString(chars, length);
  push(OBJ_VAL(result));
}

static inline Value readConstantLong() {
  uint8_t constant_lo = READ_BYTE();
  uint8_t constant_mid = READ_BYTE();
  uint8_t constant_hi = READ_BYTE();
  int constant = constant_lo | (constant_mid << 8) | constant_hi << 16;
  return vm.chunk->constants.values[constant];
}

static bool check_array_index(Value index, int* idx) {
  if (!IS_NUMBER(index)) {
    runtimeError("Array index must be a non-negative integer.");
    return false;
  }
  double i_double = AS_NUMBER(index);
  if (i_double < 0) {
    runtimeError("Array index must be a non-negative integer.");
    return false;
  }
  if ((int) i_double != i_double) {
    runtimeError("Array index must be a non-negative integer.");
    return false;
  }
  *idx = (int) i_double;
  return true;
}

static bool vm_index(Value index, Value val) {
  if (IS_OBJ(val)) {
    Obj* obj = AS_OBJ(val);
    switch (obj->type) {
    case OBJ_ARRAY: {
      ValueArray* valArray = AS_VARRAY(val);
      int i;
      if (!check_array_index(index, &i)) {
        return false;
      }
      if (i>=valArray->count) {
        runtimeError("Array index out of bounds");
        return false;
      }
      push(valArray->values[i]);
      break;
    }
    case OBJ_HASH: {
      Table* hashTable = AS_HASH(val);
      Value* value = malloc(sizeof(Value));
      if (tableGet(hashTable, index, value)) {
        push(*value);
      } else {
        push(NIL_VAL);
      }
      break;
    }
    default:
      runtimeError("Cannot index value.");
      return false;
    }
  } else {
    runtimeError("Cannot index value.");
    return false;
  }
  return true;
}

static InterpretResult run() {
  for (;;) {
#ifdef DEBUG_TRACE_EXECUTION
    printf("          ");
    for (int i = 0; i < vm.sp; i++) {
      printf("[ ");
      printValue(vm.stack[i]);
      printf(" ]");
    }
    printf("\n");
    disassembleInstruction(vm.chunk, (int) (vm.ip - vm.chunk->code));
#endif
    uint8_t instruction;
    switch (instruction = READ_BYTE()) {
    case OP_ARRAY: {
      int num_values = READ_BYTE();
      Value* values = malloc(sizeof(Value) * num_values);
      for (int i = num_values - 1; i >= 0; i--) {
        values[i] = pop();
      }
      ValueArray* valArray = malloc(sizeof(ValueArray));
      initValueArray(valArray);
      for (int i = 0; i < num_values; i++) {
        writeValueArray(valArray, values[i]);
      }
      ObjArray* array = takeArray(valArray);
      push(OBJ_VAL(array));
      free(values);
      break;
    }
    case OP_HASH: {
      int num_values = READ_BYTE();
      Value* values = malloc(sizeof(Value) * num_values);
      Value* keys = malloc(sizeof(Value) * num_values);
      for (int i = num_values - 1; i >= 0; i--) {
        values[i] = pop();
        keys[i] = pop();
      }
      Table* hashTable = malloc(sizeof(Table));
      initTable(hashTable);
      for (int i = 0; i < num_values; i++) {
        tableSet(hashTable, keys[i], values[i]);
      }
      ObjHash* hash = takeHash(hashTable);
      push(OBJ_VAL(hash));
      free(values);
      free(keys);
      break;
    }
    case OP_CONSTANT: {
      Value constant = READ_CONSTANT();
      push(constant);
      break;
    }
    case OP_CONSTANT_LONG: {
      Value constant = readConstantLong();
      push(constant);
      break;
    }
    case OP_NIL:
      push(NIL_VAL);
      break;
    case OP_TRUE:
      push(BOOL_VAL(true));
      break;
    case OP_FALSE:
      push(BOOL_VAL(false));
      break;
    case OP_POP:
      pop();
      break;
    case OP_INDEX: {
      Value index = pop();
      Value val = pop();
      bool res = vm_index(index, val);
      if (!res) return INTERPRET_RUNTIME_ERROR;
      break;
    }
    case OP_INDEX_FLIPPED: {
      Value val = pop();
      Value index = pop();
      bool res = vm_index(index, val);
      if (!res) return INTERPRET_RUNTIME_ERROR;
      break;
    }
    case OP_SET_INDEX: {
      Value val = pop();
      Value set_val = pop();
      Value index = pop();
      if (IS_OBJ(val)) {
        Obj* obj = AS_OBJ(val);
        switch (obj->type) {
        case OBJ_ARRAY: {
          ValueArray* valArray = AS_VARRAY(val);
          int i;
          if (!check_array_index(index, &i)) {
            printf("ERR\n");
            return INTERPRET_RUNTIME_ERROR;
          }
          if (i>=valArray->count) {
            int num_nils=i-valArray->count;
            for (;num_nils>0;num_nils--) {
              writeValueArray(valArray,NIL_VAL);
            }
            writeValueArray(valArray,set_val);
          } else {
            valArray->values[i]=set_val;
          }
          break;
        }
        case OBJ_HASH: {
          Table* hashTable = AS_HASH(val);
          tableSet(hashTable, index, set_val);
          break;
        }
        default:
          runtimeError("Cannot index value.");
          return INTERPRET_RUNTIME_ERROR;
        }
      } else {
        runtimeError("Cannot index value.");
        return INTERPRET_RUNTIME_ERROR;
      }
      push(set_val);
      break;
    }
    case OP_GET_LOCAL: {
      uint8_t slot = READ_BYTE();
      push(vm.stack[slot]);
      break;
    }
    case OP_GET_LOCAL_LONG: {
      uint32_t slot = READ_3BYTE();
      push(vm.stack[slot]);
      break;
    }
    case OP_SET_LOCAL: {
      uint8_t slot = READ_BYTE();
      vm.stack[slot] = peek(0);
      break;
    }
    case OP_SET_LOCAL_LONG: {
      uint32_t slot = READ_3BYTE();
      vm.stack[slot] = peek(0);
      break;
    }
    case OP_GET_GLOBAL: {
      ObjString* name = READ_STRING();
      Value value;
      if (!tableGet(&vm.globals, OBJ_VAL(name), &value)) {
        runtimeError("Undefined variable '%s'.", name->chars);
        return INTERPRET_RUNTIME_ERROR;
      }
      push(value);
      break;
    }
    case OP_GET_GLOBAL_LONG: {
      ObjString* name = READ_STRING_LONG();
      Value value;
      if (!tableGet(&vm.globals, OBJ_VAL(name), &value)) {
        runtimeError("Undefined variable '%s'.", name->chars);
        return INTERPRET_RUNTIME_ERROR;
      }
      push(value);
      break;
    }
    case OP_DEFINE_GLOBAL: {
      ObjString* name = READ_STRING();
      tableSet(&vm.globals, OBJ_VAL(name), peek(0));
      pop();
      break;
    }
    case OP_DEFINE_GLOBAL_LONG: {
      ObjString* name = READ_STRING_LONG();
      tableSet(&vm.globals, OBJ_VAL(name), peek(0));
      pop();
      break;
    }
    case OP_SET_GLOBAL: {
      ObjString* name = READ_STRING();
      if (tableSet(&vm.globals, OBJ_VAL(name), peek(0))) {
        tableDelete(&vm.globals, OBJ_VAL(name));
        runtimeError("Undefined variable '%s'.", name->chars);
        return INTERPRET_RUNTIME_ERROR;
      }
      break;
    }
    case OP_SET_GLOBAL_LONG: {
      ObjString* name = READ_STRING_LONG();
      if (tableSet(&vm.globals, OBJ_VAL(name), peek(0))) {
        tableDelete(&vm.globals, OBJ_VAL(name));
        runtimeError("Undefined variable '%s'.", name->chars);
        return INTERPRET_RUNTIME_ERROR;
      }
      break;
    }
    case OP_EQUAL: {
      Value b = pop();
      Value a = pop();
      push(BOOL_VAL(valuesEqual(a, b)));
      break;
    }
    case OP_GREATER:
      BINARY_OP(BOOL_VAL, >);
      break;
    case OP_LESS:
      BINARY_OP(BOOL_VAL, <);
      break;
    case OP_ADD: {
      if (IS_STRING(peek(0)) && IS_STRING(peek(1))) {
        concatenate();
      } else if (IS_NUMBER(peek(0)) && IS_NUMBER(peek(1))) {
        double b = AS_NUMBER(pop());
        double a = AS_NUMBER(pop());
        push(NUMBER_VAL(a + b));
      } else {
        runtimeError("Operands must be two numbers or two strings.");
        return INTERPRET_RUNTIME_ERROR;
      }
      break;
    }
    case OP_SUBTRACT:
      BINARY_OP(NUMBER_VAL, -);
      break;
    case OP_MULTIPLY:
      BINARY_OP(NUMBER_VAL, *);
      break;
    case OP_DIVIDE:
      BINARY_OP(NUMBER_VAL, /);
      break;
    case OP_NOT:
      push(BOOL_VAL(isFalsey(pop())));
      break;
    case OP_NEGATE:
      if (!IS_NUMBER(peek(0))) {
        runtimeError("Operand must be a number.");
        return INTERPRET_RUNTIME_ERROR;
      }

      push(NUMBER_VAL(-AS_NUMBER(pop())));
      break;
    case OP_PRINT: {
      printValue(pop());
      printf("\n");
      break;
    }
    case OP_JUMP: {
      uint16_t offset = READ_SHORT();
      vm.ip += offset;
      break;
    }
    case OP_JUMP_IF_FALSE: {
      uint16_t offset = READ_SHORT();
      if (isFalsey(peek(0))) vm.ip += offset;
      break;
    }
    case OP_LOOP: {
      uint16_t offset = READ_SHORT();
      vm.ip -= offset;
      break;
    }
    case OP_RETURN:
      return INTERPRET_OK;
      break;
    }
  }
  return INTERPRET_OK;
}

InterpretResult interpret(const char* source, bool repl) {
  Chunk chunk;
  initChunk(&chunk);

  if (!compile(source, &chunk, repl)) {
    freeChunk(&chunk);
    return INTERPRET_COMPILE_ERROR;
  }

  vm.chunk = &chunk;
  vm.ip = vm.chunk->code;

  InterpretResult result = run();

  freeChunk(&chunk);
  return result;
}