// Copyright 2026 Dolphin Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#include <array>
#include <memory>

#include "Common/x64ABI.h"
#include "Core/ConfigManager.h"
#include "Core/Core.h"
#include "Core/HW/Memmap.h"
#include "Core/PowerPC/Jit64/Jit.h"
#include "Core/PowerPC/Jit64Common/Jit64AsmCommon.h"
#include "Core/PowerPC/Jit64Common/Jit64Constants.h"
#include "Core/PowerPC/MMU.h"
#include "Core/PowerPC/PowerPC.h"
#include "Core/System.h"

#include <gtest/gtest.h>

namespace
{
using namespace Gen;

// Execute the real SafeLoad/SafeWrite emitter, including its fallback and register contract.
// Testing only the C++ MMU helpers cannot detect corruption introduced at the native call site.
class DCacheCode : public CommonAsmRoutines
{
public:
  explicit DCacheCode(Core::System& system) : CommonAsmRoutines(jit), jit(system)
  {
    jit.jo = {};
    jit.js = {};
    AllocCodeSpace(512 * 1024);
    old_read = dcache32_read_hit_dbat;
    old_write = dcache32_write_hit_dbat;
    dcache32_read_hit_dbat = AlignCode4();
    GenDCache32Hit(false);
    dcache32_write_hit_dbat = AlignCode4();
    GenDCache32Hit(true);
  }

  ~DCacheCode() override
  {
    dcache32_read_hit_dbat = old_read;
    dcache32_write_hit_dbat = old_write;
  }

  void Access(bool write, u32 address, u32 value, X64Reg address_reg, X64Reg value_reg,
              s32 offset = 0, int flags = 0)
  {
    run = reinterpret_cast<void (*)()>(AlignCode4());
    ABI_PushRegistersAndAdjustStack(ABI_ALL_CALLEE_SAVED, 8);
    MOV(64, R(RPPCSTATE), ImmPtr(reinterpret_cast<u8*>(&jit.m_ppc_state) + 0x80));
    MOV(32, R(RDX), Imm32(0x12345678));
    MOV(32, R(RCX), Imm32(0x87654321));
    MOV(32, R(address_reg), Imm32(address));
    if (write)
      MOV(32, R(value_reg), Imm32(value));
    const BitSet32 live{RCX, RDX, R8, R9};
    if (flags & SAFE_LOADSTORE_NO_PROLOG)
      SUB(64, R(RSP), Imm8(8));
    if (write)
      SafeWriteRegToReg(value_reg, address_reg, 32, offset, live, flags);
    else
      SafeLoadToReg(value_reg, R(address_reg), 32, offset, live, false, flags);
    if (flags & SAFE_LOADSTORE_NO_PROLOG)
      ADD(64, R(RSP), Imm8(8));

    MOV(64, R(R11), ImmPtr(result.data()));
    for (const auto reg : {RAX, RCX, RDX, R8, R9})
      MOV(64, MDisp(R11, static_cast<int>(reg) * sizeof(u64)), R(reg));
    ABI_PopRegistersAndAdjustStack(ABI_ALL_CALLEE_SAVED, 8);
    RET();
    ASSERT_FALSE(HasWriteFailed());
    run();
  }

  std::array<u64, 16> result{};
  void (*run)() = nullptr;
  Jit64 jit;
  const u8* old_read;
  const u8* old_write;
};

class Jit64DCache : public testing::Test
{
protected:
  static void SetUpTestSuite()
  {
    SConfig::Init();
    auto& system = Core::System::GetInstance();
    system.SetIsWii(true);
    system.GetMemory().Init();
    Core::DeclareAsCPUThread();
    system.GetPPCState().dCache.Init(system.GetMemory());
  }

  static void TearDownTestSuite()
  {
    auto& system = Core::System::GetInstance();
    system.GetPPCState().m_enable_dcache = false;
    system.GetMemory().Shutdown();
    system.SetIsWii(false);
    Core::UndeclareAsCPUThread();
    SConfig::Shutdown();
  }

  void SetUp() override
  {
    state.dCache.Reset();
    state.m_enable_dcache = true;
    state.msr.DR = 1;
    state.feature_flags = FEATURE_FLAG_MSR_DR;
    state.Exceptions = 0;
    state.spr[SPR_HID0] = 0;
    bats.fill(0);
    bats[0x80000000 >> PowerPC::BAT_INDEX_SHIFT] = PowerPC::BAT_MAPPED_BIT;
    bats[0x90000000 >> PowerPC::BAT_INDEX_SHIFT] = 0x10000000 | PowerPC::BAT_MAPPED_BIT;
    code = std::make_unique<DCacheCode>(system);
  }

  void TearDown() override { code.reset(); }

