#include <errno.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>

#define MEMORY_SIZE 65536 /* 64KB memory (adjustable) */
typedef union {
  float f;
  uint32_t u;
} FloatUint32Union;
FloatUint32Union memory[MEMORY_SIZE]; /* Memory array */

typedef enum {
  OP_HALT,    /* terminate execution */
  OP_ADD,     /* dest = src1 + src2  */
  OP_SUB,     /* dest = src1 - src2  */
  OP_MUL,     /* dest = src1 * src2  */
  OP_DIV,     /* dest = src1 / src2  */
  OP_ADD_F32, /* dest = src1 + src2  */
  OP_SUB_F32, /* dest = src1 - src2  */
  OP_MUL_F32, /* dest = src1 * src2  */
  OP_DIV_F32, /* dest = src1 / src2  */
  OP_MOV,     /* dest = src1	     */
  OP_JMP,     /* jump to address src1 unconditionally */
  OP_JGZ,     /* jump to address dest if src1 > 0 */
  OP_READ_STRING,
  OP_PRINT_STRING,
} Opcode;

uint8_t get_char(uint32_t word, int index) {
  return (word >> (8 * index)) & 0xFF;
}

uint32_t set_char(uint32_t word, int index, uint8_t ch) {
  return (word & ~(0xFF << (8 * index))) | (ch << (8 * index));
}
void run_vm() {
  uint32_t pc = 0; /* Program counter */

  while (pc < MEMORY_SIZE - 4) {

    Opcode opcode = memory[pc].u;
    uint32_t src1_addr = memory[pc + 1].u;
    uint32_t src2_addr = memory[pc + 2].u;
    uint32_t dest_addr = memory[pc + 3].u;
    pc += 4; /* Advance to next instruction */

    if (src1_addr >= MEMORY_SIZE || src2_addr >= MEMORY_SIZE ||
        dest_addr >= MEMORY_SIZE) {
      printf("Invalid memory address!\n");
      exit(1);
    }

    switch (opcode) {
    case OP_ADD:
      memory[dest_addr].u = memory[src1_addr].u + memory[src2_addr].u;
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      break;

    case OP_SUB:
      memory[dest_addr].u = memory[src1_addr].u - memory[src2_addr].u;
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      break;

    case OP_MUL:
      memory[dest_addr].u = memory[src1_addr].u * memory[src2_addr].u;
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      break;

    case OP_DIV:
      memory[dest_addr].u = memory[src1_addr].u / memory[src2_addr].u;
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      break;

    case OP_ADD_F32:
      memory[dest_addr].f = memory[src1_addr].f + memory[src2_addr].f;
      printf("ADD;src1[%d]:%f;src2[%d]:%f;dest[%d]:%f\n", src1_addr,
             memory[src1_addr].f, src2_addr, memory[src2_addr].f, dest_addr,
             memory[dest_addr].f);
      break;

    case OP_SUB_F32:
      memory[dest_addr].f = memory[src1_addr].f - memory[src2_addr].f;
      printf("ADD;src1[%d]:%f;src2[%d]:%f;dest[%d]:%f\n", src1_addr,
             memory[src1_addr].f, src2_addr, memory[src2_addr].f, dest_addr,
             memory[dest_addr].f);
      break;

    case OP_MUL_F32:
      memory[dest_addr].f = memory[src1_addr].f * memory[src2_addr].f;
      printf("ADD;src1[%d]:%f;src2[%d]:%f;dest[%d]:%f\n", src1_addr,
             memory[src1_addr].f, src2_addr, memory[src2_addr].f, dest_addr,
             memory[dest_addr].f);
      break;

    case OP_DIV_F32:
      if (memory[src2_addr].f == 0.0f) {
        printf("Division by zero error at address %d\n", pc - 4);
        exit(1);
      }
      memory[dest_addr].f = memory[src1_addr].f / memory[src2_addr].f;
      printf("ADD;src1[%d]:%f;src2[%d]:%f;dest[%d]:%f\n", src1_addr,
             memory[src1_addr].f, src2_addr, memory[src2_addr].f, dest_addr,
             memory[dest_addr].f);
      break;
    case OP_HALT:
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      return;
    case OP_MOV:
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      memory[dest_addr] = memory[src1_addr];
      break;
    case OP_JMP:
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      pc = src1_addr; /* Jump to address */
      break;
    case OP_JGZ: {
      printf("ADD;src1[%d]:%d;src2[%d]:%d;dest[%d]:%d\n", src1_addr,
             memory[src1_addr].u, src2_addr, memory[src2_addr].u, dest_addr,
             memory[dest_addr].u);
      uint32_t value = memory[src1_addr].u;
      uint32_t jump_target = src2_addr;

      /* Branchless greater-than-zero check */
      int32_t mask = -((uint32_t)(value > 0));
      pc = (jump_target & mask) | (pc & ~mask);
      break;
    }
    case OP_PRINT_STRING: {
      printf("PRINT_STRING;");
      uint32_t string_addr = src1_addr;
      int i = 0;
      while (1) {
        int j = 0;
        uint32_t word = memory[string_addr + (i++)].u;
        for (j = 0; j < 4; j++) {
          char ch = (word >> (8 * j)) & 0xFF;
          if (ch == '\0')
            goto done;
          putchar(ch);
        }
      }
    done:
      putchar('\n');
      break;
    }
    case OP_READ_STRING: {
      printf("> ");
      uint32_t buffer_addr = dest_addr;
      int word_index = 0;
      int char_index = 0;

      while (1) {
        int ch = getchar();
        if (ch == '\n' || ch == EOF) {
          /* Store null terminator */
          uint32_t word = memory[buffer_addr + word_index].u;
          word = set_char(word, char_index, '\0');
          memory[buffer_addr + word_index].u = word;
          break;
        }
        uint32_t word = memory[buffer_addr + word_index].u;
        word = set_char(word, char_index, ch);
        memory[buffer_addr + word_index].u = word;

        char_index++;
        if (char_index == 4) {
          char_index = 0;
          word_index++;
        }
      }
      break;
    }
    default:
      printf("Unknown opcode: %d\n", opcode);
      return;
    }
  }
}

int main() {
  memory[0].u = OP_READ_STRING;
  memory[1].u = 0;
  memory[2].u = 0;
  memory[3].u = 104;
  memory[4].u = OP_ADD_F32;
  memory[5].u = 102;
  memory[6].u = 103;
  memory[7].u = 103;
  memory[8].u = OP_SUB;
  memory[9].u = 100;
  memory[10].u = 101;
  memory[11].u = 100;
  memory[12].u = OP_JGZ;
  memory[13].u = 100;
  memory[14].u = 4;
  memory[15].u = 4;
  memory[16].u = OP_PRINT_STRING;
  memory[17].u = 104;
  memory[18].u = 0;
  memory[19].u = 0;
  memory[20].u = OP_HALT;
  memory[100].u = 5;
  memory[101].u = 1;
  memory[102].f = 5.f;
  memory[103].f = 5.f;

  run_vm();

  printf("Dest at address 103: %f\n", memory[103].f);
  return 0;
}