#include "vm.h"

#include <assert.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

// TODO: Make all these arguments of vm_new().

// Program's main memory stack size.
#define PROGRAM_STACK_SIZE (64 * 1024)

// Locals stack size.
#define LOCALS_STACK_SIZE 1024

// Frame stack size.
#define FRAME_STACK_SIZE (16 * 1024)

// Block stack size.
#define BLOCK_STACK_SIZE 1024

#define IMPLICIT_LABEL -1

// Locals index.
typedef size_t Index;

// Bools are internally I32s. 0 = false, non-zero = true.
static Type     Bool = I32;
typedef int32_t bool_t;

/// Function frame.
///
/// Every Frame implicitly starts a Block. The function's locals are inside this
/// implicit block.
typedef struct Frame {
  Label label;
} Frame;

/// A block of code, used for control flow.
///
/// Blocks have a label that the machine can jump to. Jumps are constrained to
/// the block labels that are in scope.
///
/// Block termination automatically frees the block's locals.
typedef struct Block {
  Label  label;
  size_t addr;         // Address (saved program counter) of the block.
  size_t locals_start; // Offset into the locals stack.
  size_t locals_size;  // Total size in bytes of local variables.
} Block;

typedef struct Vm {
  struct {
    bool exit : 1;
  } flags;
  int32_t  exit_code;
  size_t   pc;     // Program instruction counter.
  size_t   sp;     // Program stack pointer. Points to next available slot.
  size_t   lsp;    // Locals stack pointer. Points to next available slot.
  size_t   fsp;    // Frame stack pointer.
  size_t   bsp;    // Block stack pointer. Points to current Block.
  uint8_t* stack;  // Program stack. Program's main memory.
  uint8_t* locals; // Locals stack. Stores locals for each Block.
  Frame*   frames; // Frame stack for function calls.
  Block*   blocks; // Block stack for control flow.
} Vm;

Vm* vm_new() {
  Vm* vm = calloc(1, sizeof(Vm));
  if (!vm) {
    goto cleanup;
  }
  if (!(vm->stack = calloc(1, PROGRAM_STACK_SIZE))) {
    goto cleanup;
  }
  if (!(vm->locals = calloc(1, LOCALS_STACK_SIZE))) {
    goto cleanup;
  }
  if (!(vm->frames = calloc(FRAME_STACK_SIZE, sizeof(Frame)))) {
    goto cleanup;
  }
  if (!(vm->blocks = calloc(BLOCK_STACK_SIZE, sizeof(Block)))) {
    goto cleanup;
  }
  return vm;

cleanup:
  vm_del(&vm);
  return 0;
}

void vm_del(Vm** pVm) {
  if (pVm && *pVm) {
    Vm* vm = *pVm;
    if (vm->stack) {
      free(vm->stack);
    }
    if (vm->locals) {
      free(vm->locals);
    }
    if (vm->frames) {
      free(vm->frames);
    }
    if (vm->blocks) {
      free(vm->blocks);
    }
    free(vm);
    *pVm = 0;
  }
}

// static size_t get_size(Type type) {
//   switch (type) {
//   case I32:
//     return sizeof(int32_t);
//   }
//   assert(false);
//   return 0;
// }

// TODO: Not used?
#define VM_ASSERT(expr) \
  {}

#define TOP(vm, type)              \
  (assert(vm->sp >= sizeof(type)), \
   *((const type*)&vm->stack[vm->sp - sizeof(type)]))

#define PUSH(vm, value, type)                          \
  assert(vm->sp + sizeof(type) <= PROGRAM_STACK_SIZE); \
  *((type*)(&vm->stack[vm->sp])) = value;              \
  vm->sp += sizeof(type);

#define POP(vm, type)                                      \
  (assert(vm->sp >= sizeof(type)), vm->sp -= sizeof(type), \
   *((type*)(&vm->stack[vm->sp])))

