#include <Core.hpp>
#include <jit/helpers.hpp>

namespace n64 {
#ifndef __aarch64__
JIT::JIT(ParallelRDP& parallel, Mem& mem, Registers& regs) : mem(mem), regs(regs) {
  regs.SetJIT(this);
  blockCache.resize(kUpperSize);
  if (cs_open(CS_ARCH_MIPS, static_cast<cs_mode>(CS_MODE_MIPS64 | CS_MODE_BIG_ENDIAN), &disassemblerMips) !=
      CS_ERR_OK) {
    panic("Failed to initialize MIPS disassembler");
  }

  if (cs_open(CS_ARCH_X86, static_cast<cs_mode>(CS_MODE_64 | CS_MODE_LITTLE_ENDIAN), &disassemblerX86) != CS_ERR_OK) {
    panic("Failed to initialize x86 disassembler");
  }
}

bool JIT::ShouldServiceInterrupt() const {
  const bool interrupts_pending = (regs.cop0.status.im & regs.cop0.cause.interruptPending) != 0;
  const bool interrupts_enabled = regs.cop0.status.ie == 1;
  const bool currently_handling_exception = regs.cop0.status.exl == 1;
  const bool currently_handling_error = regs.cop0.status.erl == 1;

  return interrupts_pending && interrupts_enabled && !currently_handling_exception && !currently_handling_error;
}

void JIT::CheckCompareInterrupt() {
  regs.cop0.count++;
  regs.cop0.count &= 0x1FFFFFFFF;
  if (regs.cop0.count == static_cast<u64>(regs.cop0.compare) << 1) {
    regs.cop0.cause.ip7 = 1;
    Core::GetMem().mmio.mi.UpdateInterrupt();
  }
}

void JIT::InvalidateBlock(const u32 paddr) {
  if (const u32 index = paddr >> kUpperShift; !blockCache[index].empty())
    blockCache[index] = {};
}

u32 JIT::FetchInstruction() {
  u32 paddr = 0;

  if (check_address_error(0b11, u64(blockPC))) [[unlikely]] {
    /*regs.cop0.HandleTLBException(blockPC);
    regs.cop0.FireException(ExceptionCode::AddressErrorLoad, 0, blockPC);
    return 1;*/
    Util::Error::GetInstance().Throw(
      {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::UNHANDLED_EXCEPTION}, blockPC, {},
      "[JIT]: Unhandled exception ADL due to unaligned PC virtual value!");
    return 0;
  }

  if (!regs.cop0.MapVAddr(Cop0::LOAD, blockPC, paddr)) {
    /*regs.cop0.HandleTLBException(blockPC);
    regs.cop0.FireException(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD), 0, blockPC);
    return 1;*/
    Util::Error::GetInstance().Throw(
      {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::UNHANDLED_EXCEPTION}, blockPC, {},
      "[JIT]: Unhandled exception TLB exception {} when retrieving PC physical address!",
      static_cast<int>(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD)));
    return 0;
  }

  return Core::GetMem().Read<u32>(paddr);
}

void JIT::SetPC32(s32 val) {
  code.mov(code.SCR1, REG(qword, pc));
  code.mov(REG(qword, oldPC), code.SCR1);
  code.mov(code.SCR1, s64(val));
  code.mov(REG(qword, pc), code.SCR1);
  code.mov(code.SCR1, s64(val) + 4);
  code.mov(REG(qword, nextPC), code.SCR1);
}

void JIT::SetPC64(s64 val) {
  code.mov(code.SCR1, REG(qword, pc));
  code.mov(REG(qword, oldPC), code.SCR1);
  code.mov(code.SCR1, val);
  code.mov(REG(qword, pc), code.SCR1);
  code.mov(code.SCR1, val + 4);
  code.mov(REG(qword, nextPC), code.SCR1);
}

void JIT::SetPC32(const Xbyak::Reg32& val) {
  code.mov(code.SCR1, REG(qword, pc));
  code.mov(REG(qword, oldPC), code.SCR1);
  code.movsxd(val.cvt64(), val);
  code.mov(REG(qword, pc), val);
  code.add(val, 4);
  code.mov(REG(qword, nextPC), val);
}

void JIT::SetPC64(const Xbyak::Reg64& val) {
  code.mov(code.SCR1, REG(qword, pc));
  code.mov(REG(qword, oldPC), code.SCR1);
  code.mov(REG(qword, pc), val);
  code.add(val, 4);
  code.mov(REG(qword, nextPC), val);
}