  Core::System& system = Core::System::GetInstance();
  PowerPC::PowerPCState& state = system.GetPPCState();
  PowerPC::MMU& mmu = system.GetMMU();
  Memory::MemoryManager& memory = system.GetMemory();
  PowerPC::BatTable& bats = const_cast<PowerPC::BatTable&>(mmu.GetDBATTable());
  std::unique_ptr<DCacheCode> code;
};

TEST_F(Jit64DCache, StoreMissRetainsAddressInRDX)
{
  // The data is deliberately also a mapped address: a broken fallback writes there instead.
  memory.Write_U32(0, 0x1000);
  memory.Write_U32(0, 0x4000);
  code->Access(true, 0x80001000, 0x80004000, RDX, RCX);
  EXPECT_EQ(mmu.ReadForJit<u32>(0x80001000), 0x80004000);
  EXPECT_EQ(mmu.ReadForJit<u32>(0x80004000), 0U);
  EXPECT_EQ(code->result[RDX], 0x80001000);
  EXPECT_EQ(code->result[RCX], 0x80004000);
}

TEST_F(Jit64DCache, LoadMissRetainsAddressInRDX)
{
  memory.Write_U32(0x89abcdef, 0x1000);
  code->Access(false, 0x80001000, 0, RDX, R8);
  EXPECT_EQ(code->result[R8], 0x89abcdef);
  EXPECT_EQ(code->result[RDX], 0x80001000);
}

TEST_F(Jit64DCache, HitPreservesLiveScratchRegisters)
{
  mmu.WriteForJit<u32>(0xffffffff, 0x80001000);
  code->Access(false, 0x80001000, 0, R9, R8);
  EXPECT_EQ(code->result[R8], 0xffffffff);
  EXPECT_EQ(code->result[RDX], 0x12345678U);
  EXPECT_EQ(code->result[RCX], 0x87654321U);
  code->Access(true, 0x80001000, 0x11223344, R9, R8);
  EXPECT_EQ(mmu.ReadForJit<u32>(0x80001000), 0x11223344U);
  EXPECT_EQ(code->result[RDX], 0x12345678U);
  EXPECT_EQ(code->result[RCX], 0x87654321U);
}

TEST_F(Jit64DCache, PhysicalAccessIgnoresDBAT)
{
  state.msr.DR = 0;
  state.feature_flags = CPUEmuFeatureFlags{};
  bats[0] = 0x20000 | PowerPC::BAT_MAPPED_BIT;
  mmu.WriteForJit<u32>(0x11111111, 0x1000);
  mmu.WriteForJit<u32>(0x22222222, 0x21000);
  code->Access(false, 0x1000, 0, R9, R8);
  EXPECT_EQ(code->result[R8], 0x11111111U);
  code->Access(true, 0x1000, 0x33333333, R9, R8);
  EXPECT_EQ(mmu.ReadForJit<u32>(0x1000), 0x33333333U);
  EXPECT_EQ(mmu.ReadForJit<u32>(0x21000), 0x22222222U);
}

TEST_F(Jit64DCache, LoadRegisterAndOffsetCombinations)
{
  for (const bool hit : {false, true})
  {
    for (const auto address_reg : {RAX, RCX, RDX, R8, R9})
    {
      for (const auto value_reg : {RAX, RCX, RDX, R8, R9})
      {
        for (const s32 offset : {0, 4, -4})
        {
          SCOPED_TRACE(testing::Message()
                       << hit << ' ' << address_reg << ' ' << value_reg << ' ' << offset);
          state.dCache.Reset();
          memory.Write_U32(0xfedcba98, 0x1020);
          if (hit)
            ASSERT_EQ(mmu.ReadForJit<u32>(0x80001020), 0xfedcba98);
          code->Access(false, 0x80001020 - offset, 0, address_reg, value_reg, offset);
          EXPECT_EQ(code->result[value_reg], 0xfedcba98);
        }
      }
    }
  }
}

TEST_F(Jit64DCache, StoreRegisterAndOffsetCombinations)
{
  for (const bool hit : {false, true})
  {
    for (const auto address_reg : {RAX, RCX, RDX, R8, R9})
    {
      for (const auto value_reg : {RCX, RDX, R8, R9})
      {
        if (address_reg == value_reg)
          continue;
        for (const s32 offset : {0, 4, -4})
        {
          SCOPED_TRACE(testing::Message()
                       << hit << ' ' << address_reg << ' ' << value_reg << ' ' << offset);
          state.dCache.Reset();
          memory.Write_U32(0, 0x1020);
          if (hit)
            ASSERT_EQ(mmu.ReadForJit<u32>(0x80001020), 0U);
          code->Access(true, 0x80001020 - offset, 0xfedcba98, address_reg, value_reg, offset,
                       EmuCodeBlock::SAFE_LOADSTORE_CLOBBER_RSCRATCH_INSTEAD_OF_ADDR);
          EXPECT_EQ(mmu.ReadForJit<u32>(0x80001020), 0xfedcba98);
          EXPECT_EQ(code->result[value_reg], 0xfedcba98);
        }
      }
    }
  }
}

TEST_F(Jit64DCache, SplitLineAndPageAccess)
{
  for (const u32 offset : {29U, 30U, 31U, 4093U, 4094U, 4095U})
  {
    SCOPED_TRACE(offset);
    const u32 address = 0x80001000 + offset;
    mmu.WriteForJit<u32>(0x11223344, address);
    code->Access(false, address, 0, RDX, R8);
    EXPECT_EQ(code->result[R8], 0x11223344U);
    code->Access(true, address, 0x55667788, RDX, RCX);
    EXPECT_EQ(mmu.ReadForJit<u32>(address), 0x55667788U);
  }
}

TEST_F(Jit64DCache, InhibitedBATBypassesCachedData)
{
  mmu.WriteForJit<u32>(0x11111111, 0x80001000);
  memory.Write_U32(0x22222222, 0x1000);
  bats[0x80000000 >> PowerPC::BAT_INDEX_SHIFT] |= PowerPC::BAT_WI_BIT;
  code->Access(false, 0x80001000, 0, RDX, R8);
  EXPECT_EQ(code->result[R8], 0x22222222U);
  code->Access(true, 0x80001000, 0x33333333, RDX, RCX);
  EXPECT_EQ(memory.Read_U32(0x1000), 0x33333333U);
  bats[0x80000000 >> PowerPC::BAT_INDEX_SHIFT] &= ~PowerPC::BAT_WI_BIT;
  EXPECT_EQ(mmu.ReadForJit<u32>(0x80001000), 0x11111111U);
}

TEST_F(Jit64DCache, SharedRoutineStackAndDRFlag)
{
  // Shared paired-load routines are emitted without a known block's feature flags and run
  // with a return address on the stack. DR_ON is their explicit translation contract.
  state.feature_flags = CPUEmuFeatureFlags{};
  constexpr int flags = EmuCodeBlock::SAFE_LOADSTORE_NO_PROLOG |
                        EmuCodeBlock::SAFE_LOADSTORE_NO_UPDATE_PC |
                        EmuCodeBlock::SAFE_LOADSTORE_DR_ON;
  memory.Write_U32(0x89abcdef, 0x1000);
  for (int repeat = 0; repeat < 2; ++repeat)
  {
    code->Access(false, 0x80001000, 0, RDX, R8, 0, flags);
    EXPECT_EQ(code->result[R8], 0x89abcdef);
    code->Access(true, 0x80001000, 0x89abcdef, RDX, RCX, 0, flags);
    EXPECT_EQ(mmu.ReadForJit<u32>(0x80001000), 0x89abcdef);
  }
}

TEST_F(Jit64DCache, NativePLRUAndDirtyStateMatchMMU)
{
  for (const u32 base : {0x80000000U, 0x90000000U})
  {
    state.dCache.Reset();
    constexpr u32 set = 126;
    for (u32 way = 0; way < PowerPC::CACHE_WAYS; ++way)
      mmu.WriteForJit<u32>(0xffffffff, base + way * 4096 + set * 32);
    for (u32 way = 0; way < PowerPC::CACHE_WAYS; ++way)
    {
      const u32 address = base + way * 4096 + set * 32;
      for (const bool write : {false, true})
      {
        code->Access(write, address, 0xffffffff, RDX, R8);
        for (u32 old_plru = 0; old_plru < 128; ++old_plru)
        {
          SCOPED_TRACE(testing::Message() << base << ' ' << way << ' ' << write << ' ' << old_plru);
          auto& cache = state.dCache;
          cache.plru[set] = static_cast<u8>(old_plru);
          cache.modified[set] = 0x80;
          if (write)
            ASSERT_TRUE(mmu.TryWriteDCache32ForJit(0xffffffff, address));
          else
            ASSERT_EQ(mmu.TryReadDCache32ForJit(address), 0x100000000ULL);
          const auto expected_plru = cache.plru[set];
          const auto expected_dirty = cache.modified[set];
          const auto expected_data = cache.data[set];
          cache.plru[set] = static_cast<u8>(old_plru);
          cache.modified[set] = 0x80;
          code->run();
          EXPECT_EQ(cache.plru[set], expected_plru);
          EXPECT_EQ(cache.modified[set], expected_dirty);
          EXPECT_EQ(cache.data[set], expected_data);
          if (!write)
            EXPECT_EQ(code->result[R8], 0xffffffff);
        }
      }
    }
  }
}
}  // namespace