#include "vm.h"
#include "lib/commons.h"
#include "lib/mod.h"
#include "lib/type.h"
#include "lib/tysy.h"

void vm_init(vm_t* vm, sym_t* sym, tysy_t* tysy, err_t* err)
{
  assert(vm);
  assert(sym);
  assert(tysy);

  vm->top_frame = NULL;
  vm_frame_push(vm);

  vm->tysy = tysy;
  vm->err = err;
}

void vm_free(vm_t* vm)
{
  assert(vm);

  while (vm_frame_pop(vm))
    {
    }

  vm->top_frame = NULL;
}

void vm_frame_init(frame_t* frame)
{
  assert(frame);

  frame->pc = 0;
  frame->bp = 0;
  frame->sp = 0;

  frame->prev = NULL;
  frame->locals.cap = 0;
  frame->locals.size = 0;
  frame->locals.data = NULL;
}

void vm_frame_free(frame_t* frame)
{
  for (size_t i=0; i<frame->locals.size; i++)
    {
      free(frame->locals.data[i]);
    }

  free(frame->locals.data);
  frame->locals.data = NULL;
  frame->locals.size = 0;
  frame->locals.cap = 0;

  if (frame->prev != NULL)
    {
      vm_frame_free((frame_t*) frame->prev);
      frame->prev = NULL;
    }
}

void vm_frame_push(vm_t* vm)
{
  assert(vm);

  frame_t* frame = malloc(sizeof(frame_t));
  vm_frame_init(frame);

  frame->prev = (struct frame*) vm->top_frame;

  vm->top_frame = frame;
}

int vm_frame_pop(vm_t* vm)
{
  assert(vm);

  if (vm->top_frame)
    {
      frame_t* frame = vm->top_frame;
      vm->top_frame = (frame_t*) vm->top_frame->prev;

      vm_frame_free(frame);
      free(frame);
      return 1;
    }

  return 0;
}

local_t* vm_try_local_load(vm_t* vm, int id)
{
  assert(vm);

  for (size_t i=0; i<vm->top_frame->locals.size; i++)
    {
      if (vm->top_frame->locals.data[i]->id == id)
        {
          return vm->top_frame->locals.data[i];
        }
    }

  return NULL;
}

void vm_local_store(vm_t* vm, int id, param_t addr)
{
  assert(vm);
  local_t* loc = vm_try_local_load(vm, id);

  if (!loc)
    {
      if (vm->top_frame->locals.data == NULL)
        {
          vm->top_frame->locals.cap = 2;
          vm->top_frame->locals.data = malloc(sizeof(local_t*)
                                              * vm->top_frame->locals.cap);
        }

      if (vm->top_frame->locals.size >= vm->top_frame->locals.cap)
        {
          vm->top_frame->locals.cap *= 2;
          vm->top_frame->locals.data = realloc(vm->top_frame->locals.data,
                                               sizeof(local_t*)
                                               * vm->top_frame->locals.cap);
        }

      vm->top_frame->locals.data[vm->top_frame->locals.size]
        = malloc(sizeof(local_t));

      vm->top_frame->locals.data[vm->top_frame->locals.size]->id = id;
      vm->top_frame->locals.data[vm->top_frame->locals.size]->addr = addr;
      vm->top_frame->locals.size++;
    }
  else
    {
      loc->id = id;
      loc->addr = addr;
    }
}

void vm_push_new_value(vm_t* vm, mod_t* mod, value_t* value)
{
  assert(vm);
  assert(value);

  param_t param = mod_push_new_value(mod, value);
  vm_push(vm, param);
}

void vm_push(vm_t* vm, param_t param)
{
  assert(vm);
  assert(vm->top_frame->sp + 1 < RZ_STACK_LIMIT);

  vm->top_frame->stack[vm->top_frame->sp] = param;
  vm->top_frame->sp++;
}

param_t vm_pop(vm_t* vm)
{
  assert(vm);
  assert(vm->top_frame->sp > 0);
  param_t param = vm->top_frame->stack[vm->top_frame->sp - 1];
  vm->top_frame->sp--;

  return param;
}

size_t vm_stack_str(vm_t* vm, char* buffer, size_t size)
{
  assert(vm);
  assert(buffer);

  size_t sz = 0;

  if (vm->top_frame->sp == 0)
    {
      sz += snprintf(buffer + sz, size - sz, "** empty **\n");
      return sz;
    }

  for (size_t i=0; i<vm->top_frame->sp; i++)
    {
      sz += snprintf(buffer + sz, size - sz, "| %d |\n", vm->top_frame->stack[vm->top_frame->sp - 1 - i]);
    }

  return sz;
}

