#include "vm.h"
#include "device.h"
#include "opcodes.h"
#include "str.h"

#define COMPARE_AND_JUMP(type, op)                                             \
  do {                                                                         \
    i32 cond;                                                                  \
    u32 mask, target;                                                          \
    type value;                                                                \
    type value2;                                                               \
    target = read_u32(vm, code, vm->pc);                                       \
    vm->pc += 4;                                                               \
    value2 = (type)vm->stack[--vm->sp];                                         \
    value = (type)vm->stack[--vm->sp];                                        \
    cond = !!(value op value2);                                                \
    mask = -(u32)cond;                                                         \
    vm->pc = (target & mask) | (vm->pc & ~mask);                               \
    return true;                                                               \
  } while (0)

#define MATH_OP(type, op)                                                      \
  do {                                                                         \
    type b = (type)vm->stack[--vm->sp];                                        \
    type a = (type)vm->stack[--vm->sp];                                        \
    vm->stack[vm->sp++] = (type)(a op b);                                      \
    return true;                                                               \
  } while (0)

#define BIT_OP(op)                                                             \
  do {                                                                         \
    u32 a = vm->stack[--vm->sp];                                               \
    u32 b = vm->stack[--vm->sp];                                               \
    vm->stack[vm->sp++] = a op b;                                              \
    return true;                                                               \
  } while (0)

u32 str_alloc(VM *vm, Frame *frame, const char *str, u32 length) {
  u32 str_addr = vm->mp;
  u32 i = 0;
  vm->mp += 4;
  while (i < length) {
    vm->memory[vm->mp++] = str[i++];
  }
  vm->memory[vm->mp++] = '\0';

  write_u32(vm, memory, str_addr, length);
  frame->end = vm->mp;
  return str_addr;
}

/**
 * Step to the next opcode in the vm.
 */
