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

#include <algorithm>
#include <array>
#include <iomanip>
#include <memory>
#include <optional>

#include <gtest/gtest.h>

#include "Common/Swap.h"
#include "Core/Core.h"
#include "Core/HW/WII_IPC.h"
#include "Core/IOS/Starlet/ARMCore.h"
#include "Core/IOS/Starlet/StarletMemory.h"
#include "Core/System.h"

namespace IOS::LLE
{
#if defined(_M_X86_64)
namespace
{
// Only code and page tables use synthetic RAM. Data reads exercise the actual SRAM/boot0 bus,
// without initializing any firmware, NAND image, profile, or live peripheral scheduler.
class SRAMReadBus final : public ARMBus
{
public:
  SRAMReadBus() : memory(Core::System::GetInstance())
  {
    memory.Reset();
    for (u32 i = 0; i < 0x18000; ++i)
      memory.GetFastmemSRAMBase()[i] = static_cast<u8>((i * 13) ^ (i >> 8) ^ 0xa5);
  }

  u8 Read8(u32 address) override
  {
    if (address < ram.size())
      return ram[address];
    ++data_reads;
    return memory.Read8(address);
  }
  u16 Read16(u32 address) override
  {
    if (address < ram.size())
      return ARMBus::Read16(address);
    ++data_reads;
    return memory.Read16(address);
  }
  u32 Read32(u32 address) override
  {
    if (address < ram.size())
      return ARMBus::Read32(address);
    ++data_reads;
    return memory.Read32(address);
  }
  void Write8(u32 address, u8 value) override
  {
    if (address < ram.size())
      ram[address] = value;
    else
      memory.Write8(address, value);
  }
  void Write16(u32 address, u16 value) override
  {
    if (address < ram.size())
      ARMBus::Write16(address, value);
    else
      memory.Write16(address, value);
  }
  void Write32(u32 address, u32 value) override
  {
    if (address < ram.size())
      ARMBus::Write32(address, value);
    else
      memory.Write32(address, value);
  }
  u8* GetFastmemBase() const override { return const_cast<u8*>(ram.data()); }
  u8* GetFastmemSRAMBase() const override { return memory.GetFastmemSRAMBase(); }
  const bool* GetFastmemBoot0Mapped() const override { return memory.GetFastmemBoot0Mapped(); }
  const bool* GetFastmemSRAMSplitMode() const override { return memory.GetFastmemSRAMSplitMode(); }
  const u8* GetDirectMemoryPointer(u32 address, u32 size) const override
  {
    if (address < ram.size() && size <= ram.size() - address)
      return ram.data() + address;
    return memory.GetDirectMemoryPointer(address, size);
  }

  void SetMapping(bool split, bool boot0)
  {
    memory.Write32(0x0d800060, split ? 0x20 : 0);    // HW_SRNPROT
    memory.Write32(0x0d80018c, boot0 ? 0 : 0x1000);  // HW_BOOT0
  }

  void InstallReads(bool thumb, bool big_endian)
  {
    if (thumb)
    {
      constexpr std::array<u16, 6> code = {
          0x6830,  // ldr r0, [r6]
          0x7831,  // ldrb r1, [r6]
          0x8832,  // ldrh r2, [r6]
          0x57f3,  // ldrsb r3, [r6, r7]
          0x5ff4,  // ldrsh r4, [r6, r7]
          0xe7fe,  // b .
      };
      for (u32 i = 0; i < code.size(); ++i)
        Write16(i * 2, big_endian ? code[i] : Common::swap16(code[i]));
    }
    else
    {
      constexpr std::array<u32, 6> code = {
          0xe5960000,  // ldr r0, [r6]
          0xe5d61000,  // ldrb r1, [r6]
          0xe1d620b0,  // ldrh r2, [r6]
          0xe1d630d0,  // ldrsb r3, [r6]
          0xe1d640f0,  // ldrsh r4, [r6]
          0xeafffffe,  // b .
      };
      for (u32 i = 0; i < code.size(); ++i)
        Write32(i * 4, big_endian ? code[i] : Common::swap32(code[i]));
    }
  }

  void InstallStores(bool thumb, bool big_endian)
  {
    if (thumb)
    {
      constexpr std::array<u16, 4> code = {
          0x6030,  // str r0, [r6]
          0x7131,  // strb r1, [r6, #4]
          0x80f2,  // strh r2, [r6, #6]
          0xe7fe,  // b .
      };
      for (u32 i = 0; i < code.size(); ++i)
        Write16(i * 2, big_endian ? code[i] : Common::swap16(code[i]));
    }
    else
    {
      constexpr std::array<u32, 4> code = {
          0xe5860000,  // str r0, [r6]
          0xe5c61004,  // strb r1, [r6, #4]
          0xe1c620b6,  // strh r2, [r6, #6]
          0xeafffffe,  // b .
      };
      for (u32 i = 0; i < code.size(); ++i)
        Write32(i * 4, big_endian ? code[i] : Common::swap32(code[i]));
    }
  }