int JIT::Step() {
  blockOldPC = regs.oldPC;
  blockPC = regs.pc;
  blockNextPC = regs.nextPC;
  u32 paddr = 0;

  if (!regs.cop0.MapVAddr(Cop0::LOAD, blockPC, paddr)) {
    /*regs.cop0.HandleTLBException(blockPC);
    regs.cop0.FireException(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD), 0, blockPC);
    return 1;*/
    Util::Error::GetInstance().Throw(
      {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::UNHANDLED_EXCEPTION}, blockPC, {},
      "[JIT]: Unhandled exception TLB exception {} when retrieving PC physical address!",
      static_cast<int>(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD)));
    return 0;
  }

  u32 upperIndex = paddr >> kUpperShift;
  u32 lowerIndex = paddr & kLowerMask;

  if (!blockCache[upperIndex].empty()) {
    if (blockCache[upperIndex][lowerIndex]) {
      // trace("[JIT]: Executing already compiled block @ 0x{:016X}", blockPC);
      return blockCache[upperIndex][lowerIndex]();
    }
  } else {
    blockCache[upperIndex].resize(kLowerSize);
  }

  trace("[JIT]: Compiling block @ 0x{:016X}:", blockPC);
  const auto blockInfo = code.getCurr();
  const auto block = code.getCurr<BlockFn>();
  blockCache[upperIndex][lowerIndex] = block;

  code.setProtectModeRW();

  u32 instructionsInBlock = 0;
  u32 instruction = 0;

  bool instrEndsBlock = false;

  code.sub(code.rsp, 8);
  code.push(code.rbp);
  code.mov(code.rbp, reinterpret_cast<uintptr_t>(this)); // Load context pointer

  //cs_insn *insn;
  trace("\tMIPS code (guest PC = 0x{:016X}):", blockPC);
  while (!instrEndsBlock) {
    // CheckCompareInterrupt();
    paddr = 0;

    if (check_address_error(0b11, u64(blockPC))) [[unlikely]] {
      /*regs.cop0.HandleTLBException(blockPC);
      regs.cop0.FireException(ExceptionCode::AddressErrorLoad, 0, blockPC);
      return 1;*/
      Util::Error::GetInstance().Throw(
        {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::UNHANDLED_EXCEPTION}, blockPC, {},
        "[JIT]: Unhandled exception ADL due to unaligned PC virtual value!");
      return 0;
    }

    if (!regs.cop0.MapVAddr(Cop0::LOAD, blockPC, paddr)) {
      /*regs.cop0.HandleTLBException(blockPC);
      regs.cop0.FireException(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD), 0, blockPC);
      return 1;*/
      Util::Error::GetInstance().Throw(
        {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::UNHANDLED_EXCEPTION}, blockPC, {},
        "[JIT]: Unhandled exception TLB exception {} when retrieving PC physical address!",
        static_cast<int>(Cop0::GetTLBExceptionCode(regs.cop0.tlbError, Cop0::LOAD)));
      return 0;
    }

    instruction = Core::GetMem().Read<u32>(paddr);
    instructionsInBlock++;

    blockOldPC = blockPC;
    blockPC = blockNextPC;
    blockNextPC += 4;

    if((instrEndsBlock = InstrEndsBlock(instruction)))
      continue;

    trace("{}", Disassembler::GetInstance().DisassembleSimple(paddr, instruction).full);

    /*if(ShouldServiceInterrupt()) {
      regs.cop0.FireException(ExceptionCode::Interrupt, 0, blockPC);
      return 1;
    }*/

    Emit(instruction);
  }

  instructionsInBlock++;
  const u32 delay_instruction = FetchInstruction();
  instrEndsBlock = InstrEndsBlock(delay_instruction);
  if(instrEndsBlock) {
    Util::Error::GetInstance().Throw(
      {Util::Error::Severity::NON_FATAL}, {Util::Error::Type::JIT_BRANCH_INSIDE_DELAY_SLOT},
      blockPC, {}, "[JIT]: Unhandled case of branch from delay slot!");
    return 0;
  }

  blockOldPC = blockPC;
  blockPC = blockNextPC;
  blockNextPC += 4;

  Emit(delay_instruction);

  Emit(instruction);
  
  if(!regs.delaySlot) {
    code.mov(code.SCR1, blockOldPC);
    code.mov(REG(qword, oldPC), code.SCR1);
    code.mov(code.SCR1, blockPC);
    code.mov(REG(qword, pc), code.SCR1);
    code.mov(code.SCR1, blockNextPC);
    code.mov(REG(qword, nextPC), code.SCR1);
  }

  code.mov(code.rax, instructionsInBlock);
  code.pop(code.rbp);
  code.add(code.rsp, 8);
  code.ret();
  code.setProtectModeRE();
  //static auto blockInfoSize = 0;
  //blockInfoSize = code.getSize() - blockInfoSize;

  //trace("\tX86 code (block address = 0x{:016X}):", (uintptr_t)block);
  //auto count = cs_disasm(disassemblerX86, blockInfo, blockInfoSize, (uintptr_t)block, 0, &insn);
  //if (count > 0) {
  //  for (size_t j = 0; j < count; j++) {
  //    trace("\t\t0x{:016X}:\t{}\t\t{}\n", insn[j].address, insn[j].mnemonic, insn[j].op_str);
  //  }
  //
  //  cs_free(insn, count);
  //}
  // const auto dump = code.getCode();
  // Util::WriteFileBinary(dump, code.getSize(), "jit.dump");
  //  panic("");
  return block();
}
#endif
} // namespace n64