#include "vm.hpp"

#include <cstdint>
#include <cstring>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <vector>

inline int32_t sign_extend(int32_t value, int bits) {
  int32_t mask = 1 << (bits - 1);
  return (value ^ mask) - mask;
}

std::vector<uint8_t> load_program(const std::string& filename,
                                  size_t memory_size) {
  std::vector<uint8_t> memory(memory_size, 0);

  std::ifstream file(filename, std::ios::binary | std::ios::ate);

  if (!file.is_open()) {
    throw std::runtime_error("Failed to open file: " + filename);
  }

  size_t file_size = file.tellg();

  if (file_size > memory_size) {
    throw std::runtime_error("File is too big");
  }

  file.seekg(0, std::ios::beg);

  file.read(reinterpret_cast<char*>(&memory[0]), file_size);

  if (!file) {
    throw std::runtime_error(
        "Failed to read the complete program into memory.");
  }

  file.close();

  return memory;
}

VM::VM(std::vector<uint8_t> memory) : memory_(memory) {}

std::vector<uint8_t> VM::read_memory(size_t start, size_t size) {
  if (start + size > memory_.size()) {
    return std::vector<uint8_t>(size, 0);
  }
  return std::vector<uint8_t>(memory_.begin() + start,
                              memory_.begin() + start + size);
}

uint32_t VM::read_register(size_t regnum) {
  if (regnum == 32) return pc;

  if (regnum >= NUM_REGISTERS) {
    throw std::runtime_error("Register out of range");
  }

  return registers[regnum];
}