bool step_vm(VM *vm) {
  /* Get current instruction & Advance to next instruction */
  u8 opcode = vm->code[vm->pc++];
  Frame *frame = &vm->frames[vm->fp];

  switch (opcode) {
  case OP_HALT: {
    return false;
  }
  case OP_CALL: {
    u32 jmp = read_u32(vm, code, vm->pc); /* location of function in code */
    vm->pc += 4;
    vm->return_stack[vm->rp++] = vm->pc; /* set return address */
    vm->fp++;                            /* increment to the next free frame */
    vm->frames[vm->fp].start = vm->mp;   /* set start of new memory block */
    vm->pc = jmp;
    return true;
  }
  case OP_RETURN: {
    vm->pc = vm->return_stack[--vm->rp]; /* set pc to return address */
    vm->mp =
        vm->frames[vm->fp--].start; /* reset memory pointer to start
                                                 of old slice, pop the frame */
    return true;
  }
  case OP_MALLOC: {
    u32 size = read_u32(vm, code, vm->pc);
    vm->pc++;
    vm->stack[vm->sp++] = vm->mp;
    write_u32(vm, memory, vm->mp, size);
    vm->mp += (size + 4);
    return true;
  }
  case OP_MEMSET_32: {
    u32 i, start, end;
    u32 dest = vm->stack[--vm->sp];
    u32 value = vm->stack[--vm->sp];
    u32 count = vm->stack[--vm->sp];

    if (count == 0) {
      vm->flag = 1;
      return true;
    }

    start = dest;
    end = dest + count;

    if (start >= vm->mp || count > vm->mp || end > vm->mp) {
      vm->flag = 0;
      return true;
    }

    for (i = start; i < end; i += 4) {
      write_u32(vm, memory, i, value);
    }

    vm->stack[vm->sp++] = dest;
    vm->flag = 1;
    return true;
  }
  case OP_MEMSET_16: {
    u32 i, start, end;
    u32 dest = vm->stack[--vm->sp];
    u16 value = (u16)(vm->stack[--vm->sp]);
    u32 count = vm->stack[--vm->sp];

    if (count == 0) {
      vm->flag = 1;
      return true;
    }

    start = dest;
    end = dest + count;

    if (start >= vm->mp || count > vm->mp || end > vm->mp) {
      vm->flag = 0;
      return true;
    }

    for (i = start; i < end; i += 2) {
      write_u16(vm, memory, i, value);
    }

    vm->stack[vm->sp++] = dest;
    vm->flag = 1;
    return true;
  }
  case OP_MEMSET_8: {
    u32 i, start, end;
    u32 dest = vm->stack[--vm->sp];
    u8 value = (u8)(vm->stack[--vm->sp]);
    u32 count = vm->stack[--vm->sp];

    if (count == 0) {
      vm->flag = 1;
      return true;
    }

    start = dest;
    end = dest + count;

    if (start >= vm->mp || count > vm->mp || end > vm->mp) {
      vm->flag = 0;
      return true;
    }

    for (i = start; i < end; i++) {
      write_u8(vm, memory, i, value);
    }

    vm->stack[vm->sp++] = dest;
    vm->flag = 1;
    return true;
  }
  case OP_LOAD_8: {
    u32 addr = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = read_u8(vm, memory, addr);
    return true;
  }
  case OP_LOAD_16: {
    u32 addr = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = read_u16(vm, memory, addr);
    return true;
  }
  case OP_LOAD_32: {
    u32 addr = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = read_u32(vm, memory, addr);
    return true;
  }
  case OP_STORE_8: {
    u8 value = vm->stack[--vm->sp];
    u32 addr = vm->stack[--vm->sp];
    write_u8(vm, memory, addr, value);
    return true;
  }
  case OP_STORE_16: {
    u16 value = vm->stack[--vm->sp];
    u32 addr = vm->stack[--vm->sp];
    write_u16(vm, memory, addr, value);
    return true;
  }
  case OP_STORE_32: {
    u32 value = vm->stack[--vm->sp];
    u32 addr = vm->stack[--vm->sp];
    write_u32(vm, memory, addr, value);
    return true;
  }
  case OP_PUSH: {
    u32 val = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    vm->stack[vm->sp++] = val;
    return true;
  }
  case OP_POP: {
    vm->sp--;
    return true;
  }
  case OP_DUP: {
    u32 a = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = a;
    vm->stack[vm->sp++] = a;
    return true;
  }
  case OP_EXCH: {
    u32 a = vm->stack[--vm->sp];
    u32 b = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = b;
    vm->stack[vm->sp++] = a;
    return true;
  }
  case OP_OVER: {
    u32 a = vm->stack[vm->sp - 1];
    vm->stack[vm->sp++] = a;
    return true;
  }
  case OP_PICK: {
    u32 n = vm->stack[--vm->sp];
    u32 b = vm->stack[vm->sp - n];
    vm->stack[vm->sp++] = b;
    return true;
  }
  case OP_ROT: {
    return false;
  }
  case OP_DEPTH: {
    u32 a = vm->sp;
    vm->stack[vm->sp++] = a; 
    return true;
  }
  case OP_JMP: {
    u32 dest = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    vm->pc = dest; /* Jump to address */
    return true;
  }
  case OP_JMPF: { /* error handling for syscall, jump if flag == 0 */
    u32 mask, dest = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    mask = -(u32)(vm->flag == 0);
    vm->pc = (dest & mask) | (vm->pc & ~mask);
    return true;
  }
  case OP_SYSCALL: {
    u32 syscall_id;

    syscall_id = read_u32(vm, code, vm->pc);
    vm->pc += 4;

    switch (syscall_id) {
    case SYSCALL_DEVICE_OPEN: {
      Device *dev;
      u32 path_ptr = vm->stack[--vm->sp];
      u32 mode = vm->stack[--vm->sp];
      dev = find_device_by_path(vm, (const char *)&vm->memory[path_ptr + 4]);
      if (dev) {
        if (dev->ops->open) {
          vm->flag = dev->ops->open(dev->data, mode);
          vm->stack[vm->sp++] = dev->handle;
        } else {
          vm->flag = 1; /* success, no open needed */
        }
      } else {
        vm->flag = 0; /* error */
      }

      return true;
    }

    case SYSCALL_DEVICE_READ: {
      Device *dev;

      u32 handle = vm->stack[--vm->sp]; /* path pointer */
      u32 size = vm->stack[--vm->sp];     /* size */
      u32 buffer_ptr = vm->stack[--vm->sp];

      dev = &vm->devices[handle];
      if (dev && dev->ops->read) {
        vm->flag = dev->ops->read(dev->data, &vm->memory[buffer_ptr + 4], size);
        vm->stack[vm->sp++] = buffer_ptr;
      } else {
        vm->flag = 0;
      }

      return true;
    }

    case SYSCALL_DEVICE_WRITE: {
      Device *dev;
      u32 handle = vm->stack[--vm->sp];   /* path pointer */
      u32 size = vm->stack[--vm->sp];       /* size */
      u32 buffer_ptr = vm->stack[--vm->sp]; /* buffer pointer */

      dev = &vm->devices[handle];
      if (dev && dev->ops->write) {
        vm->flag = dev->ops->write(
            dev->data, (const u8 *)&vm->memory[buffer_ptr + 4], size);
      } else {
        vm->flag = 0;
      }

      return true;
    }

    case SYSCALL_DEVICE_CLOSE: {
      Device *dev;
      u32 handle = vm->stack[--vm->sp]; /* path pointer */

      dev = &vm->devices[handle];
      if (dev && dev->ops->close) {
        i32 result = dev->ops->close(dev->data);
        vm->flag = result;
      } else {
        vm->flag = 0;
      }

      return true;
    }

    case SYSCALL_DEVICE_IOCTL: {
      Device *dev;
      u32 handle = vm->stack[--vm->sp]; /* device path */
      u32 cmd = vm->stack[--vm->sp];      /* ioctl command */
      u32 args_ptr = vm->stack[--vm->sp]; /* args pointer */

      dev = &vm->devices[handle];
      if (dev && dev->ops && dev->ops->ioctl) {
        i32 result = dev->ops->ioctl(dev->data, cmd, &vm->memory[args_ptr]);
        vm->flag = result;
      } else {
        vm->flag = 0; /* error or no ioctl support */
      }

      return true;
    }

    case SYSCALL_EXIT: {
      return false;
    }

    default: {
      vm->flag = 0; /* unknown syscall */
      return true;
    }
    }
    return true;
  }
  case OP_SLL:
    BIT_OP(<<);
  case OP_SRL:
    BIT_OP(>>);
  case OP_SRE:
    MATH_OP(i32, >>);
  case OP_BAND:
    BIT_OP(&);
  case OP_BOR:
    BIT_OP(|);
  case OP_BXOR:
    BIT_OP(^);
  case OP_ADD_INT:
    MATH_OP(i32, +);
  case OP_SUB_INT:
    MATH_OP(i32, -);
  case OP_MUL_INT:
    MATH_OP(i32, *);
  case OP_DIV_INT:
    MATH_OP(i32, /);
  case OP_ADD_UINT:
    MATH_OP(u32, +);
  case OP_SUB_UINT:
    MATH_OP(u32, -);
  case OP_MUL_UINT:
    MATH_OP(u32, *);
  case OP_DIV_UINT:
    MATH_OP(u32, /);
  case OP_MUL_REAL: {
    u32 a = vm->stack[--vm->sp];
    u32 b = (vm->stack[--vm->sp] >> 16);
    vm->stack[vm->sp++] = a * b;
    return true;
  }
  case OP_DIV_REAL: {
    u32 a = vm->stack[--vm->sp] << 16;
    u32 b = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = a / b;
    return true;
  }
  case OP_ADD_REAL: {
    u32 a = vm->stack[--vm->sp];
    u32 b = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = a + b;
    return true;
  }
  case OP_SUB_REAL: {
    u32 a = vm->stack[--vm->sp];
    u32 b = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = a - b;
    return true;
  }
  case OP_REAL_TO_INT: {
    i32 value = vm->stack[--vm->sp];
    if (value >= 0) {
      vm->stack[vm->sp++] = value >> 16;
    } else {
      vm->stack[vm->sp++] = -((-value) >> 16);
    }
    return true;
  }
  case OP_INT_TO_REAL: {
    i32 value = (vm->stack[--vm->sp] << 16);
    vm->stack[vm->sp++] = value;
    return true;
  }
  case OP_REAL_TO_UINT: {
    i32 value = vm->stack[--vm->sp];
    if (value < 0) {
      vm->stack[vm->sp++] = 0;
    } else {
      vm->stack[vm->sp++] = AS_UINT(value >> 16);
    }
    return true;
  }
  case OP_UINT_TO_REAL: {
    i32 value = AS_INT(vm->stack[--vm->sp] << 16);
    vm->stack[vm->sp++] = value;
    return true;
  }
  case OP_JEQ_UINT: {
    COMPARE_AND_JUMP(u32, ==);
  }
  case OP_JNEQ_UINT: {
    COMPARE_AND_JUMP(u32, !=);
  }
  case OP_JGT_UINT: {
    COMPARE_AND_JUMP(u32, >);
  }
  case OP_JLT_UINT: {
    COMPARE_AND_JUMP(u32, <);
  }
  case OP_JLE_UINT: {
    COMPARE_AND_JUMP(u32, <=);
  }
  case OP_JGE_UINT: {
    COMPARE_AND_JUMP(u32, >=);
  }
  case OP_JEQ_INT: {
    COMPARE_AND_JUMP(i32, ==);
  }
  case OP_JNEQ_INT: {
    COMPARE_AND_JUMP(i32, !=);
  }
  case OP_JGT_INT: {
    COMPARE_AND_JUMP(i32, >);
  }
  case OP_JLT_INT: {
    COMPARE_AND_JUMP(i32, <);
  }
  case OP_JLE_INT: {
    COMPARE_AND_JUMP(i32, <=);
  }
  case OP_JGE_INT: {
    COMPARE_AND_JUMP(i32, >=);
  }
  case OP_JEQ_REAL: {
    COMPARE_AND_JUMP(i32, ==);
  }
  case OP_JNEQ_REAL: {
    COMPARE_AND_JUMP(i32, !=);
  }
  case OP_JGT_REAL: {
    COMPARE_AND_JUMP(i32, >);
  }
  case OP_JLT_REAL: {
    COMPARE_AND_JUMP(i32, <);
  }
  case OP_JGE_REAL: {
    COMPARE_AND_JUMP(i32, >=);
  }
  case OP_JLE_REAL: {
    COMPARE_AND_JUMP(i32, <=);
  }
  case OP_INT_TO_STRING: {
    char buffer[32];
    int_to_string(AS_INT(vm->stack[--vm->sp]), buffer);
    vm->stack[vm->sp++] = str_alloc(vm, frame, buffer, strlength(buffer));
    return true;
  }
  case OP_UINT_TO_STRING: {
    char buffer[32];
    uint_to_string(vm->stack[--vm->sp], buffer);
    vm->stack[vm->sp++] = str_alloc(vm, frame, buffer, strlength(buffer));
    return true;
  }
  case OP_REAL_TO_STRING: {
    char buffer[32];
    fixed_to_string(AS_INT(vm->stack[--vm->sp]), buffer);
    vm->stack[vm->sp++] = str_alloc(
        vm, frame, buffer, strlength(buffer)); /* copy buffer to dest */
    return true;
  }
  case OP_STRLEN: {
    u32 ptr;
    ptr = vm->stack[--vm->sp];
    vm->stack[vm->sp++] = read_u32(vm, memory, ptr);
    return true;
  }
  case OP_STRCAT: {
    /* not implemented yet */
    return false;
  }
  case OP_STR_GET_CHAR: {
    /* not implemented yet */
    return false;
  }
  case OP_STR_FIND_CHAR: {
    /* not implemented yet */
    return false;
  }
  case OP_STR_SLICE: {
    /* not implemented yet */
    return false;
  }
  case OP_STRING_TO_INT: {
    /* not implemented yet */
    return false;
  }
  case OP_STRING_TO_UINT: {
    /* not implemented yet */
    return false;
  }
  case OP_STRING_TO_REAL: {
    /* not implemented yet */
    return false;
  }
  }
  return false; /* something bad happened */
}