#include <common/fzl_builtins.h>
#include <common/fzl_opcodes.h>
#include <common/fzl_types.h>
#include <common/fzl_value.h>
#include <compiler/utils/fzl_const_pool.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vm/fzl_vm.h>
#include <vm/fzl_vm_env.h>
#include <vm/fzl_vm_object.h>

char last_err[256];

static char *expect_hdr = "FZL\0\1";

fzl_vm_t *fzl_mkvm(const char *filepath)
{
  fzl_vm_t *vm = malloc(sizeof(fzl_vm_t));
  if (!vm)
  {
    snprintf(last_err, sizeof(last_err), "failed to allocate vm");
    return NULL;
  }

  vm->consts        = NULL;
  vm->consts_count  = 0;
  vm->ops           = NULL;
  vm->ops_count     = 0;
  vm->ip            = 0;
  vm->dtop          = 0;
  vm->ctop          = 0;
  vm->etop          = 0;
  vm->objects_count = 0;

  memset(vm->dstack, 0, sizeof(vm->dstack));
  memset(vm->cstack, 0, sizeof(vm->cstack));
  memset(vm->estack, 0, sizeof(vm->estack));
  memset(vm->objects, 0, sizeof(vm->objects));

  FILE *fd = fopen(filepath, "rb");
  if (!fd)
  {
    snprintf(last_err, sizeof(last_err), "failed to open file '%s'", filepath);
    return NULL;
  }

  char hdr[5];
  fread(hdr, sizeof(char), 5, fd);
  if (memcmp(hdr, expect_hdr, 5) != 0)
  {
    snprintf(last_err, sizeof(last_err), "magic check error");
    return NULL;
  }

  size_t consts_count, ops_count;
  fread(&consts_count, sizeof(size_t), 1, fd);
  fread(&ops_count, sizeof(size_t), 1, fd);

  if (ops_count < 1)
  {
    snprintf(last_err, sizeof(last_err), "there is nothing to do (0 ops)");
    return NULL;
  }

  if (consts_count > 0)
  {
    vm->consts_count = consts_count;
    vm->consts       = malloc(consts_count * sizeof(fzl_value_t));

    for (size_t i = 0; i < consts_count; i++)
    {
      fzl_consttype_t type;
      fread(&type, sizeof(fzl_consttype_t), 1, fd);

      if (type == FZL_CONSTTYPE_STRING)
      {
        size_t length;
        char *chars;

        fread(&length, sizeof(size_t), 1, fd);

        chars = malloc(length * sizeof(char) + 1);
        fread(chars, sizeof(char), length, fd);
        chars[length] = 0;

        vm->consts[i] = (fzl_const_t){
            .type = type, .as.string.length = length, .as.string.chars = chars};
      }
      else
      {
        snprintf(last_err, sizeof(last_err), "unknown const type at index %ld",
                 i);
        return NULL;
      }
    }
  }

  vm->ops_count = ops_count;
  vm->ops       = malloc(ops_count * sizeof(uint64_t));
  for (size_t i = 0; i < ops_count; i++)
  {
    uint64_t op;
    fread(&op, sizeof(uint64_t), 1, fd);
    vm->ops[i] = op;
  }

  return vm;
}