int vm_exec_mod(vm_t* vm, mod_t* mod)
{
  assert(vm);
  vm->top_frame->pc = 0;

  while (vm->top_frame->pc < mod->program.size)
    {
      int status = vm_exec_instr(vm,
                                 mod,
                                 mod->program.ops[vm->top_frame->pc],
                                 mod->program.params[vm->top_frame->pc]);

      if (status != 0)
        {
          return status;
        }
    }

  return 0;
}

int vm_exec_instr(vm_t* vm, mod_t* mod, Opcode op, param_t param)
{
  assert(vm);

  switch (op)
    {
    case OP_PUSH: {
      if (vm->top_frame->sp + 1 >= RZ_STACK_LIMIT)
        {
          fprintf(stderr, "Stack overflow\n");
          abort();
        }

      vm->top_frame->stack[vm->top_frame->sp] = param;
      vm->top_frame->sp++;
      vm->top_frame->pc++;
    } break;

    case OP_ASSERT: {
      param_t top = vm_pop(vm);
      value_t* value = mod->values.data[top];

      if (value->type->kind != TYPE_BOOL)
        {
          vm_ensure_type(vm, value, TYPE_BOOL);

          vm->top_frame->pc++;
          break;
        }

      if (!value->value.bool)
        {
          fprintf(stderr, "\33[31mASSERT\33[0m[:%d] assertion failed\n",
                  value->line);
          return -1;
        }

      vm_push(vm, top);
      vm->top_frame->pc++;

    } break;

    case OP_NE:
    case OP_EQ: {
      int l = vm_pop(vm);
      int r = vm_pop(vm);
      assert(l != -1);
      assert(r != -1);

      value_t* rhs = mod->values.data[l];
      value_t* lhs = mod->values.data[r];

      int same_type = lhs->type->kind == rhs->type->kind;
      int eq = value_eq(lhs, rhs);

      value_t* res = tysy_new_bool(vm->tysy,
                                   (op == OP_EQ
                                    ? (same_type && eq)
                                    : (!same_type || !eq)),
                                   rhs->line);

      vm_push_new_value(vm, mod, res);

      vm->top_frame->pc++;
    } break;

          case OP_LT:
    case OP_LE:
    case OP_GT:
    case OP_GE: {
      int r = vm_pop(vm);
      int l = vm_pop(vm);

      value_t* lhs = mod->values.data[l];
      value_t* rhs = mod->values.data[r];

      vm_ensure_type(vm, lhs, TYPE_NUM);
      vm_ensure_type(vm, rhs, TYPE_NUM);

      int val = 0;

      switch (op)
        {
        case OP_LT: val = lhs->value.num < rhs->value.num; break;
        case OP_LE: val = lhs->value.num <= rhs->value.num; break;
        case OP_GT: val = lhs->value.num > rhs->value.num; break;
        case OP_GE: val = lhs->value.num >= rhs->value.num; break;
        default: break;
        }

      value_t* res = tysy_new_bool(vm->tysy, val, lhs->line);
      vm_push_new_value(vm, mod, res);

      vm->top_frame->pc++;
    } break;

    case OP_ADD:
    case OP_SUB:
    case OP_MUL:
    case OP_DIV:
    case OP_MODULO:
    case OP_POW: {
      int r = vm_pop(vm);
      int l = vm_pop(vm);

      value_t* lhs = mod->values.data[l];
      value_t* rhs = mod->values.data[r];

      vm_ensure_type(vm, lhs, TYPE_NUM);
      vm_ensure_type(vm, rhs, TYPE_NUM);

      double val = 0;

      switch (op)
        {
        case OP_ADD: val = lhs->value.num + rhs->value.num; break;
        case OP_SUB: val = lhs->value.num - rhs->value.num; break;
        case OP_MUL: val = lhs->value.num * rhs->value.num; break;
        case OP_DIV: val = lhs->value.num / rhs->value.num; break;
        case OP_MODULO: val = fmod(lhs->value.num, rhs->value.num); break;
        case OP_POW: val = powf(lhs->value.num, rhs->value.num); break;
        default: assert(0); break;
        }

      value_t* res = tysy_new_num(vm->tysy, val, lhs->line);
      vm_push_new_value(vm, mod, res);

      vm->top_frame->pc++;
    } break;

    case OP_STRCAT: {
      int r = vm_pop(vm);
      int l = vm_pop(vm);

      value_t* lhs = mod->values.data[l];
      value_t* rhs = mod->values.data[r];

      vm_ensure_type(vm, lhs, TYPE_STR);
      vm_ensure_type(vm, rhs, TYPE_STR);

      str_t res;
      str_init(&res);
      str_append(&res, lhs->value.str);
      str_append(&res, rhs->value.str);

      value_t* value = tysy_new_str(vm->tysy, res.data, lhs->line);
      str_free(&res);

      vm_push_new_value(vm, mod, value);

      vm->top_frame->pc++;
    } break;

    case OP_STRDUP: {
      int r = vm_pop(vm);
      int l = vm_pop(vm);

      value_t* lhs = mod->values.data[l];
      value_t* rhs = mod->values.data[r];

      if (lhs->type->kind != TYPE_NUM)
        {
          value_t* tmp = lhs;
          lhs = rhs;
          rhs = tmp;
        }

      vm_ensure_type(vm, lhs, TYPE_NUM);
      vm_ensure_type(vm, rhs, TYPE_STR);

      str_t res;
      str_init(&res);

      for (int i=0; i<(int) lhs->value.num; i++)
        {
          str_append(&res, rhs->value.str);
        }

      value_t* value = tysy_new_str(vm->tysy, res.data, lhs->line);
      str_free(&res);

      vm_push_new_value(vm, mod, value);

      vm->top_frame->pc++;
    } break;

    case OP_USUB: {
      int l = vm_pop(vm);

      value_t* lhs = mod->values.data[l];

      vm_ensure_type(vm, lhs, TYPE_NUM);

      double val = -lhs->value.num;

      value_t* res = tysy_new_num(vm->tysy, val, lhs->line);
      vm_push_new_value(vm, mod, res);

      vm->top_frame->pc++;
    } break;

    case OP_NOT: {
      int l = vm_pop(vm);
      value_t* lhs = mod->values.data[l];
      vm_ensure_type(vm, lhs, TYPE_BOOL);

      int val = !lhs->value.bool;

      value_t* res = tysy_new_bool(vm->tysy, val, lhs->line);
      vm_push_new_value(vm, mod, res);

      vm->top_frame->pc++;
    } break;

    case OP_BRT:
    case OP_BRF: {
      param_t p = vm_pop(vm);
      value_t* value = mod->values.data[p];
      vm_ensure_type(vm, value, TYPE_BOOL);

      if (op == OP_BRT && value->value.bool)
        {
          vm->top_frame->pc = param;
        }
      else if (op == OP_BRF && !value->value.bool)
        {
          vm->top_frame->pc = param;
        }
      else
        {
          vm->top_frame->pc++;
        }

    } break;

    case OP_BR: {
      vm->top_frame->pc = param;
    } break;

    case OP_LOAD: {
      local_t* local = vm_try_local_load(vm, param);

      if (!local)
        {
          char msg[RZ_STR_LIMIT];

          snprintf(msg, RZ_STR_LIMIT, "Cannot load local variable '%d'.",
                   param);

          param_t p = vm->top_frame->stack[vm->top_frame->sp - 1];
          int line = mod->values.data[p]->line;

          err_fatal(vm->err, msg, line);
          err_dump(vm->err);
        }
      else
        {
          vm_push(vm, vm->top_frame->stack[local->addr]);
        }

      vm->top_frame->pc++;
    } break;

    case OP_STORE: {
      assert(vm->top_frame->sp > 0);
      vm_local_store(vm, param, vm->top_frame->sp - 1);
      vm->top_frame->pc++;
    } break;

    default: {
      fprintf(stderr, "Cannot execute unknown opcode '%s'.\n",
              OpcodeStr[op]);
      abort();
    };
    }

  return 0;
}

void vm_ensure_type(vm_t* vm, value_t* value, TypeKind want)
{
  assert(vm);
  assert(value);
  TypeKind got = value->type->kind;
  int line = value->line;

  if (got != want)
    {
      char msg[RZ_STR_LIMIT];
      snprintf(msg, RZ_STR_LIMIT,
               "type mismatch: %s expected, got %s.",
               TypeKindStr[want] + strlen("TYPE_"),
               TypeKindStr[got] + strlen("TYPE_"));

      err_fatal(vm->err, msg, line);
      err_dump(vm->err);
    }
}