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

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

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

u32 str_alloc(VM *vm, 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);
  vm->frames[vm->fp].end = vm->mp;
  return str_addr;
}

/**
 * Step to the next opcode in the vm.
 */
bool step_vm(VM *vm) {
  u8 opcode, dest, src1, src2;
  u32 v, ptr;
  i32 value;

  /* Get current instruction & Advance to next instruction */
  opcode = vm->code[vm->pc++];

  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->frames[vm->fp].rp = 0;           /* reset register ptr */
    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_LOAD_IMM: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    v = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    vm->frames[vm->fp].registers[dest] = v;
    return true;
  }
  case OP_LOAD: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    ptr = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    v = read_u32(vm, memory, ptr);
    vm->frames[vm->fp].registers[dest] = v;
    return true;
  }
  case OP_STORE: {
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    ptr = read_u32(vm, code, vm->pc);
    vm->pc += 4;
    v = vm->frames[vm->fp].registers[src1];
    write_u32(vm, memory, ptr, v);
    return true;
  }
  case OP_PUSH: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->stack[++vm->sp] = vm->frames[vm->fp].registers[dest];
    return true;
  }
  case OP_POP: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] = vm->stack[vm->sp--];
    return true;
  }
  case OP_REG_MOV: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] = vm->frames[vm->fp].registers[src1];
    return true;
  }
  case OP_JMP: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->pc = vm->frames[vm->fp].registers[dest]; /* Jump to address */
    return true;
  }
  case OP_SYSCALL: {
    u32 syscall_id, arg_count;
    u8 first_reg;

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

    switch (syscall_id) {
    case SYSCALL_DEVICE_OPEN: {
      if (arg_count >= 2) {
        Device *dev;
        u32 path_ptr, mode, str_src;
        path_ptr =
            vm->frames[vm->fp].registers[first_reg]; /* R0: path pointer */
        mode = vm->frames[vm->fp].registers[first_reg + 1]; /* R1: mode */
        str_src = path_ptr + 1;

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

        if (dev) {
          if (dev->ops->open) {
            vm->acc = dev->ops->open(dev->data, mode);
          } 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) {
        Device *dev;
        u32 path_ptr, buffer_ptr, size, str_src;

        path_ptr =
            vm->frames[vm->fp].registers[first_reg]; /* R0: path pointer */
        buffer_ptr = vm->frames[vm->fp]
                         .registers[first_reg + 1]; /* R1: buffer pointer */
        size = vm->frames[vm->fp].registers[first_reg + 2]; /* R2: size */
        str_src = path_ptr + 1;

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

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

    case SYSCALL_DEVICE_WRITE: {
      if (arg_count >= 2) {
        Device *dev;
        u32 path_ptr, buffer_ptr, size, str_src;
        path_ptr =
            vm->frames[vm->fp].registers[first_reg]; /* R0: path pointer */
        buffer_ptr = vm->frames[vm->fp]
                         .registers[first_reg + 1]; /* R1: buffer pointer */
        size = read_u32(vm, memory, buffer_ptr); /* R2: size */
        str_src = path_ptr += 4;

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

    case SYSCALL_DEVICE_CLOSE: {
      if (arg_count >= 1) {
        Device *dev;
        u32 path_ptr, str_src;
        path_ptr =
            vm->frames[vm->fp].registers[first_reg]; /* R0: path pointer */
        str_src = path_ptr + 1;

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

        if (dev && dev->ops->close) {
          i32 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) {
        Device *dev;
        u32 path_ptr, args_ptr, cmd, str_src;
        path_ptr =
            vm->frames[vm->fp].registers[first_reg]; /* R0: device path */
        cmd =
            vm->frames[vm->fp].registers[first_reg + 1]; /* R1: ioctl command */
        args_ptr =
            vm->frames[vm->fp].registers[first_reg + 2]; /* R2: args pointer */
        str_src = path_ptr + 1;

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

        if (dev && dev->ops && dev->ops->ioctl) {
          i32 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(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: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    src2 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] = (vm->frames[vm->fp].registers[src1] *
                                          vm->frames[vm->fp].registers[src2]) >>
                                         16;
    return true;
  }

  case OP_DIV_REAL: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    src2 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] =
        (vm->frames[vm->fp].registers[src1] << 16) /
        vm->frames[vm->fp].registers[src2];
    return true;
  }

  case OP_ADD_REAL: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    src2 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] =
        vm->frames[vm->fp].registers[src1] + vm->frames[vm->fp].registers[src2];
    return true;
  }

  case OP_SUB_REAL: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    src2 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] =
        vm->frames[vm->fp].registers[src1] - vm->frames[vm->fp].registers[src2];
    return true;
  }
  case OP_REAL_TO_INT: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    value = vm->frames[vm->fp].registers[src1];

    if (value >= 0) {
      vm->frames[vm->fp].registers[dest] = value >> 16;
    } else {
      vm->frames[vm->fp].registers[dest] = -((-value) >> 16);
    }

    return true;
  }
  case OP_INT_TO_REAL: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] =
        (vm->frames[vm->fp].registers[src1] << 16);
    return true;
  }
  case OP_REAL_TO_UINT: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    value = vm->frames[vm->fp].registers[src1];
    if (value < 0) {
      vm->frames[vm->fp].registers[dest] = 0;
    } else {
      vm->frames[vm->fp].registers[dest] = AS_UINT(value >> 16);
    }
    return true;
  }
  case OP_UINT_TO_REAL: {
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    vm->frames[vm->fp].registers[dest] =
        AS_INT(vm->frames[vm->fp].registers[src1] << 16);
    return true;
  }
  case OP_JEQ_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_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_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];
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    int_to_string(AS_INT(vm->frames[vm->fp].registers[src1]), buffer);
    ptr = str_alloc(vm, buffer, strlen(buffer));
    vm->frames[vm->fp].registers[dest] = ptr;
    return true;
  }
  case OP_UINT_TO_STRING: {
    char buffer[32];
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    uint_to_string(vm->frames[vm->fp].registers[src1], buffer);
    ptr = str_alloc(vm, buffer, strlen(buffer));
    vm->frames[vm->fp].registers[dest] = ptr;
    return true;
  }
  case OP_REAL_TO_STRING: {
    char buffer[32];
    dest = read_u8(vm, code, vm->pc);
    vm->pc++;
    src1 = read_u8(vm, code, vm->pc);
    vm->pc++;
    fixed_to_string(AS_INT(vm->frames[vm->fp].registers[src1]), buffer);
    ptr = str_alloc(vm, buffer, strlen(buffer)); /* copy buffer to dest */
    vm->frames[vm->fp].registers[dest] = ptr;
    return true;
  }
  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 */
}