int fzl_vm_run(fzl_vm_t *vm)
{
  if (!vm)
  {
    snprintf(last_err, sizeof(last_err), "vm is null");
    return -1;
  }

  vm->ip = 0;
  while (vm->ip < vm->ops_count)
  {
    fzl_opcode_t opcode = vm->ops[vm->ip];

    switch (opcode)
    {
      case FZL_OPCODE_ALLOC_OBJECT:
      {
        uint64_t values_count = vm->ops[++vm->ip];
        fzl_vm_object_t *obj  = malloc(sizeof(fzl_vm_object_t));
        obj->values           = malloc(values_count * sizeof(fzl_value_t));
        obj->values_count     = values_count;

        vm->objects[vm->objects_count++] = obj;
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_OBJECT, .as.obj = obj};
        break;
      }

      case FZL_OPCODE_GET_FIELD:
      {
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected value in data stack");
          return -1;
        }

        uint64_t field_slot = vm->ops[++vm->ip];
        fzl_value_t val     = vm->dstack[--vm->dtop];
        if (val.type != FZL_TYPE_OBJECT)
        {
          snprintf(
              last_err, sizeof(last_err),
              "runtime error - cannot access field #%lu of non-object type",
              field_slot);
          return -1;
        }

        fzl_vm_object_t *obj = val.as.obj;
        if (obj->values_count <= field_slot)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - cannot access field #%lu", field_slot);
          return -1;
        }

        vm->dstack[vm->dtop++] = obj->values[field_slot];
        break;
      }

      case FZL_OPCODE_SET_FIELD:
      {
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        uint64_t field_slot = vm->ops[++vm->ip];
        fzl_value_t objval  = vm->dstack[--vm->dtop];
        fzl_value_t wrval   = vm->dstack[--vm->dtop];
        if (objval.type != FZL_TYPE_OBJECT)
        {
          snprintf(
              last_err, sizeof(last_err),
              "runtime error - cannot access field #%lu of non-object type",
              field_slot);
          return -1;
        }

        fzl_vm_object_t *obj = objval.as.obj;
        if (obj->values_count <= field_slot)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - cannot access field #%lu", field_slot);
          return -1;
        }

        obj->values[field_slot] = wrval;
        break;
      }

      case FZL_OPCODE_CALL:
      {
        uint64_t ip            = vm->ops[++vm->ip];
        vm->cstack[vm->ctop++] = vm->ip;
        vm->ip                 = ip - 1;
        break;
      }

      case FZL_OPCODE_RET:
      {
        if (vm->ctop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - no place to jump");
          return -1;
        }

        vm->ip = vm->cstack[--vm->ctop];
        break;
      }

      case FZL_OPCODE_CMP_EQ:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t req = vm->dstack[--vm->dtop];
        fzl_value_t leq = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = leq.as.u == req.as.u};
        break;

      case FZL_OPCODE_CMP_NEQ:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t rneq       = vm->dstack[--vm->dtop];
        fzl_value_t lneq       = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] = (fzl_value_t){.type = FZL_TYPE_BOOL,
                                               .as.b = lneq.as.u != rneq.as.u};
        break;

      case FZL_OPCODE_CMP_LT:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t rlt = vm->dstack[--vm->dtop];
        fzl_value_t llt = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = llt.as.u < rlt.as.u};
        break;

      case FZL_OPCODE_CMP_LTEQ:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t rlteq      = vm->dstack[--vm->dtop];
        fzl_value_t llteq      = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] = (fzl_value_t){
            .type = FZL_TYPE_BOOL, .as.b = llteq.as.u <= rlteq.as.u};
        break;

      case FZL_OPCODE_CMP_GT:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t rgt = vm->dstack[--vm->dtop];
        fzl_value_t lgt = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = lgt.as.u > rgt.as.u};
        break;

      case FZL_OPCODE_CMP_GTEQ:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t rgteq      = vm->dstack[--vm->dtop];
        fzl_value_t lgteq      = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] = (fzl_value_t){
            .type = FZL_TYPE_BOOL, .as.b = lgteq.as.u >= rgteq.as.u};
        break;

      case FZL_OPCODE_NOT:
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected value in data stack");
          return -1;
        }

        fzl_value_t notval = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = !notval.as.b};
        break;

      case FZL_OPCODE_JZ:
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected value in data stack");
          return -1;
        }

        fzl_value_t jzval = vm->dstack[--vm->dtop];
        uint64_t jzip     = vm->ops[++vm->ip] - 1;
        if (jzval.as.b == 0)
          vm->ip = jzip;
        break;

      case FZL_OPCODE_JMP:
        vm->ip = vm->ops[++vm->ip] - 1;
        break;

      case FZL_OPCODE_AND:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t land       = vm->dstack[--vm->dtop];
        fzl_value_t rand       = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] = (fzl_value_t){.type = FZL_TYPE_BOOL,
                                               .as.b = land.as.b && rand.as.b};
        break;

      case FZL_OPCODE_OR:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected two values in data stack");
          return -1;
        }

        fzl_value_t lor = vm->dstack[--vm->dtop];
        fzl_value_t ror = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = lor.as.b || ror.as.b};
        break;

      case FZL_OPCODE_STORE:
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected var name in data stack");
          return -1;
        }

        fzl_value_t varname = vm->dstack[--vm->dtop];
        fzl_env_storevalue(vm->estack[vm->etop - 1], (char *)varname.as.s.value,
                           (fzl_value_t){});

        break;

      case FZL_OPCODE_WRITE:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in data stack");
          return -1;
        }

        fzl_value_t writevar = vm->dstack[--vm->dtop];
        fzl_value_t value    = vm->dstack[--vm->dtop];

        fzl_vm_env_wresult_t wresult = fzl_env_writevalue(
            vm->estack[vm->etop - 1], (char *)writevar.as.s.value, value);
        if (wresult != FZL_VM_ENV_WRESULT_OK)
        {
          switch (wresult)
          {
            case FZL_VM_ENV_WRESULT_ERR:
              snprintf(last_err, sizeof(last_err),
                       "runtime error - failed to write variable");
              return -1;

            case FZL_VM_ENV_WRESULT_ERR_EXISTS:
              snprintf(last_err, sizeof(last_err),
                       "runtime error - redefenition of '%s'",
                       writevar.as.s.value);
              return -1;

            case FZL_VM_ENV_WRESULT_ERR_NOT_FOUND:
              snprintf(last_err, sizeof(last_err),
                       "runtime error - cannot find name '%s'",
                       writevar.as.s.value);
              return -1;

            default:
              return -1;
          }
        }

        // vm->dstack[vm->dtop++] = value;
        // NOTE: add POP opcode to prevent
        // stack overflow!
        break;

      case FZL_OPCODE_MK_ENVTABLE:
        vm->estack[vm->etop++] = fzl_mkenv();
        break;

      case FZL_OPCODE_RM_ENVTABLE:
        fzl_freeenv(vm->estack[--vm->etop]);
        break;

      case FZL_OPCODE_LOOKUP_ENV:
      {
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - LOOKUP_ENV: empty data stack");
          return -1;
        }

        fzl_value_t val = vm->dstack[--vm->dtop];
        fzl_value_t env_val;
        bool result;
        for (ssize_t i = vm->etop - 1; i >= 0; i--)
        {
          if ((result = fzl_env_lookup(vm->estack[i], (char *)val.as.s.value,
                                       &env_val)))
            break;
        }
        vm->dstack[vm->dtop++] = env_val;

        break;
      }

      case FZL_OPCODE_CALL_BUILTIN:
      {
        uint64_t fn_id = vm->ops[++vm->ip];
        if (fn_id == FZL_BUILTIN_PRINTLINE)
        {
          fzl_value_t arg = vm->dstack[--vm->dtop];
          if (arg.type == FZL_TYPE_STRING)
            printf("%s\n", arg.as.s.value);
          else if (arg.type == FZL_TYPE_CHAR)
            printf("%c\n", arg.as.c);
          else if (arg.type == FZL_TYPE_INT)
            printf("%ld\n", arg.as.i);
          else if (arg.type == FZL_TYPE_UINT)
            printf("%lu\n", arg.as.u);
          else if (arg.type == FZL_TYPE_BOOL)
            printf("%b\n", arg.as.b);
          else if (arg.type == FZL_TYPE_NULL)
            printf("<null>\n");
        }
        else
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - undefined builtin %lu", fn_id);
          return -1;
        }
      }
      break;

      case FZL_OPCODE_LOADCONST:
      {
        uint64_t cn = vm->ops[++vm->ip];
        if (vm->consts_count <= cn)
        {
          snprintf(last_err, sizeof(last_err), "runtime error - unknown const");
          return -1;
        }
        fzl_const_t c = vm->consts[cn];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type        = FZL_TYPE_STRING,
                          .as.s.length = c.as.string.length,
                          .as.s.value  = c.as.string.chars};
      }
      break;

      case FZL_OPCODE_ADD:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t add1 = vm->dstack[--vm->dtop];
        fzl_value_t add2 = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = add1.type, .as.u = add1.as.u + add2.as.u};
        break;

      case FZL_OPCODE_SUB:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t sub2 = vm->dstack[--vm->dtop];
        fzl_value_t sub1 = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = sub1.type, .as.u = sub1.as.u - sub2.as.u};
        break;

      case FZL_OPCODE_MUL:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t mul1 = vm->dstack[--vm->dtop];
        fzl_value_t mul2 = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = mul1.type, .as.u = mul1.as.u * mul2.as.u};
        break;

      case FZL_OPCODE_DIV:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t div1 = vm->dstack[--vm->dtop];
        fzl_value_t div2 = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = div1.type, .as.u = div2.as.u / div1.as.u};
        break;

      case FZL_OPCODE_POW:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t pow2 = vm->dstack[--vm->dtop];
        fzl_value_t pow1 = vm->dstack[--vm->dtop];

        uint64_t pow_n     = pow1.as.u;
        uint64_t pow_res   = pow_n;
        uint64_t pow_times = pow2.as.u;
        for (size_t i = 1; i < pow_times; i++)
          pow_res *= pow_n;

        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = pow1.type, .as.u = pow_res};
        break;

      case FZL_OPCODE_MOD:
        if (vm->dtop < 2)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected 2 values in stack");
          return -1;
        }
        fzl_value_t mod1 = vm->dstack[--vm->dtop];
        fzl_value_t mod2 = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = mod1.type, .as.u = mod2.as.u % mod1.as.u};
        break;

      case FZL_OPCODE_NEG:
        if (vm->dtop < 1)
        {
          snprintf(last_err, sizeof(last_err),
                   "runtime error - expected value in stack");
          return -1;
        }
        fzl_value_t neg_n = vm->dstack[--vm->dtop];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = neg_n.type, .as.u = -neg_n.as.u};
        break;

      case FZL_OPCODE_PUSHNUM:
      {
        uint64_t push_n = vm->ops[++vm->ip];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_UINT, .as.u = push_n};
        break;
      }

      case FZL_OPCODE_PUSHBOOL:
      {
        uint64_t push_b = vm->ops[++vm->ip];
        vm->dstack[vm->dtop++] =
            (fzl_value_t){.type = FZL_TYPE_BOOL, .as.b = push_b};
        break;
      }

      default:
        snprintf(last_err, sizeof(last_err),
                 "runtime error - unknown opcode %u", opcode);
        return -1;
    }

    vm->ip++;
  }

  return 0;
}

char *fzl_vm_geterror() { return last_err; }

void fzl_freevm(fzl_vm_t *vm)
{
  if (!vm)
    return;

  for (size_t i = 0; i < vm->consts_count; i++)
  {
    if (vm->consts[i].type == FZL_CONSTTYPE_STRING)
      free((void *)vm->consts[i].as.string.chars);
  }

  for (size_t i = 0; i < vm->objects_count; i++)
  {
    free(vm->objects[i]->values);
    free(vm->objects[i]);
  }

  free(vm->consts);
  free(vm->ops);
  free(vm);
}