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

uint8_t UART::read_register(uint32_t address) {
  switch (address) {
    case UART_LSR:
      // Always ready to transmit
      return LSR_TRANSMITTER_EMPTY;
    default:
      return 0;
  }
}

void UART::write_register(uint32_t address, uint8_t value) {
  switch (address) {
    case UART_THR:
      std::cout.put(static_cast<char>(value));
      break;
  }
}

bool UART::is_transmitter_ready() {
  return read_register(UART_LSR) & LSR_TRANSMITTER_EMPTY;
}

VM::VM(const std::vector<uint8_t>& memory, const std::string& file_path)
    : memory_(memory), pc(PROGRAM_ADDR), file_path(file_path) {}

void VM::setreg(int regnum, uint32_t value) {
  if (regnum == 0) {
    return;
  }

  registers[regnum] = value;
}

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

uint32_t VM::read_memory_word(size_t pos) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos + 3 >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }
    return *(uint32_t*)&memory_[pos];
  } else {
    throw std::runtime_error("Memory access out of bounds");
  }
}

uint16_t VM::read_memory_half_word(size_t pos) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos + 1 >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }

    return *(uint16_t*)&memory_[pos];
  } else {
    throw std::runtime_error("Memory access out of bounds");
  }
}

uint8_t VM::read_memory_byte(size_t pos) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }

    return memory_[pos];
  } else {
    throw std::runtime_error("Memory access out of bounds");
  }
}

void VM::write_memory_word(size_t pos, uint32_t value) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos + 1 >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }
    *(uint32_t*)&memory_[pos] = value;
  } else {
    throw std::runtime_error("Memory access out of bounds");
  }
}

void VM::write_memory_half_word(size_t pos, uint16_t value) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos + 3 >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }
    *(uint16_t*)&memory_[pos] = value;
  } else {
    throw std::runtime_error("Memory access out of bounds");
  }
}

void VM::write_memory_byte(size_t pos, uint8_t value) {
  if (pos >= PROGRAM_ADDR) {
    pos -= PROGRAM_ADDR;
    if (pos >= memory_.size()) {
      throw std::runtime_error("Memory access out of bounds");
    }
    memory_[pos] = value;
  } else {
    if (is_mmap(pos, 1)) {
      if (pos >= UART_ADDR && pos < UART_ADDR + 8) {
        uart.write_register(pos - UART_ADDR, value);
      }
      return;
    }
  }
}

bool VM::is_mmap(size_t pos, size_t size) {
  if (pos + size < UART_ADDR) return false;
  if (pos >= UART_ADDR + 8) return false;

  return (pos < UART_ADDR + 8) && (pos + size >= UART_ADDR - 1);
}

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

const std::string& VM::get_file_path() { return file_path; }