#define BLOCK_PUSH(vm, block)         \
  assert(vm->bsp < BLOCK_STACK_SIZE); \
  vm->blocks[++vm->bsp] = block;

#define BLOCK_POP(vm)                            \
  assert(vm->bsp > 0);                           \
  vm->locals -= vm->blocks[vm->bsp].locals_size; \
  vm->bsp--;

#define PUSH_LOCAL(vm, type)                           \
  assert(vm->lsp + sizeof(type) <= LOCALS_STACK_SIZE); \
  /* Auto-initialize locals to 0. */                   \
  *((type*)(&vm->locals[vm->lsp])) = 0;                \
  vm->lsp += sizeof(type);                             \
  vm->blocks[vm->bsp].locals_size += sizeof(type);

#define POP_LOCAL(vm, type)                                  \
  (assert(vm->lsp >= sizeof(type)), vm->lsp -= sizeof(type), \
   vm->blocks[vm->bsp].locals -= sizeof(type),               \
   *((type*)(&vm->locals[vm->lsp])))

// TODO: Should be an offset from the current frame, not block.
#define GET_LOCAL_PTR(vm, idx, type) \
  ((const type*)(&vm->locals[vm->blocks[vm->bsp].locals_start + idx]))
// TODO: Same here.
#define GET_LOCAL_PTR_MUT(vm, idx, type) \
  ((type*)(&vm->locals[vm->blocks[vm->bsp].locals_start + idx]))

#define LOCAL_RD(vm, idx, type) (*GET_LOCAL_PTR(vm, idx, type))

#define LOCAL_WR(vm, idx, val, type) (*GET_LOCAL_PTR_MUT(vm, idx, type) = val)

static void push(Vm* vm, Inst inst) {
  switch (inst.type) {
  case I32:
    PUSH(vm, inst.payload.i32, int32_t);
    break;
  case F32:
    PUSH(vm, inst.payload.f32, float);
    break;
  }
}

static Value pop(Vm* vm, Type type) {
  Value val;
  switch (type) {
  case I32:
    val.i32 = POP(vm, int32_t);
    break;
  case F32:
    val.f32 = POP(vm, float);
    break;
  }
  return val;
}

static void vm_exit(Vm* vm, Inst inst) {
  vm->exit_code  = vm->sp == 0 ? 0 : POP(vm, int32_t);
  vm->flags.exit = true;
}

#define ADD(vm, a, b, field, type)         \
  {                                        \
    type result = ((a.field) + (b.field)); \
    PUSH(vm, result, type);                \
  }

static void add(Vm* vm, Type type) {
  Value opr1 = pop(vm, type);
  Value opr2 = pop(vm, type);
  switch (type) {
  case I32:
    ADD(vm, opr1, opr2, i32, int32_t);
    break;

  case F32:
    ADD(vm, opr1, opr2, f32, float);
    break;
  }
}

static void dec(Vm* vm, Type type) {
  Value top = pop(vm, type);
  switch (type) {
  case I32:
    PUSH(vm, top.i32 - 1, int32_t);
    break;
  case F32:
    PUSH(vm, top.f32 - 1.0f, float);
    break;
  }
}

static void empty(Vm* vm) { PUSH(vm, vm->sp == 0, bool_t); }

#define CMP(vm, val, type) (POP(vm, type) == val)

static void cmp(Vm* vm, Inst inst) {
  switch (inst.type) {
  case I32:
    PUSH(vm, CMP(vm, inst.payload.i32, int32_t), bool_t);
    break;
  case F32:
    PUSH(vm, CMP(vm, inst.payload.f32, float), bool_t);
    break;
  }
}

static void end(Vm* vm) { BLOCK_POP(vm); }

static void loop(Vm* vm, Inst inst) {
  const Block block = (Block){.label = inst.payload.i32, .addr = vm->pc};
  BLOCK_PUSH(vm, block);
}

