#include <errno.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)
uint32_t memory[MEMORY_SIZE]; // Memory array
#define DEBUG_PRINT                                                            \
  printf("dest[%d]=%d, src1[%d]=%d, src2[%d]=%d\n", dest_addr,                 \
         memory[dest_addr], src1_addr, memory[src1_addr], src2_addr,           \
         memory[src2_addr]);

typedef enum {
  OP_HALT, // terminate execution
  OP_ADD,  // dest = src1 + src2
  OP_SUB,  // 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) {
    // Fetch instruction
    Opcode opcode = memory[pc];
    uint32_t src1_addr = memory[pc + 1];
    uint32_t src2_addr = memory[pc + 2];
    uint32_t dest_addr = memory[pc + 3];
    pc += 4; // Advance to next instruction

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

    printf("opcode: ");
    // Execute instruction
    switch (opcode) {
    case OP_ADD:
      printf("ADD, ");
      DEBUG_PRINT
      memory[dest_addr] = memory[src1_addr] + memory[src2_addr];
      break;
    case OP_SUB:
      printf("SUB, ");
      DEBUG_PRINT
      memory[dest_addr] = memory[src1_addr] - memory[src2_addr];
      break;
    case OP_HALT:
      printf("HALT, ");
      return;
    case OP_MOV:
      printf("MOV, ");
      DEBUG_PRINT
      memory[dest_addr] = memory[src1_addr];
      break;
    case OP_JMP:
      printf("JMP, ");
      DEBUG_PRINT
      pc = src1_addr; // Jump to address
      break;
    case OP_JGZ: {
      printf("JGZ, ");
      DEBUG_PRINT
      uint32_t value = memory[src1_addr];
      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_STR, ");
      uint32_t string_addr = src1_addr;
      int i = 0;
      while (1) {
        uint32_t word = memory[string_addr + (i++)];
        for (int 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("READ_STR, ");
      uint32_t buffer_addr = dest_addr;
      char buffer[4];
      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];
          word = set_char(word, char_index, '\0');
          memory[buffer_addr + word_index] = word;
          break;
        }
        uint32_t word = memory[buffer_addr + word_index];
        word = set_char(word, char_index, ch);
        memory[buffer_addr + word_index] = word;

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

int main() {
  // Initialize memory
  memory[0] = OP_READ_STRING;
  memory[1] = 0;   // unused
  memory[2] = 0;   // unused
  memory[3] = 104; // dest
  memory[4] = OP_ADD;
  memory[5] = 102; // A (src1)
  memory[6] = 103; // B (src2)
  memory[7] = 103; // C (dest)
  memory[8] = OP_SUB;
  memory[9] = 100;  // counter (src1)
  memory[10] = 101; // value (src2)
  memory[11] = 100; // counter (dest)
  memory[12] = OP_JGZ;
  memory[13] = 100; // (src1)
  memory[14] = 4;   // (src2)
  memory[15] = 4;   // (dest)
  memory[16] = OP_PRINT_STRING;
  memory[17] = 104;     // String address
  memory[18] = 0;       // Unused
  memory[19] = 0;       // Unused
  memory[20] = OP_HALT; // Terminate after ADD
  memory[100] = 5;      // Value of A
  memory[101] = 1;      // Value of B
  memory[102] = 5;
  memory[103] = 5;
  /* memcpy(&memory[104], "hell", 4); */
  /* memcpy(&memory[105], "o\n\0\0", 4); */

  run_vm();

  printf("Result at address 103: %u\n", memory[103]); // Output: 12
  return 0;
}