#include "vm.h"
#include <string.h>
#include <unistd.h>

int core_dump(Data *memory, uint32_t memory_size) {
  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));
}

void set_memory_char(Data *memory, uint8_t ch, uint32_t buffer_addr,
                     uint32_t length) {
  uint32_t word = memory[buffer_addr + (length / 4)].u;
  word = set_char(word, length % 4, ch);
  memory[buffer_addr + (length / 4)].u = word;
}

/* Pack string into union-based memory */
void pack_string(Data *memory, const char *str, uint32_t length,
                 uint32_t dest_addr) {
  memory[dest_addr].u = length;
  uint32_t buffer_addr = dest_addr + 1;
  uint32_t i = 0;
  while (i < length) {
    char ch = str[i];
    if (ch == '\0') {
      break;
    }
    set_memory_char(memory, ch, buffer_addr, i);
    i++;
  }
}

void run_vm(Data *memory, uint32_t memory_size) {
  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: {
      memory[dest_addr].u = (uint32_t)memory[src1_addr].f;
      break;
    }
    case OP_INT_TO_F32: {
      memory[dest_addr].f = (float)memory[src1_addr].u;
      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(memory, 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(memory, 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; i++) {
        uint32_t word = memory[string_addr + (i / 4)].u;
        uint8_t ch = get_char(word, i % 4);
        if (ch == '\0')
          break;
        putchar(ch);
      }
      putchar('\n');
      break;
    }
    case OP_READ_STRING: {
      uint32_t buffer_addr = dest_addr + 1;
      uint32_t length = 0;
      while (1) {
        int ch = getchar();
        if (ch == '\n' || ch == EOF) {
          set_memory_char(memory, '\0', buffer_addr, length);
          break;
        }
        set_memory_char(memory, ch, buffer_addr, length);
        length++;
      }
      memory[dest_addr].u = length;
      break;
    }
    case OP_CMP_STRING: {
      uint32_t addr1 = src1_addr;
      uint32_t addr2 = src2_addr;
      uint32_t length1 = memory[src1_addr - 1].u;
      uint32_t length2 = memory[src2_addr - 1].u;
      int equal = 1;

      if (length1 != length2) {
        equal = 0;
      } else {
        uint32_t i = 0;
        while (i < length1) {
          uint32_t word1 = memory[addr1 + i].u;
          uint32_t word2 = memory[addr2 + i].u;
          if (word1 != word2) {
            equal = 0;
            break;
          }
          if ((word1 & 0xFF) == '\0' && (word2 & 0xFF) == '\0')
            break;
        }
      }
      memory[dest_addr].u = equal;
      break;
    }
    default:
      printf("Unknown opcode: %d\n", opcode);
      return;
    }
  }
}