static void br(Vm* vm, Inst inst) {
  const Branch  branch         = inst.payload.branch;
  const int32_t label          = branch.label;
  const bool    value          = branch.expected;
  const bool    is_conditional = branch.conditional;
  bool should_branch = is_conditional ? POP(vm, bool_t) == value : true;
  // printf("is conditional: %d\n", is_conditional);
  // printf("value: %d\n", value);
  // printf("should branch: %d\n", should_branch);
  if (should_branch) {
    while (vm->bsp > 0) {
      const Block block = vm->blocks[vm->bsp];
      if (block.label == label) {
        vm->pc = block.addr;
        vm->pc--; // Account for increment at every step of the VM loop.
        return;
      }
      vm->bsp--;
    }
    // We should be able to find the label in the block stack.
    assert(false);
  }
}

static void vm_break(Vm* vm, Inst inst) {
  // TODO.
  // Step over instructions until an End instruction is found.
}

static void local(Vm* vm, Type type) {
  switch (type) {
  case I32:
    PUSH_LOCAL(vm, int32_t);
    break;
  case F32:
    PUSH_LOCAL(vm, float);
    break;
  }
}

static void local_rd(Vm* vm, Inst inst) {
  const Index idx = (Index)inst.payload.u64;
  switch (inst.type) {
  case I32:
    PUSH(vm, LOCAL_RD(vm, idx, int32_t), int32_t);
    break;
  case F32:
    PUSH(vm, LOCAL_RD(vm, idx, float), float);
    break;
  }
}

static void local_wr(Vm* vm, Inst inst) {
  const Index idx = (Index)inst.payload.u64;
  const Value top = pop(vm, inst.type);
  switch (inst.type) {
  case I32:
    LOCAL_WR(vm, idx, top.i32, int32_t);
    break;
  case F32:
    LOCAL_WR(vm, idx, top.f32, float);
    break;
  }
}

static void exec(Vm* vm, Inst inst) {
  switch (inst.op) {
  case Exit:
    vm_exit(vm, inst);
    break;
  case Push:
    push(vm, inst);
    break;
  case Pop:
    pop(vm, inst.type);
    break;
  case Add:
    add(vm, inst.type);
    break;
  case Sub:
    break;
  case Mul:
    break;
  case Div:
    break;
  case Dec:
    dec(vm, inst.type);
    break;
  case Empty:
    empty(vm);
    break;
  case Cmp:
    cmp(vm, inst);
    break;
  case End:
    end(vm);
    break;
  case Break:
    vm_break(vm, inst);
    break;
  case Loop:
    loop(vm, inst);
    break;
  case Br:
    br(vm, inst);
    break;
  case Func: // TODO
    break;
  case Arg: // TODO
    break;
  case Call: // TODO
    break;
  case Local:
    local(vm, inst.type);
    break;
  case LocalRd:
    local_rd(vm, inst);
    break;
  case LocalWr:
    local_wr(vm, inst);
    break;
  }
}

static void init(Vm* vm) {
  // Create an implicit frame for the start of the program.
  vm->frames[0] = (Frame){.label = IMPLICIT_LABEL};
  vm->blocks[0] = (Block){.label = IMPLICIT_LABEL};
  // TODO: Reset all Vm state.
}

int vm_run(Vm* vm, const Inst instructions[], size_t count) {
  assert(vm);
  init(vm);
  for (vm->pc = 0; !vm->flags.exit && vm->pc < count; vm->pc++) {
    const Inst inst = instructions[vm->pc];
    exec(vm, inst);
  }
  // printf("exit code: %d\n", vm->exit_code);
  return vm->exit_code;
}

void vm_print_stack(const Vm* vm) {
  printf("stack start\n");
  for (size_t i = 0; i < vm->sp; ++i) {
    const char sep = (i == vm->sp - 1) ? '\n' : ' ';
    printf("%x%c", vm->stack[i], sep);
  }
  printf("stack end\n");
}