/* Static digit lookup table (0-9) */
const char digits[10] = {'0', '1', '2', '3', '4', '5', '6', '7', '8', '9'};

/* Writes decimal digits of 'value' backwards into 'buf_end' (inclusive),
   stopping at 'buf_start'. Returns pointer to first written digit. */
char *write_digits_backwards(u32 value, char *buf_end, char *buf_start) {
  char *p = buf_end;
  if (value == 0) {
    *--p = '0';
  } else {
    u32 num = value;
    while (num && p > buf_start) {
      *--p = digits[num % 10];
      num /= 10;
    }
  }
  return p;
}

void uint_to_string(u32 value, char *buffer) {
  char temp[16];
  char *start;
  char *end = temp + sizeof(temp) - 1;
  *end = '\0';
  start = write_digits_backwards(value, end, temp);
  strcopy(buffer, start, end - start + 1); /* +1 for null terminator */
}

void int_to_string(i32 value, char *buffer) {
  char temp[17]; /* Extra space for '-' */
  i32 negative = 0;
  u32 abs_value;
  char *end = temp + sizeof(temp) - 1;
  *end = '\0';

  if (value == (-2147483647 - 1)) { /* INT32_MIN */
    strcopy(buffer, "-2147483648", 12);
    return;
  }

  if (value == 0) {
    *--end = '0';
  } else {
    if (value < 0) {
      negative = 1;
      abs_value = (u32)(-value);
    } else {
      abs_value = (u32)value;
    }

    end = write_digits_backwards(abs_value, end, temp);

    if (negative) {
      *--end = '-';
    }
  }

  strcopy(buffer, end, temp + sizeof(temp) - end);
}

void fixed_to_string(i32 value, char *buffer) {
  char temp[32];
  i32 negative;
  u32 int_part;
  u32 frac_part, frac_digits;
  char *end = temp + sizeof(temp) - 1;
  *end = '\0';

  negative = 0;
  if (value < 0) {
    negative = 1;
    value = -value;
  }

  int_part = AS_UINT(value >> 16);
  frac_part = AS_UINT(value & 0xFFFF);

  /* Convert fractional part to 5 decimal digits */
  frac_digits = (frac_part * 100000U) / 65536U;

  /* Trim trailing zeros */
  while (frac_digits > 0 && frac_digits % 10 == 0) {
    frac_digits /= 10;
  }

  if (frac_digits > 0) {
    end = write_digits_backwards(frac_digits, end, temp);
    *--end = '.';
  }

  if (int_part == 0 && frac_digits == 0) {
    *--end = '0';
  } else {
    end = write_digits_backwards(int_part, end, temp);
  }

  if (negative) {
    *--end = '-';
  }

  strcopy(buffer, end, temp + sizeof(temp) - end);
}