  StarletMemory memory;
  std::array<u8, 0x8000> ram{};
  u64 data_reads = 0;
};

class StarletSRAMRead : public testing::TestWithParam<bool>
{
protected:
  void SetUp() override
  {
    Core::DeclareAsCPUThread();
    Core::System::GetInstance().GetWiiIPC().Reset();
  }

  static void Prepare(ARMCore& core, bool thumb, bool big_endian, u32 address)
  {
    if (core.IsBigEndian() != big_endian)
      core.SetBigEndian(big_endian);
    core.SetCPSR(static_cast<u32>(ARMCore::Mode::Supervisor) | (thumb ? ARMCore::CPSR_T : 0));
    for (u32 i = 0; i < 16; ++i)
      core.SetRegister(i, 0);
    core.SetRegister(6, address);
  }
};

TEST_P(StarletSRAMRead, NewlyProfiledPagesAvoidDataBusCalls)
{
  for (const bool thumb : {false, true})
  {
    for (const bool big_endian : {false, true})
    {
      auto bus = std::make_unique<SRAMReadBus>();
      bus->SetMapping(true, false);
      bus->InstallReads(thumb, big_endian);
      ARMCore core(*bus, GetParam());
      core.SetJitEnabled(true);
      for (const u32 base : {0x0d400000u, 0xfff00000u, 0xfffe0000u})
      {
        for (const u32 page : {0x01u, 0x02u, 0x04u, 0x13u, 0x15u, 0x16u, 0x1au, 0x1cu, 0x1du})
        {
          const u32 address = base + page * 0x1000;
          SCOPED_TRACE(testing::Message() << "thumb=" << thumb << " be=" << big_endian
                                          << " address=" << std::hex << address);
          Prepare(core, thumb, big_endian, address);
          bus->data_reads = 0;
          ASSERT_EQ(core.RunCycles(6), 6u);
          ASSERT_EQ(core.GetLastUndefinedInstruction(), 0u);
          ASSERT_EQ(bus->data_reads, 0u);
          EXPECT_EQ(core.GetRegister(0), big_endian ? bus->memory.Read32(address) :
                                                      Common::swap32(bus->memory.Read32(address)));
        }
      }
      EXPECT_EQ(core.GetJitSlowReadCount(), 0u);
      EXPECT_EQ(core.GetJitFallbackInstructionCount(), 0u);
      EXPECT_EQ(core.GetJitCompiledBlockCount(), 1u);
    }
  }
}

TEST_P(StarletSRAMRead, SizedNativeStoresMatchInterpreterAcrossMappings)
{
  for (const bool thumb : {false, true})
  {
    for (const bool big_endian : {false, true})
    {
      auto reference_bus = std::make_unique<SRAMReadBus>();
      auto native_bus = std::make_unique<SRAMReadBus>();
      reference_bus->InstallStores(thumb, big_endian);
      native_bus->InstallStores(thumb, big_endian);
      ARMCore reference(*reference_bus);
      ARMCore native(*native_bus, GetParam());
      native.SetJitEnabled(true);
      for (const bool split : {false, true, false})
      {
        for (const bool boot0 : {false, true, false})
        {
          reference_bus->SetMapping(split, boot0);
          native_bus->SetMapping(split, boot0);
          for (const u32 base : {0x0d400000U, 0xfff00000U, 0xfffe0000U})
          {
            for (u32 page = 0; page < 32; ++page)
            {
              for (const u32 offset : {0U, 1U, 2U, 3U, 0xff8U, 0xffcU, 0xffeU, 0xfffU})
              {
                const u32 address = base + page * 0x1000 + offset;
                if (address > 0xfffffff7U)
                  continue;  // Do not wrap a test data store into synthetic guest code.
                SCOPED_TRACE(testing::Message()
                             << "thumb=" << thumb << " be=" << big_endian << " split=" << split
                             << " boot0=" << boot0 << " address=" << std::hex << address);
                for (ARMCore* core : {&reference, &native})
                {
                  Prepare(*core, thumb, big_endian, address);
                  core->SetRegister(0, 0x12345678);
                  core->SetRegister(1, 0x9abcdef0);
                  core->SetRegister(2, 0x7e3dac12);
                  ASSERT_EQ(core->RunCycles(4), 4U);
                  ASSERT_EQ(core->GetLastUndefinedInstruction(), 0U);
                }
                for (u32 reg = 0; reg < 16; ++reg)
                  ASSERT_EQ(native.GetRegister(reg), reference.GetRegister(reg));
                ASSERT_EQ(native.GetCPSR(), reference.GetCPSR());
                ASSERT_TRUE(std::equal(reference_bus->memory.GetFastmemSRAMBase(),
                                       reference_bus->memory.GetFastmemSRAMBase() + 0x18000,
                                       native_bus->memory.GetFastmemSRAMBase()));
                for (u32 byte = 0; byte < 8; ++byte)
                  ASSERT_EQ(native_bus->Read8(address + byte),
                            reference_bus->Read8(address + byte));
              }
            }
          }
        }
      }
    }
  }
}

TEST_P(StarletSRAMRead, MatchesInterpreterAcrossLiveMappingsAliasesAndBoundaries)
{
  for (const bool thumb : {false, true})
  {
    for (const bool big_endian : {false, true})
    {
      auto interpreter_bus = std::make_unique<SRAMReadBus>();
      auto jit_bus = std::make_unique<SRAMReadBus>();
      interpreter_bus->InstallReads(thumb, big_endian);
      jit_bus->InstallReads(thumb, big_endian);
      ARMCore interpreter(*interpreter_bus);
      ARMCore jit(*jit_bus, GetParam());
      jit.SetJitEnabled(true);
      std::array<u8, 0x18000> original;
      std::copy_n(jit_bus->GetFastmemSRAMBase(), original.size(), original.begin());
      std::optional<u64> compiled_blocks;
      // Toggle the hardware state while reusing the same compiled code: guards must consult
      // live boot0/split latches, not their values at compilation time.
      for (const bool split : {false, true, false})
      {
        for (const bool boot0 : {false, true, false})
        {
          interpreter_bus->SetMapping(split, boot0);
          jit_bus->SetMapping(split, boot0);
          for (const u32 base : {0x0d400000u, 0xfff00000u, 0xfffe0000u})
          {
            for (u32 page = 0; page < 32; ++page)
            {
              for (const u32 offset :
                   {0u, 1u, 2u, 3u, 0x3fdu, 0x3feu, 0x3ffu, 0xffcu, 0xffdu, 0xffeu, 0xfffu})
              {
                const u32 address = base + page * 0x1000 + offset;
                if (address > 0xfffffffcu)
                  continue;  // Exclude address-space wrap into the synthetic code buffer.
                SCOPED_TRACE(testing::Message()
                             << "thumb=" << thumb << " be=" << big_endian << " split=" << split
                             << " boot0=" << boot0 << " address=" << std::hex << address);
                Prepare(interpreter, thumb, big_endian, address);
                Prepare(jit, thumb, big_endian, address);
                ASSERT_EQ(interpreter.RunCycles(6), 6u);
                ASSERT_EQ(jit.RunCycles(6), 6u);
                ASSERT_EQ(jit.GetLastUndefinedInstruction(), 0u);
                for (u32 reg = 0; reg < 16; ++reg)
                  ASSERT_EQ(jit.GetRegister(reg), interpreter.GetRegister(reg)) << "r" << reg;
                ASSERT_EQ(jit.GetCPSR(), interpreter.GetCPSR());
              }
            }
          }
          // Unaligned fallbacks also compile continuations. Once the first boundary sweep has
          // warmed them, changing the hardware mapping must reuse those blocks, not rebuild them.
          if (!compiled_blocks)
            compiled_blocks = jit.GetJitLifetimeCompiledBlockCount();
          else
            EXPECT_EQ(jit.GetJitLifetimeCompiledBlockCount(), *compiled_blocks);
        }
      }
      EXPECT_TRUE(std::equal(original.begin(), original.end(), jit_bus->GetFastmemSRAMBase()));
      EXPECT_TRUE(
          std::equal(original.begin(), original.end(), interpreter_bus->GetFastmemSRAMBase()));
    }
  }
}

TEST_P(StarletSRAMRead, PreservedBlockUsesRemappedSRAMAfterGuestTLBInvalidation)
{
  auto bus = std::make_unique<SRAMReadBus>();
  bus->SetMapping(true, false);
  bus->Write32(0, 0xe5960000);       // ldr r0, [r6]
  bus->Write32(4, 0xeafffffe);       // b .
  bus->Write32(0x40, 0xee080f17);    // mcr p15, 0, r0, c8, c7, 0 (invalidate unified TLB)
  bus->Write32(0x4000, 0x00000c02);  // Identity section for code/page tables.
  bus->Write32(0x4c00, 0x00005001);  // VA 0x30000000 -> coarse page table at 0x5000.
  bus->Write32(0x5000, 0x0d404ff2);  // Small page -> SRAM B, aperture page 04.
  bus->memory.Write32(0x0d404000, 0x12345678);
  bus->memory.Write32(0xffff6000, 0x87654321);
  ARMCore core(*bus, GetParam());
  core.GetCP15State().translation_table_base = 0x4000;
  core.GetCP15State().domain_access_control = 3;
  core.GetCP15State().control |= 1;
  core.SetRegister(6, 0x30000000);
  core.SetJitEnabled(true);
  ASSERT_EQ(core.RunCycles(2), 2u);
  ASSERT_EQ(core.GetRegister(0), 0x12345678u);
  const auto old_blocks = core.GetJitCompiledBlockCount();
  bus->Write32(0x5000, 0xffff6ff2);  // Remap to SRAM A, aperture page 16.
  core.SetRegister(15, 0x40);
  ASSERT_EQ(core.RunCycles(1), 1u);
  core.SetRegister(15, 0);
  ASSERT_EQ(core.RunCycles(2), 2u);
  EXPECT_EQ(core.GetRegister(0), 0x87654321u);
  EXPECT_EQ(core.GetJitCompiledBlockCount(), old_blocks + 1);
  EXPECT_EQ(core.GetJitSlowReadCount(), 0u);
}
TEST_P(StarletSRAMRead, OptionalProfilingPreservesBusCountsResultsAndReset)
{
  for (const bool thumb : {false, true})
  {
    for (const bool big_endian : {false, true})
    {
      for (const u32 address : {0xfff1f000U, 0x0d800214U})
      {
        auto bus = std::make_unique<SRAMReadBus>();
        auto reference_bus = std::make_unique<SRAMReadBus>();
        for (SRAMReadBus* target : {bus.get(), reference_bus.get()})
        {
          target->SetMapping(true, false);
          target->memory.Write32(address, 0x89abcdef);
          if (thumb)
          {
            target->Write16(0, big_endian ? 0x6830 : Common::swap16(0x6830));  // ldr r0,[r6]
            target->Write16(2, big_endian ? 0x6030 : Common::swap16(0x6030));  // str r0,[r6]
            target->Write16(4, big_endian ? 0xe7fc : Common::swap16(0xe7fc));  // b 0
          }
          else
          {
            target->Write32(0, big_endian ? 0xe5960000 : Common::swap32(0xe5960000));
            target->Write32(4, big_endian ? 0xe5860000 : Common::swap32(0xe5860000));
            target->Write32(8, big_endian ? 0xeafffffc : Common::swap32(0xeafffffc));
          }
        }
        ARMCore core(*bus, GetParam());
        ARMCore reference(*reference_bus);
        EXPECT_FALSE(reference.IsJitMemoryProfilingEnabled());
        core.SetJitEnabled(true);
        // Recompile after reset and endian changes with the same host-only profiling choice.
        for (int reset = 0; reset < 2; ++reset)
        {
          core.Reset();
          reference.Reset();
          EXPECT_EQ(core.IsJitMemoryProfilingEnabled(), GetParam());
          const u64 reads_before = core.GetJitSlowReadCount();
          const u64 writes_before = core.GetJitSlowWriteCount();
          for (ARMCore* cpu : {&core, &reference})
          {
            Prepare(*cpu, thumb, big_endian, address);
            ASSERT_EQ(cpu->RunCycles(768), 768U);
          }
          EXPECT_EQ(core.GetJitSlowReadCount() - reads_before, 256U);
          EXPECT_EQ(core.GetJitSlowWriteCount() - writes_before, 256U);
          for (u32 reg = 0; reg < 16; ++reg)
            EXPECT_EQ(core.GetRegister(reg), reference.GetRegister(reg));
          EXPECT_EQ(core.GetCPSR(), reference.GetCPSR());
          EXPECT_EQ(bus->memory.Read32(address), reference_bus->memory.Read32(address));
          EXPECT_EQ(bus->memory.Read32(address), 0x89abcdefU);
          const auto samples = core.TakeHotJitSlowMemorySamples(8);
          if (GetParam())
          {
            ASSERT_EQ(samples.size(), 2U);
            for (const auto& [key, count] : samples)
            {
              EXPECT_EQ(static_cast<u32>(key), address);
              EXPECT_EQ(count, 1U);
            }
            EXPECT_GT(address == 0xfff1f000 ? core.GetJitSlowSRAMAccessCount() :
                                              core.GetJitSlowMMIOAccessCount(),
                      0U);
          }
          else
          {
            EXPECT_TRUE(samples.empty());
            EXPECT_EQ(core.GetJitSlowSRAMAccessCount(), 0U);
            EXPECT_EQ(core.GetJitSlowMMIOAccessCount(), 0U);
            EXPECT_EQ(core.GetJitSlowSRAMPageAccessCount(0x1f), 0U);
          }
        }
      }
    }
  }
}

INSTANTIATE_TEST_SUITE_P(MemoryProfiling, StarletSRAMRead, testing::Bool());
}  // namespace
#endif
}  // namespace IOS::LLE