#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/* #define MEMORY_SIZE 65536 /\* 64KB memory (adjustable) *\/ */
#define MEMORY_SIZE 1024
typedef union {
  float f;
  uint32_t u;
} Data;
Data 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_F32_TO_INT,    /* dest = src1 as int  */
  OP_INT_TO_F32,    /* dest = src1 as f32  */
  OP_MOV,           /* dest = src1	   */
  OP_JMP,           /* jump to address src1 unconditionally */
  OP_JGZ,           /* jump to address dest if src1 > 0 */
  OP_INT_TO_STRING, /* dest = src1 as str */
  OP_F32_TO_STRING, /* dest = src2 as str */
  OP_READ_STRING,
  OP_PRINT_STRING,
} Opcode;

int core_dump() {
  FILE *file = fopen("memory_dump.bin", "wb");
  if (!file) {
    perror("Failed to open file");
    return EXIT_FAILURE;
  }
  size_t written = fwrite(memory, 1, MEMORY_SIZE, file);
  if (written != MEMORY_SIZE) {
    fprintf(stderr, "Incomplete write: %zu bytes written out of %u\n", written,
            MEMORY_SIZE);
    fclose(file);
    return EXIT_FAILURE;
  }

  fclose(file);
  return EXIT_SUCCESS;
}

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));
}

/* Pack string into union-based memory */
void pack_string(const char *str, uint32_t length, uint32_t dest_addr) {
  memory[dest_addr].u = length;
  uint32_t buffer_addr = dest_addr + 1;
  int word_index = 0;
  int char_index = 0;

  uint32_t i = 0;
  while (i < length) {
    char ch = str[i++];
    if (ch == '\0') {
      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++;
    }
  }
}

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;
      break;

    case OP_SUB:
      memory[dest_addr].u = memory[src1_addr].u - memory[src2_addr].u;
      break;

    case OP_MUL:
      memory[dest_addr].u = memory[src1_addr].u * memory[src2_addr].u;
      break;

    case OP_DIV:
      memory[dest_addr].u = memory[src1_addr].u / memory[src2_addr].u;
      break;

    case OP_ADD_F32:
      memory[dest_addr].f = memory[src1_addr].f + memory[src2_addr].f;
      break;

    case OP_SUB_F32:
      memory[dest_addr].f = memory[src1_addr].f - memory[src2_addr].f;
      break;

    case OP_MUL_F32:
      memory[dest_addr].f = memory[src1_addr].f * memory[src2_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;
      break;
    case OP_F32_TO_INT: {
      float tmp = memory[src1_addr].f;
      memory[dest_addr].u = (uint32_t)tmp;
      break;
    }
    case OP_INT_TO_F32: {
      uint32_t tmp = memory[src1_addr].u;
      memory[dest_addr].f = (float)tmp;
      break;
    }
    case OP_HALT:
      return;
    case OP_MOV:
      memory[dest_addr] = memory[src1_addr];
      break;
    case OP_JMP:
      pc = src1_addr; /* Jump to address */
      break;
    case OP_JGZ: {
      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_INT_TO_STRING: {
      int32_t a = (int32_t)memory[src1_addr].u;
      char buffer[32];
      sprintf(buffer, "%d", a);
      pack_string(buffer, strlen(buffer), dest_addr);
      break;
    }
    case OP_F32_TO_STRING: {
      float a = memory[src1_addr].f;
      char buffer[32];
      sprintf(buffer, "%f", a);
      pack_string(buffer, strlen(buffer), dest_addr);
      break;
    }
    case OP_PRINT_STRING: {
      uint32_t string_addr = src1_addr;
      uint32_t length = memory[src1_addr - 1].u;
      uint32_t i;
      for (i = 0; i < length;) {
        uint32_t word = memory[string_addr + (i / 4)].u;
        uint8_t ch = get_char(word, i % 4);
        if (ch == '\0')
          break;
        putchar(ch);
        i++;
      }
      putchar('\n');
      break;
    }
    case OP_READ_STRING: {
      putchar('>');
      putchar(' ');
      uint32_t buffer_addr = dest_addr + 1;
      uint32_t length = 0;
      int word_index = 0;
      int char_index = 0;

      while (1) {
        int ch = getchar();
        if (ch == '\n' || ch == EOF) {
          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++;
        }
        length++;
      }
      memory[dest_addr].u = length;
      break;
    }
    default:
      printf("Unknown opcode: %d\n", opcode);
      return;
    }
  }
}

int main() {
  int i = 0;
  memory[i++].u = OP_ADD_F32;
  memory[i++].u = 102;
  memory[i++].u = 103;
  memory[i++].u = 103;
  memory[i++].u = OP_SUB;
  memory[i++].u = 100;
  memory[i++].u = 101;
  memory[i++].u = 100;
  memory[i++].u = OP_JGZ;
  memory[i++].u = 100;
  memory[i++].u = 0;
  memory[i++].u = 0;
  memory[i++].u = OP_F32_TO_INT;
  memory[i++].u = 103;
  memory[i++].u = 0;
  memory[i++].u = 103;
  memory[i++].u = OP_INT_TO_STRING;
  memory[i++].u = 103;
  memory[i++].u = 0;
  memory[i++].u = 104;
  memory[i++].u = OP_PRINT_STRING;
  memory[i++].u = 105;
  memory[i++].u = 0;
  memory[i++].u = 0;
  memory[i++].u = OP_READ_STRING;
  memory[i++].u = 0;
  memory[i++].u = 0;
  memory[i++].u = 109;
  memory[i++].u = OP_PRINT_STRING;
  memory[i++].u = 110;
  memory[i++].u = 0;
  memory[i++].u = 0;
  memory[i++].u = OP_HALT;
  memory[100].u = 5;
  memory[101].u = 1;
  memory[102].f = 5.f;
  memory[103].f = 5.f;

  run_vm();

  return core_dump();
}