void VM::step() {
  size_t memory_size = memory_.size();
  uint8_t* memory = &memory_[0];

  uint32_t instr = *(uint32_t*)&memory[pc];
  // std::cout << "pc: " << std::hex << pc << std::dec << "\n";
  //  std::cout << "instr: " << std::hex << instr << "\n";
  pc += 4;

  // Decode instruction
  uint32_t opcode = instr & 0x7F;
  uint32_t rd = (instr >> 7) & 0x1F;
  uint32_t funct3 = (instr >> 12) & 0x7;
  uint32_t rs1 = (instr >> 15) & 0x1F;
  uint32_t rs2 = (instr >> 20) & 0x1F;
  uint32_t funct7 = (instr >> 25);
  int32_t imm;

  switch (opcode) {
    case 0x33: {  // R-type
      if (funct7 == 0x00) {
        if (funct3 == 0x0) {  // ADD
          registers[rd] = registers[rs1] + registers[rs2];
        } else if (funct3 == 0x04) {  // XOR
          registers[rd] = registers[rs1] ^ registers[rs2];
        } else if (funct3 == 0x06) {  // OR
          registers[rd] = registers[rs1] | registers[rs2];
        } else if (funct3 == 0x07) {  // AND
          registers[rd] = registers[rs1] & registers[rs2];
        } else if (funct3 == 0x01) {  // SLL
          registers[rd] = registers[rs1] << registers[rs2];
        } else if (funct3 == 0x05) {  // SRL
          uint32_t value = registers[rs1];
          uint32_t shift_amount = registers[rs2] & 0x1F;
          registers[rd] = value << shift_amount;
        } else if (funct3 == 0x02) {  // SLT
          registers[rd] = (static_cast<int32_t>(registers[rs1]) <
                           static_cast<int32_t>(registers[rs2]))
                              ? 0
                              : 1;
        } else if (funct3 == 0x03) {  // SLTU
          registers[rd] = (registers[rs1] < registers[rs2]) ? 1 : 0;
        } else {
          throw std::runtime_error("Unknown R-type instruction");
        }
      } else if (funct7 == 0x20) {
        if (funct3 == 0x0) {  // SUB
          registers[rd] = registers[rs1] - registers[rs2];
        } else if (funct3 == 0x05) {  // SRA
          // Only the lower 5 bits are used for shift
          int32_t value = static_cast<int32_t>(registers[rs1]);
          int32_t shift_amount = registers[rs2] & 0x1F;
          registers[rd] = value >> shift_amount;
        } else {
          throw std::runtime_error("Unknown R-type instruction");
        }
      } else if (funct7 == 0x01) {
        if (funct3 == 0x0) {  // MUL
          int64_t result =
              static_cast<int64_t>(static_cast<int32_t>(registers[rs1])) *
              static_cast<int64_t>(static_cast<int32_t>(registers[rs2]));
          registers[rd] = static_cast<uint32_t>(result);
        } else if (funct3 == 0x1) {  // MULH
          int64_t result =
              static_cast<int64_t>(static_cast<int32_t>(registers[rs1])) *
              static_cast<int64_t>(static_cast<int32_t>(registers[rs2]));
          registers[rd] = static_cast<uint32_t>(result >> 32);
        } else if (funct3 == 0x2) {  // MULSU
          int64_t result =
              static_cast<int64_t>(static_cast<int32_t>(registers[rs1])) *
              static_cast<uint64_t>(registers[rs2]);
          registers[rd] = static_cast<uint32_t>(result >> 32);
        } else if (funct3 == 0x3) {  // MULU
          uint64_t result = static_cast<uint64_t>(registers[rs1]) *
                            static_cast<uint64_t>(registers[rs2]);
          registers[rd] = static_cast<uint32_t>(result >> 32);  // Upper 32 bits
        } else if (funct3 == 0x4) {                             // DIV
          int32_t dividend = static_cast<int32_t>(registers[rs1]);
          int32_t divisor = static_cast<int32_t>(registers[rs2]);
          if (divisor == 0) {
            registers[rd] = -1;  // Division by zero result
          } else if (dividend == INT32_MIN && divisor == -1) {
            registers[rd] = dividend;  // Overflow case
          } else {
            registers[rd] = dividend / divisor;
          }
        } else if (funct3 == 0x5) {  // DIVU
          uint32_t dividend = registers[rs1];
          uint32_t divisor = registers[rs2];
          registers[rd] = (divisor == 0) ? UINT32_MAX : dividend / divisor;
        } else if (funct3 == 0x6) {  // REM
          int32_t dividend = static_cast<int32_t>(registers[rs1]);
          int32_t divisor = static_cast<int32_t>(registers[rs2]);
          if (divisor == 0) {
            registers[rd] =
                dividend;  // Remainder with zero divisor is the dividend
          } else if (dividend == INT32_MIN && divisor == -1) {
            registers[rd] = 0;  // Overflow case
          } else {
            registers[rd] = dividend % divisor;
          }
        } else if (funct3 == 0x7) {  // REMU
          uint32_t dividend = registers[rs1];
          uint32_t divisor = registers[rs2];
          registers[rd] = (divisor == 0) ? dividend : dividend % divisor;
        } else {
          throw std::runtime_error("Unknown R-type instruction");
        }
      } else {
        throw std::runtime_error("Unknown R-type instruction");
      }
      break;
    }
    case 0x13: {                           // I-type (ADDI)
      imm = sign_extend(instr >> 20, 12);  // Extract 12-bit immediate
      if (funct3 == 0x0) {                 // ADDI
        registers[rd] = registers[rs1] + imm;
      } else {
        throw std::runtime_error("Unknown I-type instruction");
      }
      break;
    }
    case 0x63: {  // B-type (branches)
      imm = ((instr >> 7) & 0x1E) | ((instr >> 20) & 0x7E0) |
            ((instr >> 19) & 0x800) | ((instr >> 31) << 12);
      imm = sign_extend(imm, 13);  // Sign-extend 13-bit immediate
      if (funct3 == 0x0) {         // BEQ
        if (registers[rs1] == registers[rs2]) {
          pc += imm - 4;  // Offset PC (adjust for pre-increment)
        }
      } else if (funct3 == 0x1) {  // BNE
        if (registers[rs1] != registers[rs2]) {
          pc += imm - 4;  // Offset PC
        }
      } else if (funct3 == 0x4) {  // BLT
        if (static_cast<int32_t>(registers[rs1]) <
            static_cast<int32_t>(registers[rs2])) {
          pc += imm - 4;
        }
      } else if (funct3 == 0x5) {  // BGE
        if (static_cast<int32_t>(registers[rs1]) >=
            static_cast<int32_t>(registers[rs2])) {
          pc += imm - 4;
        }
      } else if (funct3 == 0x6) {  // BLTU
        if (registers[rs1] < registers[rs2]) pc += imm - 4;
      } else if (funct3 == 0x7) {  // BGEU
        if (registers[rs1] >= registers[rs2]) pc += imm - 4;
      } else {
        throw std::runtime_error("Unknown B-type instruction");
      }
      break;
    }
    case 0x03: {                           // I-type (loads)
      imm = sign_extend(instr >> 20, 12);  // Extract 12-bit immediate
      if (funct3 == 0x00) {                // LB
        uint32_t addr = registers[rs1] + imm;
        if (addr + 1 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        registers[rd] = 0;
        std::memcpy(&registers[rd], memory + addr, sizeof(uint8_t));
      } else if (funct3 == 0x01) {  // LH
        uint32_t addr = registers[rs1] + imm;
        if (addr + 2 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        registers[rd] = 0;
        std::memcpy(&registers[rd], memory + addr, sizeof(uint16_t));
      } else if (funct3 == 0x2) {  // LW
        uint32_t addr = registers[rs1] + imm;
        if (addr + 4 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        std::memcpy(&registers[rd], memory + addr, sizeof(uint32_t));
      } else {
        throw std::runtime_error("Unknown load instruction");
      }
      break;
    }
    case 0x23: {  // S-type (SW)
      imm = ((instr >> 7) & 0x1F) | (((instr >> 25) & 0x7F) << 5);
      imm = sign_extend(imm, 12);  // Sign-extend 12-bit immediate
      if (funct3 == 0x0) {         // SB
        uint32_t addr = registers[rs1] + imm;
        if (addr + 1 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        std::memcpy(memory + addr, &registers[rs2], sizeof(uint8_t));
      } else if (funct3 == 0x1) {  // SH
        uint32_t addr = registers[rs1] + imm;
        if (addr + 2 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        std::memcpy(memory + addr, &registers[rs2], sizeof(uint16_t));
      } else if (funct3 == 0x2) {  // SW
        uint32_t addr = registers[rs1] + imm;
        if (addr + 4 > memory_size) {
          throw std::runtime_error("Memory access out of bounds");
        }
        std::memcpy(memory + addr, &registers[rs2], sizeof(uint32_t));
      } else {
        throw std::runtime_error("Unknown store instruction");
      }
      break;
    }
    case 0x6F: {  // JAL
      int32_t offset =
          ((instr & 0x80000000) ? 0xFFF00000
                                : 0) |    // Sign-extension for imm[20]
          ((instr >> 21) & 0x3FF) << 1 |  // imm[10:1]
          ((instr >> 20) & 0x1) << 11 |   // imm[11]
          ((instr & 0xFF000));            // imm[19:12]

      registers[rd] = pc;  // Save return address
      pc += offset - 4;
      break;
    }
    case 0x67: {                       // JALR
      int32_t offset = (instr >> 20);  // Sign-extended 12-bit immediate
      uint32_t target =
          (registers[rs1] + offset) & ~1;  // Target address (LSB cleared)

      registers[rd] = pc;  // Save return address
      pc = target;
      break;
    }
    case 0x37: {                               // LUI
      uint32_t imm = (instr >> 12) & 0xFFFFF;  // Extract 20-bit immediate
      registers[rd] =
          imm
          << 12;  // Shift the immediate to the upper 20 bits of the register
      break;
    }
    case 0x17: {                               // AUIPC
      uint32_t imm = (instr >> 12) & 0xFFFFF;  // Extract 20-bit immediate
      registers[rd] =
          pc +
          (imm << 12);  // Add the immediate (shifted left) to the current PC
      break;
    }
    case 0x73: {        // EBREAK
      running = false;  // Flag to stop the emulator
      break;
    }
    default:
      throw std::runtime_error("Unknown opcode");
  }
}

void VM::eval() {
  while (running) {
    step();
  }
}