#include "vm.h"
#include "device.h"
#include "opcodes.h"
#include <string.h>

/* no inline fn in ANSI C :( */
#define COMPARE_AND_JUMP(type, accessor, op)                                   \
  do {                                                                         \
    type value = vm->frames[vm->fp].registers[src1].accessor;                  \
    type value2 = vm->frames[vm->fp].registers[src2].accessor;                 \
    vm->pc =                                                                   \
        (value op value2) ? vm->frames[vm->fp].registers[dest].u : vm->pc;     \
    return true;                                                               \
  } while (0)

#define MATH_OP(accessor, op)                                                  \
  do {                                                                         \
    vm->frames[vm->fp].registers[dest].accessor =                              \
        vm->frames[vm->fp]                                                     \
            .registers[src1]                                                   \
            .accessor op vm->frames[vm->fp]                                    \
            .registers[src2]                                                   \
            .accessor;                                                         \
    return true;                                                               \
  } while (0)

/**
 * Embeds a string into the VM
 */
uint32_t str_alloc(VM *vm, const char *str, uint32_t length) {
  uint32_t str_addr = vm->mp++;
  uint32_t temp = 0;

  while (temp < length) {
    uint32_t idx = temp % 4;
    vm->memory[vm->mp].c[idx] = str[temp];
    temp++;

    if (idx == 3)
      vm->mp++;
  }

  uint32_t i, rem = temp % 4;
  if (rem != 0) {
    for (i = rem; i < 4; i++) {
      vm->memory[vm->mp].c[i] = '\0';
    }
    vm->mp++;
  }

  vm->memory[str_addr].u = length;
  vm->frames[vm->fp].allocated.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 */
  uint32_t instruction = vm->code[vm->pc++].u;

  uint8_t opcode = (instruction >> 24) & 0xFF;
  uint8_t dest = (instruction >> 16) & 0xFF;
  uint8_t src1 = (instruction >> 8) & 0xFF;
  uint8_t src2 = instruction & 0xFF;

  switch (opcode) {
  case OP_HALT: {
    return false;
  }
  case OP_CALL: {
    uint32_t jmp = vm->code[vm->pc++].u;   /* location of function in code */
    vm->return_stack[vm->rp++].u = vm->pc; /* set return address */
    vm->fp++; /* increment to the next free frame */
    vm->frames[vm->fp].allocated.start =
        vm->mp; /* set start of new memory block */
    vm->pc = jmp;
    return true;
  }
  case OP_RETURN: {
    vm->frames[vm->fp].rp = 0;             /* reset register ptr */
    vm->pc = vm->return_stack[--vm->rp].u; /* set pc to return address */
    vm->mp =
        vm->frames[vm->fp--].allocated.start; /* reset memory pointer to start
                                                 of old slice, pop the frame */
    return true;
  }
  case OP_LOAD: {
    uint32_t ptr = vm->code[vm->pc++].u;
    Value v = vm->memory[ptr];
    vm->frames[vm->fp].registers[dest] = v;
    return true;
  }
  case OP_STORE: {
    uint32_t ptr = vm->code[vm->pc++].u;
    vm->memory[ptr] = vm->frames[vm->fp].registers[src1];
    return true;
  }
  case OP_PUT: {
    uint32_t addr = vm->frames[vm->fp].registers[dest].u;
    vm->memory[addr] = vm->frames[vm->fp].registers[src1];
    return true;
  }
  case OP_GET: {
    uint32_t addr = vm->frames[vm->fp].registers[src1].u;
    vm->frames[vm->fp].registers[dest] = vm->memory[addr];
    return true;
  }
  case OP_OFFSET: {
    uint32_t ptr = vm->frames[vm->fp].registers[src1].u;
    uint32_t offset = vm->frames[vm->fp].registers[src2].u;
    uint32_t result = ptr + offset;
    vm->frames[vm->fp].registers[dest].u = result;
    return true;
  }
  case OP_PUSH: {
    vm->stack[++vm->sp] = vm->frames[vm->fp].registers[dest];
    return true;
  }
  case OP_POP: {
    vm->frames[vm->fp].registers[dest] = vm->stack[vm->sp--];
    return true;
  }
  case OP_MEM_ALLOC: {
    uint32_t mem_dest = (uint32_t)vm->frames[vm->fp]
                            .allocated.end; /* get start of unallocated */
    vm->frames[vm->fp].registers[dest].u =
        mem_dest; /* store ptr of array to dest register */
    uint32_t size = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    vm->memory[mem_dest].u = size;
    vm->frames[vm->fp].allocated.end +=
        size; /* increment to end of allocated */
    return true;
  }
  case OP_MEM_MOV: {
    uint32_t dest_addr = (uint32_t)vm->frames[vm->fp].registers[dest].u;
    uint32_t src_addr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    vm->memory[dest_addr] = vm->memory[src_addr];
    return true;
  }
  case OP_MEM_SWAP: {
    uint32_t dest_addr = (uint32_t)vm->frames[vm->fp].registers[dest].u;
    uint32_t src_addr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    vm->memory[dest_addr].u ^= vm->memory[src_addr].u;
    vm->memory[src_addr].u ^= vm->memory[dest_addr].u;
    vm->memory[dest_addr].u ^= vm->memory[src_addr].u;
    return true;
  }
  case OP_REG_SWAP: {
    vm->frames[vm->fp].registers[dest].u ^=
        vm->frames[vm->fp].registers[src1].u;
    vm->frames[vm->fp].registers[src1].u ^=
        vm->frames[vm->fp].registers[dest].u;
    vm->frames[vm->fp].registers[dest].u ^=
        vm->frames[vm->fp].registers[src1].u;
    return true;
  }
  case OP_REG_MOV: {
    vm->frames[vm->fp].registers[dest].i = vm->frames[vm->fp].registers[src1].i;
    return true;
  }
  case OP_GET_ACC: {
    vm->frames[vm->fp].registers[dest].u = vm->acc;
    return true;
  }
  case OP_GET_PC: {
    vm->frames[vm->fp].registers[dest].u = vm->pc;
    return true;
  }
  case OP_JMP: {
    vm->pc = vm->frames[vm->fp].registers[dest].u; /* Jump to address */
    return true;
  }  
  case OP_SYSCALL: {
    uint32_t syscall_id = dest;
    uint32_t arg_count = src2;
    uint32_t first_reg = src1;

    switch (syscall_id) {
    case SYSCALL_DEVICE_OPEN: {
      if (arg_count >= 2) {
        uint32_t path_ptr =
            vm->frames[vm->fp].registers[first_reg].u; /* R0: path pointer */
        uint32_t mode =
            vm->frames[vm->fp].registers[first_reg + 1].u; /* R1: mode */
        /* uint32_t length = vm->memory[path_ptr].u; */
        uint32_t str_src = path_ptr + 1;
        const char *path = (const char *)&vm->memory[str_src];

        Device *dev = find_device_by_path(vm, path);

        if (dev) {
          if (dev->ops->open) {
            int result = dev->ops->open(dev->data, mode);
            vm->acc = result;
          } else {
            vm->acc = 1; /* success, no open needed */
          }
        } else {
          vm->acc = 0; /* error */
        }
      } else {
        vm->acc = 0; /* error: not enough arguments */
      }
      return true;
    }

    case SYSCALL_DEVICE_READ: {
      if (arg_count >= 3) {
        uint32_t path_ptr =
            vm->frames[vm->fp].registers[first_reg].u; /* R0: path pointer */
        uint32_t buffer_ptr = vm->frames[vm->fp]
                                  .registers[first_reg + 1]
                                  .u; /* R1: buffer pointer */
        uint32_t size =
            vm->frames[vm->fp].registers[first_reg + 2].u; /* R2: size */
        /* uint32_t length = vm->memory[path_ptr].u; */
        uint32_t str_src = path_ptr + 1;

        const char *path = (const char *)&vm->memory[str_src];
        Device *dev = find_device_by_path(vm, path);

        if (dev && dev->ops->read) {
          int result = dev->ops->read(dev->data,
                                      (uint8_t *)&vm->memory[buffer_ptr], size);
          vm->acc = result;
        } else {
          vm->acc = 0;
        }
      } else {
        vm->acc = 0; /* error: not enough arguments */
      }
      return true;
    }

    case SYSCALL_DEVICE_WRITE: {
      if (arg_count >= 3) {
        uint32_t path_ptr =
            vm->frames[vm->fp].registers[first_reg].u; /* R0: path pointer */
        uint32_t buffer_ptr = vm->frames[vm->fp]
                                  .registers[first_reg + 1]
                                  .u; /* R1: buffer pointer */
        uint32_t size =
            vm->frames[vm->fp].registers[first_reg + 2].u; /* R2: size */
        /* uint32_t length = vm->memory[path_ptr].u; */
        uint32_t str_src = path_ptr + 1;

        const char *path = (const char *)&vm->memory[str_src];
        Device *dev = find_device_by_path(vm, path);

        if (dev && dev->ops->write) {
          int result = dev->ops->write(
              dev->data, (const uint8_t *)&vm->memory[buffer_ptr], size);
          vm->acc = result;
        } else {
          vm->acc = 0;
        }
      } else {
        vm->acc = 0; /* error: not enough arguments */
      }
      return true;
    }

    case SYSCALL_DEVICE_CLOSE: {
      if (arg_count >= 1) {
        uint32_t path_ptr =
            vm->frames[vm->fp].registers[first_reg].u; /* R0: path pointer */
        /* uint32_t length = vm->memory[path_ptr].u; */
        uint32_t str_src = path_ptr + 1;

        const char *path = (const char *)&vm->memory[str_src];
        Device *dev = find_device_by_path(vm, path);

        if (dev && dev->ops->close) {
          int result = dev->ops->close(dev->data);
          vm->acc = result;
        } else {
          vm->acc = 0;
        }
      } else {
        vm->acc = 0; /* error: not enough arguments */
      }
      return true;
    }

    case SYSCALL_DEVICE_IOCTL: {
      if (arg_count >= 3) {
        uint32_t path_ptr =
            vm->frames[vm->fp].registers[first_reg].u; /* R0: device path */
        uint32_t cmd = vm->frames[vm->fp]
                           .registers[first_reg + 1]
                           .u; /* R1: ioctl command */
        uint32_t args_ptr = vm->frames[vm->fp]
                                .registers[first_reg + 2]
                                .u; /* R2: args pointer */
        /* uint32_t length = vm->memory[path_ptr].u; */
        uint32_t str_src = path_ptr + 1;

        const char *path = (const char *)&vm->memory[str_src];
        Device *dev = find_device_by_path(vm, path);

        if (dev && dev->ops && dev->ops->ioctl) {
          int result = dev->ops->ioctl(dev->data, cmd, &vm->memory[args_ptr]);
          vm->acc = result;
        } else {
          vm->acc = 0; /* error or no ioctl support */
        }
      } else {
        vm->acc = 0; /* error: not enough arguments */
      }
      return true;
    }

    case SYSCALL_EXIT: {
      return false;
    }

    default: {
      vm->acc = 0; /* unknown syscall */
      return true;
    }
    }
    return true;
  }
  case OP_ADD_INT:
    MATH_OP(i, +);
  case OP_SUB_INT:
    MATH_OP(i, -);
  case OP_MUL_INT:
    MATH_OP(i, *);
  case OP_DIV_INT:
    MATH_OP(i, /);
  case OP_ADD_UINT:
    MATH_OP(u, +);
  case OP_SUB_UINT:
    MATH_OP(u, -);
  case OP_MUL_UINT:
    MATH_OP(u, *);
  case OP_DIV_UINT:
    MATH_OP(u, /);
  case OP_ADD_REAL:
    MATH_OP(f, +);
  case OP_SUB_REAL:
    MATH_OP(f, -);
  case OP_MUL_REAL:
    MATH_OP(f, *);
  case OP_DIV_REAL:
    MATH_OP(f, /);
  case OP_REAL_TO_INT: {
    vm->frames[vm->fp].registers[dest].i =
        (int32_t)(vm->frames[vm->fp].registers[src1].f);
    return true;
  }
  case OP_INT_TO_REAL: {
    vm->frames[vm->fp].registers[dest].f =
        (float)(vm->frames[vm->fp].registers[src1].i);
    return true;
  }
  case OP_REAL_TO_UINT: {
    vm->frames[vm->fp].registers[dest].u =
        (uint32_t)(vm->frames[vm->fp].registers[src1].f);
    return true;
  }
  case OP_UINT_TO_REAL: {
    vm->frames[vm->fp].registers[dest].f =
        (float)(vm->frames[vm->fp].registers[src1].u);
    return true;
  }
  case OP_JEQ_UINT: {
    COMPARE_AND_JUMP(uint32_t, u, ==);
  }
  case OP_JGT_UINT: {
    COMPARE_AND_JUMP(uint32_t, u, >);
  }
  case OP_JLT_UINT: {
    COMPARE_AND_JUMP(uint32_t, u, <);
  }
  case OP_JLE_UINT: {
    COMPARE_AND_JUMP(uint32_t, u, <=);
  }
  case OP_JGE_UINT: {
    COMPARE_AND_JUMP(uint32_t, u, >=);
  }
  case OP_JEQ_INT: {
    COMPARE_AND_JUMP(int32_t, i, ==);
  }
  case OP_JGT_INT: {
    COMPARE_AND_JUMP(int32_t, i, >);
  }
  case OP_JLT_INT: {
    COMPARE_AND_JUMP(int32_t, i, <);
  }
  case OP_JLE_INT: {
    COMPARE_AND_JUMP(int32_t, i, <=);
  }
  case OP_JGE_INT: {
    COMPARE_AND_JUMP(int32_t, i, >=);
  }
  case OP_JEQ_REAL: {
    COMPARE_AND_JUMP(int32_t, i, ==);
  }
  case OP_JGT_REAL: {
    COMPARE_AND_JUMP(float, u, >);
  }
  case OP_JLT_REAL: {
    COMPARE_AND_JUMP(float, u, <);
  }
  case OP_JGE_REAL: {
    COMPARE_AND_JUMP(float, u, >=);
  }
  case OP_JLE_REAL: {
    COMPARE_AND_JUMP(float, u, <=);
  }
  case OP_INT_TO_STRING: {
    int32_t a = (int32_t)vm->frames[vm->fp].registers[src1].i; /* get value */
    char buffer[32];
    int len = sprintf(buffer, "%d", a);
    uint32_t ptr = str_alloc(vm, buffer, len); /* copy buffer to dest */
    vm->frames[vm->fp].registers[dest].u = ptr;
    return true;
  }
  case OP_UINT_TO_STRING: {
    uint32_t a = (uint32_t)vm->frames[vm->fp].registers[src1].u; /* get value */
    char buffer[32];
    int len = sprintf(buffer, "%d", a);
    uint32_t ptr = str_alloc(vm, buffer, len); /* copy buffer to dest */
    vm->frames[vm->fp].registers[dest].u = ptr;
    return true;
  }
  case OP_REAL_TO_STRING: {
    float a = (float)vm->frames[vm->fp].registers[src1].f; /* get value */
    char buffer[32];
    int len = sprintf(buffer, "%f", a);
    uint32_t ptr = str_alloc(vm, buffer, len); /* copy buffer to dest */
    vm->frames[vm->fp].registers[dest].u = ptr;
    return true;
  }
  case OP_STRING_TO_INT: {
    uint32_t src_addr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    uint32_t dest_addr = (uint32_t)vm->frames[vm->fp].registers[dest].u;
    char *endptr;
    int32_t value = (int32_t)strtol((char*)&vm->memory[src_addr + 1], &endptr, 10);
    vm->memory[dest_addr].i = value;
    return true;
  }
  case OP_STRING_TO_UINT: {
    uint32_t src_addr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    uint32_t dest_addr = (uint32_t)vm->frames[vm->fp].registers[dest].u;
    long value = atol((char*)&vm->memory[src_addr + 1]);
    vm->memory[dest_addr].u = value;
    return true;
  }
  case OP_STRING_TO_REAL: {
    uint32_t src_addr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    uint32_t dest_addr = (uint32_t)vm->frames[vm->fp].registers[dest].u;
    float value = atof((char*)&vm->memory[src_addr + 1]);
    vm->memory[dest_addr].u = value;
    return true;
  }
  case OP_DBG_READ_STRING: {
    uint32_t str_addr = vm->mp++;
    uint32_t length = 0;
    int ch;

    while ((ch = getchar()) != '\n' && ch != EOF) {
      uint32_t idx = length % 4;
      vm->memory[vm->mp].c[idx] = (char)ch;
      length++;

      if (idx == 3)
        vm->mp++;
    }

    uint32_t i, rem = length % 4;
    if (rem != 0) {
      for (i = rem; i < 4; i++) {
        vm->memory[vm->mp].c[i] = '\0';
      }
      vm->mp++;
    }

    vm->memory[str_addr].u = length;
    vm->frames[vm->fp].allocated.end = vm->mp;
    vm->frames[vm->fp].registers[dest].u = str_addr;

    return true;
  }
  case OP_DBG_PRINT_STRING: {
    uint32_t ptr = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    uint32_t length = vm->memory[ptr].u;
    uint32_t str_src = ptr + 1;
    uint32_t i;
    for (i = 0; i < length; i++) {
      uint8_t ch = vm->memory[str_src + (i / 4)].c[i % 4];
      if (ch == '\0')
        break;
      putchar(ch);
    }
    putchar('\n');
    return true;
  }
  case OP_CMP_STRING: {
    uint32_t addr1 = (uint32_t)vm->frames[vm->fp].registers[src1].u;
    uint32_t addr2 = (uint32_t)vm->frames[vm->fp].registers[src2].u;
    uint32_t length1 = vm->memory[addr1 - 1].u;
    uint32_t length2 = vm->memory[addr2 - 1].u;
    uint32_t equal =
        1; /* we dont have a native boolean type so we use uint32_t */

    if (length1 != length2) {
      equal = 0;
    } else {
      uint32_t i = 0;
      while (i < length1) {
        uint32_t char1 = vm->memory[addr1 + i].u;
        uint32_t char2 = vm->memory[addr2 + i].u;
        if (char1 != char2) {
          equal = 0;
          break;
        }
        if ((char1 & 0xFF) == '\0' && (char2 & 0xFF) == '\0') {
          equal = 1;
          break;
        }
      }
    }
    vm->memory[dest].u = equal;
    return true;
  }
  }
  return false; /* something bad happened */
}