#include <Core.hpp>
#include <log.hpp>

namespace n64 {
Cop0::Cop0() { Reset(); }

void Cop0::Reset() {
  cause.raw = 0xB000007C;
  status.raw = 0;
  status.cu0 = 1;
  status.cu1 = 1;
  status.fr = 1;
  PRId = 0x00000B22;
  Config = 0x7006E463;
  EPC = 0xFFFFFFFFFFFFFFFFll;
  ErrorEPC = 0xFFFFFFFFFFFFFFFFll;
  wired = 0;
  index.raw = 63;
  badVaddr = 0xFFFFFFFFFFFFFFFF;

  kernelMode = {true};
  supervisorMode = {false};
  userMode = {false};
  is64BitAddressing = {false};
  llbit = {};

  pageMask = {};
  entryHi = {};
  entryLo0 = {}, entryLo1 = {};
  context = {};
  wired = {}, r7 = {};
  count = {};
  compare = {};
  LLAddr = {}, WatchLo = {}, WatchHi = {};
  xcontext = {};
  r21 = {}, r22 = {}, r23 = {}, r24 = {}, r25 = {};
  ParityError = {}, CacheError = {}, TagLo = {}, TagHi = {};
  ErrorEPC = {};
  r31 = {};
  memset(tlb, 0, sizeof(TLBEntry) * 32);
  tlbError = NONE;
  openbus = {};
}

u32 Cop0::GetReg32(const u8 addr) {
  switch (addr) {
  case COP0_REG_INDEX:
    return index.raw & INDEX_MASK;
  case COP0_REG_RANDOM:
    return GetRandom();
  case COP0_REG_ENTRYLO0:
    return entryLo0.raw;
  case COP0_REG_ENTRYLO1:
    return entryLo1.raw;
  case COP0_REG_CONTEXT:
    return context.raw;
  case COP0_REG_PAGEMASK:
    return pageMask.raw;
  case COP0_REG_WIRED:
    return wired;
  case COP0_REG_BADVADDR:
    return badVaddr;
  case COP0_REG_COUNT:
    return GetCount();
  case COP0_REG_ENTRYHI:
    return entryHi.raw;
  case COP0_REG_COMPARE:
    return compare;
  case COP0_REG_STATUS:
    return status.raw;
  case COP0_REG_CAUSE:
    return cause.raw;
  case COP0_REG_EPC:
    return EPC;
  case COP0_REG_PRID:
    return PRId;
  case COP0_REG_CONFIG:
    return Config;
  case COP0_REG_LLADDR:
    return LLAddr;
  case COP0_REG_WATCHLO:
    return WatchLo;
  case COP0_REG_WATCHHI:
    return WatchHi;
  case COP0_REG_XCONTEXT:
    return xcontext.raw;
  case COP0_REG_PARITY_ERR:
    return ParityError;
  case COP0_REG_CACHE_ERR:
    return CacheError;
  case COP0_REG_TAGLO:
    return TagLo;
  case COP0_REG_TAGHI:
    return TagHi;
  case COP0_REG_ERROREPC:
    return ErrorEPC;
  case 7:
  case 21:
  case 22:
  case 23:
  case 24:
  case 25:
  case 31:
    return openbus;
  default:
    panic("Unsupported word read from COP0 register {}", addr);
  }
}

u64 Cop0::GetReg64(const u8 addr) const {
  switch (addr) {
  case COP0_REG_ENTRYLO0:
    return entryLo0.raw;
  case COP0_REG_ENTRYLO1:
    return entryLo1.raw;
  case COP0_REG_CONTEXT:
    return context.raw;
  case COP0_REG_BADVADDR:
    return badVaddr;
  case COP0_REG_ENTRYHI:
    return entryHi.raw;
  case COP0_REG_STATUS:
    return status.raw;
  case COP0_REG_EPC:
    return EPC;
  case COP0_REG_PRID:
    return PRId;
  case COP0_REG_LLADDR:
    return LLAddr;
  case COP0_REG_XCONTEXT:
    return xcontext.raw & 0xFFFFFFFFFFFFFFF0;
  case COP0_REG_ERROREPC:
    return ErrorEPC;
  case 7:
  case 21:
  case 22:
  case 23:
  case 24:
  case 25:
  case 31:
    return openbus;
  default:
    panic("Unsupported dword read from COP0 register {}", addr);
  }
}

void Cop0::SetReg32(const u8 addr, const u32 value) {
  openbus = value & 0xFFFFFFFF;
  switch (addr) {
  case COP0_REG_INDEX:
    index.raw = value & INDEX_MASK;
    break;
  case COP0_REG_RANDOM:
    break;
  case COP0_REG_ENTRYLO0:
    entryLo0.raw = value & ENTRY_LO_MASK;
    break;
  case COP0_REG_ENTRYLO1:
    entryLo1.raw = value & ENTRY_LO_MASK;
    break;
  case COP0_REG_CONTEXT:
    context.raw = (s64(s32(value)) & 0xFFFFFFFFFF800000) | (context.raw & 0x7FFFFF);
    break;
  case COP0_REG_PAGEMASK:
    pageMask.raw = value & PAGEMASK_MASK;
    break;
  case COP0_REG_WIRED:
    wired = value & 63;
    break;
  case COP0_REG_BADVADDR:
    break;
  case COP0_REG_COUNT:
    count = (u64)value << 1;
    break;
  case COP0_REG_ENTRYHI:
    entryHi.raw = s64(s32(value)) & ENTRY_HI_MASK;
    break;
  case COP0_REG_COMPARE:
    compare = value;
    cause.ip7 = false;
    break;
  case COP0_REG_STATUS:
    status.raw &= ~STATUS_MASK;
    status.raw |= (value & STATUS_MASK);
    Update();
    break;
  case COP0_REG_CAUSE:
    {
      Cop0Cause tmp{};
      tmp.raw = value;
      cause.ip0 = tmp.ip0;
      cause.ip1 = tmp.ip1;
    }
    break;
  case COP0_REG_EPC:
    EPC = s64(s32(value));
    break;
  case COP0_REG_PRID:
    break;
  case COP0_REG_CONFIG:
    Config &= ~CONFIG_MASK;
    Config |= (value & CONFIG_MASK);
    break;
  case COP0_REG_LLADDR:
    LLAddr = value;
    break;
  case COP0_REG_WATCHLO:
    WatchLo = value;
    break;
  case COP0_REG_WATCHHI:
    WatchHi = value;
    break;
  case COP0_REG_XCONTEXT:
    xcontext.raw = (s64(s32(value)) & 0xFFFFFFFE00000000) | (xcontext.raw & 0x1FFFFFFFF);
    break;
  case COP0_REG_PARITY_ERR:
    ParityError = value & 0xff;
    break;
  case COP0_REG_CACHE_ERR:
    break;
  case COP0_REG_TAGLO:
    TagLo = value;
    break;
  case COP0_REG_TAGHI:
    TagHi = value;
    break;
  case COP0_REG_ERROREPC:
    ErrorEPC = s64(s32(value));
    break;
  case 7:
  case 21:
  case 22:
  case 23:
  case 24:
  case 25:
  case 31:
    break;
  default:
    panic("Unsupported word write from COP0 register {}", addr);
  }
}

void Cop0::SetReg64(const u8 addr, const u64 value) {
  openbus = value;
  switch (addr) {
  case COP0_REG_ENTRYLO0:
    entryLo0.raw = value & ENTRY_LO_MASK;
    break;
  case COP0_REG_ENTRYLO1:
    entryLo1.raw = value & ENTRY_LO_MASK;
    break;
  case COP0_REG_CONTEXT:
    context.raw = (value & 0xFFFFFFFFFF800000) | (context.raw & 0x7FFFFF);
    break;
  case COP0_REG_XCONTEXT:
    xcontext.raw = (value & 0xFFFFFFFE00000000) | (xcontext.raw & 0x1FFFFFFFF);
    break;
  case COP0_REG_ENTRYHI:
    entryHi.raw = value & ENTRY_HI_MASK;
    break;
  case COP0_REG_STATUS:
    status.raw = value;
    break;
  case COP0_REG_CAUSE:
    {
      Cop0Cause tmp{};
      tmp.raw = value;
      cause.ip0 = tmp.ip0;
      cause.ip1 = tmp.ip1;
    }
    break;
  case COP0_REG_BADVADDR:
    break;
  case COP0_REG_EPC:
    EPC = (s64)value;
    break;
  case COP0_REG_LLADDR:
    LLAddr = value;
    break;
  case COP0_REG_ERROREPC:
    ErrorEPC = (s64)value;
    break;
  default:
    panic("Unsupported dword write to COP0 register {}", addr);
  }
}

static FORCE_INLINE u64 getVPN(const u64 addr, const u64 pageMask) {
  const u64 mask = pageMask | 0x1fff;
  const u64 vpn = addr & 0xFFFFFFFFFF | addr >> 22 & 0x30000000000;

  return vpn & ~mask;
}

TLBEntry *Cop0::TLBTryMatch(const u64 vaddr, int &index) {
  for (int i = 0; i < 32; i++) {
    TLBEntry *entry = &tlb[i];
    if (!entry->initialized)
      continue;
    
    const u64 entry_vpn = getVPN(entry->entryHi.raw, entry->pageMask.raw);
    const u64 vaddr_vpn = getVPN(vaddr, entry->pageMask.raw);

    const bool vpn_match = entry_vpn == vaddr_vpn;
    const bool asid_match = entry->global || entryHi.asid == entry->entryHi.asid;

    if(!vpn_match || !asid_match)
      continue;

    index = i;
    return entry;
  }

  return nullptr;
}

TLBEntry *Cop0::TLBTryMatch(const u64 vaddr) {
  for (auto &t : tlb) {
    TLBEntry *entry = &t;
    if (!entry->initialized)
      continue;
    
    const u64 entry_vpn = getVPN(entry->entryHi.raw, entry->pageMask.raw);
    const u64 vaddr_vpn = getVPN(vaddr, entry->pageMask.raw);

    const bool vpn_match = entry_vpn == vaddr_vpn;
    const bool asid_match = entry->global || entryHi.asid == entry->entryHi.asid;
    
    if (vpn_match && asid_match)
      return entry;
  }

  return nullptr;
}

bool Cop0::ProbeTLB(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {
  const TLBEntry *entry = TLBTryMatch(vaddr);
  if (!entry) {
    tlbError = MISS;
    return false;
  }

  const u32 mask = entry->pageMask.mask << 12 | 0xFFF;
  const u32 odd = vaddr & mask + 1;

  const EntryLo entryLo = odd ? entry->entryLo1 : entry->entryLo0;

  if (!entryLo.v) {
    tlbError = INVALID;
    return false;
  }

  if (accessType == STORE && !entryLo.d) {
    tlbError = MODIFICATION;
    return false;
  }

  paddr = entryLo.pfn << 12 | vaddr & mask;

  return true;
}

void Cop0::FireException(const ExceptionCode code, const int cop, s64 pc) {
  Registers& regs = Core::GetRegs();

  u16 vectorOffset = 0x0180;
  if(tlbError == MISS && (code == ExceptionCode::TLBLoad || code == ExceptionCode::TLBStore)) {
    if(!status.exl) {
      if(is64BitAddressing) vectorOffset = 0x0080;
      else vectorOffset = 0x0000;
    }
  }

  cause.copError = cop;
  cause.exceptionCode = static_cast<u8>(code);

  if (!status.exl) {
    if ((cause.branchDelay = regs.prevDelaySlot)) {
      pc -= 4;
    }

    status.exl = true;
    EPC = pc;
  }

  if (status.bev) {
    panic("BEV bit set!");
  }

  regs.SetPC32(s32(0x80000000 + vectorOffset));
  Update();
}

void Cop0::HandleTLBException(const u64 vaddr) {
  const u64 vpn2 = vaddr >> 13 & 0x7FFFF;
  const u64 xvpn2 = vaddr >> 13 & 0x7FFFFFF;
  badVaddr = vaddr;
  context.badvpn2 = vpn2;
  xcontext.badvpn2 = xvpn2;
  xcontext.r = vaddr >> 62 & 3;
  entryHi.vpn2 = xvpn2;
  entryHi.r = vaddr >> 62 & 3;
}

ExceptionCode Cop0::GetTLBExceptionCode(const TLBError error, const TLBAccessType accessType) {
  switch (error) {
  case NONE:
    panic("Getting TLB exception with error NONE");
  case INVALID:
  case MISS:
    return accessType == LOAD ? ExceptionCode::TLBLoad : ExceptionCode::TLBStore;
  case MODIFICATION:
    return ExceptionCode::TLBModification;
  case DISALLOWED_ADDRESS:
    return accessType == LOAD ? ExceptionCode::AddressErrorLoad : ExceptionCode::AddressErrorStore;
  default:
    panic("Getting TLB exception for unknown error code! ({})", static_cast<u8>(error));
    return {};
  }
}

void Cop0::decode(const Instruction instr) {
  Registers& regs = Core::GetRegs();
  switch (instr.cop_rs()) {
    case 0x00: mfc0(instr); break;
    case 0x01: dmfc0(instr); break;
    case 0x04: mtc0(instr); break;
    case 0x05: dmtc0(instr); break;
    case 0x10 ... 0x1F:
      switch (instr.cop_funct()) {
        case 0x01: tlbr(); break;
        case 0x02: tlbw(index.i); break;
        case 0x06: tlbw(GetRandom()); break;
        case 0x08: tlbp(); break;
        case 0x18: eret(); break;
        default:
          panic("Unimplemented COP0 function {} ({:08X}) ({:016X})", instr.cop_funct(), u32(instr),
                      regs.oldPC);
      }
      break;
    default:
      panic("Unimplemented COP0 instruction {}", instr.cop_rs());
  }
}


template <>
bool Cop0::MapVirtualAddress<u32, true>(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {
  if(Util::IsInsideRange(vaddr, START_VREGION_KUSEG, END_VREGION_KUSEG))
    return ProbeTLB(accessType, s64(s32(vaddr)), paddr);
    
  tlbError = DISALLOWED_ADDRESS;
  return false;
}

template <>
bool Cop0::MapVirtualAddress<u32, false>(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {
  u8 segment = static_cast<u32>(vaddr) >> 29 & 7;
  if(Util::IsInsideRange(segment, 0, 3) || segment == 7)
    return ProbeTLB(accessType, static_cast<s32>(vaddr), paddr);
  
  if(Util::IsInsideRange(segment, 4, 5)) {
    paddr = vaddr & 0x1FFFFFFF;
    return true;
  }

  if(segment == 6)
    panic("Unimplemented virtual mapping in KSSEG! ({:08X})", vaddr);

  panic("Should never end up in base case in MapVirtualAddress! ({:08X})", vaddr);

  return false;
}

template <>
bool Cop0::MapVirtualAddress<u64, true>(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {
  if(Util::IsInsideRange(vaddr, 0x0000000000000000, 0x000000FFFFFFFFFF))
    return ProbeTLB(accessType, vaddr, paddr);
  
  tlbError = DISALLOWED_ADDRESS;
  return false;
}

template <>
bool Cop0::MapVirtualAddress<u64, false>(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {
  if(Util::IsInsideRange(vaddr, 0x0000000000000000, 0x000000FFFFFFFFFF) || // VREGION_XKUSEG
     Util::IsInsideRange(vaddr, 0x4000000000000000, 0x400000FFFFFFFFFF) || // VREGION_XKSSEG
     Util::IsInsideRange(vaddr, 0xC000000000000000, 0xC00000FF7FFFFFFF) || // VREGION_XKSEG
     Util::IsInsideRange(vaddr, 0xFFFFFFFFE0000000, 0xFFFFFFFFFFFFFFFF)) // VREGION_CKSEG3
    return ProbeTLB(accessType, vaddr, paddr);
  
  if(Util::IsInsideRange(vaddr, 0x8000000000000000, 0xBFFFFFFFFFFFFFFF)) { // VREGION_XKPHYS
    if (!kernelMode)
      panic("Access to XKPHYS address 0x{:016X} when outside kernel mode!", vaddr);
    
    const u8 high_two_bits = (vaddr >> 62) & 0b11;
    if (high_two_bits != 0b10)
      panic("Access to XKPHYS address 0x{:016X} with high two bits != 0b10!", vaddr);

    const u8 subsegment = (vaddr >> 59) & 0b11;
    bool cached = subsegment != 2; // do something with this eventually
    // If any bits in the range of 58:32 are set, the address is invalid.
    const bool valid = (vaddr & 0x07FFFFFF00000000) == 0;
    if (!valid) {
      tlbError = DISALLOWED_ADDRESS;
      return false;
    }

    paddr = vaddr & 0xFFFFFFFF;
    return true;
  }

  if(Util::IsInsideRange(vaddr, 0xFFFFFFFF80000000, 0xFFFFFFFF9FFFFFFF) || // VREGION_CKSEG0
     Util::IsInsideRange(vaddr, 0xFFFFFFFFA0000000, 0xFFFFFFFFBFFFFFFF)) { // VREGION_CKSEG1
    u32 cut = u32(vaddr) >> 28;
    u32 num = cut == 0xA;
    // Identical to ksegX in 32 bit mode.
    // Unmapped translation. Subtract the base address of the space to get the physical address.
    paddr = vaddr - (cut << 28); // Implies cutting off the high 32 bits
    trace("CKSEG{}: Translated 0x{:016X} to 0x{:08X}", num, vaddr, paddr);
    return true;
  }

  if(Util::IsInsideRange(vaddr, 0x0000010000000000, 0x3FFFFFFFFFFFFFFF) || // VREGION_XBAD1
     Util::IsInsideRange(vaddr, 0x4000010000000000, 0x7FFFFFFFFFFFFFFF) || // VREGION_XBAD2
     Util::IsInsideRange(vaddr, 0xC00000FF80000000, 0xFFFFFFFF7FFFFFFF)) { // VREGION_XBAD3
    tlbError = DISALLOWED_ADDRESS;
    return false;
  }

  panic("Resolving virtual address 0x{:016X} in 64 bit mode", vaddr);
  return false; // just to silence warning
}

bool Cop0::MapVAddr(const TLBAccessType accessType, const u64 vaddr, u32 &paddr) {  
  if(supervisorMode)
    panic("Supervisor mode memory access");

  if (is64BitAddressing) [[unlikely]] {
    if (kernelMode) [[likely]] return MapVirtualAddress<u64, false>(accessType, vaddr, paddr);
    if (userMode) return MapVirtualAddress<u64, true>(accessType, vaddr, paddr);

    panic("Unknown mode! This should never happen!");
  }

  if (kernelMode) [[likely]] return MapVirtualAddress<u32, false>(accessType, vaddr, paddr);
  if (userMode) return MapVirtualAddress<u32, true>(accessType, vaddr, paddr);

  panic("Unknown mode! This should never happen!");
}
} // namespace n64