void VM::step() {
  if (pc < PROGRAM_ADDR) throw std::runtime_error("PC out of range");
  uint32_t instr = *(uint32_t*)&memory_[pc - PROGRAM_ADDR];
  // 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
          setreg(rd, registers[rs1] + registers[rs2]);
        } else if (funct3 == 0x04) {  // XOR
          setreg(rd, registers[rs1] ^ registers[rs2]);
        } else if (funct3 == 0x06) {  // OR
          setreg(rd, registers[rs1] | registers[rs2]);
        } else if (funct3 == 0x07) {  // AND
          setreg(rd, registers[rs1] & registers[rs2]);
        } else if (funct3 == 0x01) {  // SLL
          setreg(rd, registers[rs1] << registers[rs2]);
        } else if (funct3 == 0x05) {  // SRL
          uint32_t value = registers[rs1];
          uint32_t shift_amount = registers[rs2] & 0x1F;
          setreg(rd, value >> shift_amount);
        } else if (funct3 == 0x02) {  // SLT
          setreg(rd, (static_cast<int32_t>(registers[rs1]) <
                      static_cast<int32_t>(registers[rs2]))
                         ? 0
                         : 1);
        } else if (funct3 == 0x03) {  // SLTU
          setreg(rd, (registers[rs1] < registers[rs2]) ? 1 : 0);
        } else {
          throw std::runtime_error("Unknown R-type instruction");
        }
      } else if (funct7 == 0x20) {
        if (funct3 == 0x0) {  // SUB
          setreg(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;
          setreg(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]));
          setreg(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]));
          setreg(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]);
          setreg(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]);
          setreg(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) {
            setreg(rd, -1);  // Division by zero result
          } else if (dividend == INT32_MIN && divisor == -1) {
            setreg(rd, dividend);  // Overflow case
          } else {
            setreg(rd, dividend / divisor);
          }
        } else if (funct3 == 0x5) {  // DIVU
          uint32_t dividend = registers[rs1];
          uint32_t divisor = registers[rs2];
          setreg(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) {
            setreg(rd,
                   dividend);  // Remainder with zero divisor is the dividend
          } else if (dividend == INT32_MIN && divisor == -1) {
            setreg(rd, 0);  // Overflow case
          } else {
            setreg(rd, dividend % divisor);
          }
        } else if (funct3 == 0x7) {  // REMU
          uint32_t dividend = registers[rs1];
          uint32_t divisor = registers[rs2];
          setreg(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 and friends)
      imm = sign_extend(instr >> 20, 12);  // Extract 12-bit immediate
      if (funct3 == 0x0) {                 // ADDI
        setreg(rd, registers[rs1] + imm);
      } else if (funct3 == 0x4) {  // XORI
        setreg(rd, registers[rs1] ^ imm);
      } else if (funct3 == 0x6) {  // ORI
        setreg(rd, registers[rs1] | imm);
      } else if (funct3 == 0x07) {  // ANDI
        setreg(rd, registers[rs1] & imm);
      } else if (funct3 == 0x01) {
        if (((imm >> 5) & 0x7f) == 0x0) {  // SLLI
          uint32_t value = registers[rs1];
          uint32_t shift_amount = imm & 0x1F;
          setreg(rd, value << shift_amount);
        } else {
          throw std::runtime_error("Unknown I-type instruction");
        }
      } else if (funct3 == 0x05) {
        if (((imm >> 5) & 0x7f) == 0x20) {  // SRAI
          int32_t value = static_cast<int32_t>(imm & 0x1f);
          int32_t shift_amount = imm & 0x1F;
          setreg(rd, value >> shift_amount);
        } else if (((imm >> 5) & 0x7f) == 0x0) {  // SRLI
          uint32_t value = registers[rs1];
          uint32_t shift_amount = imm & 0x1F;
          setreg(rd, value >> shift_amount);
        } else {
          throw std::runtime_error("Unknown I-type instruction");
        }
      } else if (funct3 == 0x02) {  // SLTI
        setreg(rd, (static_cast<int32_t>(registers[rs1]) <
                    static_cast<int32_t>(imm))
                       ? 0
                       : 1);
      } else if (funct3 == 0x03) {  // SLTIU
        setreg(rd, (registers[rs1] < static_cast<uint32_t>(imm)) ? 1 : 0);
      } else {
        throw std::runtime_error("Unknown I-type instruction");
      }
      break;
    }
    case 0x63: {  // B-type (branches)
      imm = ((int64_t)(int32_t)(instr & 0x80000000) >> 19) |
            ((instr & 0x80) << 4)      // imm[11]
            | ((instr >> 20) & 0x7e0)  // imm[10:5]
            | ((instr >> 7) & 0x1e);

      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;
        setreg(rd, read_memory_byte(addr));
      } else if (funct3 == 0x01) {  // LH
        uint32_t addr = registers[rs1] + imm;
        setreg(rd, read_memory_half_word(addr));
      } else if (funct3 == 0x2) {  // LW
        uint32_t addr = registers[rs1] + imm;
        setreg(rd, read_memory_word(addr));
      } else if (funct3 == 0x4) {  // LBU
        uint32_t addr = registers[rs1] + imm;
        setreg(rd, read_memory_byte(addr));
      } else if (funct3 == 0x5) {  // LHU
        uint32_t addr = registers[rs1] + imm;
        setreg(rd, read_memory_half_word(addr));
      } 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;
        write_memory_byte(addr, registers[rs2]);
      } else if (funct3 == 0x1) {  // SH
        uint32_t addr = registers[rs1] + imm;
        write_memory_half_word(addr, registers[rs2]);
      } else if (funct3 == 0x2) {  // SW
        uint32_t addr = registers[rs1] + imm;
        write_memory_word(addr, registers[rs2]);
      } 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]

      setreg(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)

      setreg(rd, pc);  // Save return address
      pc = target;
      break;
    }
    case 0x37: {                               // LUI
      uint32_t imm = (instr >> 12) & 0xFFFFF;  // Extract 20-bit immediate
      setreg(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
      setreg(rd,
             pc - 4 + (imm << 12));  // Add the immediate (shifted left) to
                                     // the current PC
      break;
    }
    case 0x73: {  // EBREAK
      pc -= 4;
      throw EbreakException();
      break;
    }
    default:
      throw std::runtime_error("Unknown opcode");
  }
}

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