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

#include <algorithm>
#include <array>
#include <limits>
#include <memory>
#include <span>
#include <vector>

#include <gtest/gtest.h>

#include "Common/Config/Config.h"
#include "Common/Crypto/SHA1.h"
#include "Common/FileUtil.h"
#include "Core/Config/MainSettings.h"
#include "Core/ConfigManager.h"
#include "Core/Core.h"
#include "Core/HW/Memmap.h"
#include "Core/HW/WII_IPC.h"
#include "Core/IOS/Starlet/StarletMemory.h"
#include "Core/System.h"

namespace IOS::LLE
{
// Compare the block helpers with the original byte bus, using synthetic memory only.
class StarletDMATest : public testing::Test
{
protected:
  static void SetUpTestSuite()
  {
    Config::Init();
    SConfig::Init();
    auto& system = Core::System::GetInstance();
    system.SetIsWii(true);
    system.GetMemory().Init();
    Core::DeclareAsCPUThread();
  }

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

  void SetUp() override
  {
    auto& system = Core::System::GetInstance();
    system.GetWiiIPC().Reset();
    m_memory = std::make_unique<StarletMemory>(system);
  }

  void TearDown() override
  {
    m_memory.reset();
    if (!m_temp_directory.empty())
      File::DeleteDirRecursively(m_temp_directory);  // Only this fixture's unique temporary SD.
  }

  std::vector<u8> ReadBytes(u32 address, size_t size)
  {
    std::vector<u8> result(size);
    for (size_t i = 0; i < size; ++i)
      result[i] = m_memory->ReadDMA8(address + static_cast<u32>(i));
    return result;
  }

  void WriteBytes(u32 address, std::span<const u8> data)
  {
    for (size_t i = 0; i < data.size(); ++i)
      m_memory->WriteDMA8(address + static_cast<u32>(i), data[i]);
  }

  void CompareRanges()
  {
    for (const bool split : {false, true})
    {
      for (const bool boot0 : {false, true})
      {
        m_memory->m_sram_split_mode = split;
        m_memory->m_boot0_mapped = boot0;
        for (const u32 address :
             {0U,          1U,          0x10003U,    0x017ffff0U, 0x01800000U, 0x0fffffffU,
              0x10000000U, 0x10000001U, 0x13fffff0U, 0x14000000U, 0x0d3ffff0U, 0x0d400000U,
              0x0d400003U, 0x0d40fff0U, 0x0d410000U, 0x0d417ff0U, 0x0d418000U, 0x0d41fff0U,
              0xfff00000U, 0xfff0fff0U, 0xfff17ff0U, 0xfffefff0U, 0xffff0000U, 0xffff7ff0U,
              0xfffffff0U, 0x0d800020U})
        {
          for (const size_t size : {0U, 1U, 2U, 3U, 4U, 7U, 16U, 63U, 512U, 2048U, 8192U})
          {
            SCOPED_TRACE(testing::Message() << "address=" << address << " size=" << size
                                            << " split=" << split << " boot0=" << boot0);
            std::vector<u8> initial(size + 32);
            std::vector<u8> payload(size);
            for (size_t i = 0; i < initial.size(); ++i)
              initial[i] = static_cast<u8>((i * 37) ^ (i >> 8));
            for (size_t i = 0; i < size; ++i)
              payload[i] = static_cast<u8>((i * 13) ^ (i >> 7) ^ 0xab);
            WriteBytes(address - 16, initial);
            std::vector<u8> actual(size);
            m_memory->ReadDMABuffer(address, actual);
            ASSERT_EQ(actual, ReadBytes(address, size));

            WriteBytes(address, payload);
            const auto expected = ReadBytes(address - 16, size + 32);
            const auto expected_sram = m_memory->m_sram;
            WriteBytes(address - 16, initial);
            m_memory->WriteDMABuffer(address, payload);
            ASSERT_EQ(ReadBytes(address - 16, size + 32), expected);
            ASSERT_EQ(m_memory->m_sram, expected_sram);
          }
        }
      }
    }
  }

  void CheckPointerGuards()
  {
    for (const u32 address : {0U, 0x10000000U, 0x0d400000U, 0x0d410000U})
    {
      EXPECT_NE(m_memory->GetDMAMemoryPointer(address, 16), nullptr);
      EXPECT_EQ(m_memory->GetDMAMemoryPointer(address, std::numeric_limits<size_t>::max()),
                nullptr);
    }
    for (const u32 address : {0x017ffff0U, 0x13fffff0U, 0x0d417ff0U})
    {
      EXPECT_NE(m_memory->GetDMAMemoryPointer(address, 16), nullptr);
      EXPECT_EQ(m_memory->GetDMAMemoryPointer(address, 17), nullptr);
    }
    for (const u32 address : {0x0d418000U, 0x0d800038U, 0xfff00000U, 0xfffffff0U})
      EXPECT_EQ(m_memory->GetDMAMemoryPointer(address, 4), nullptr);
  }

  void CheckBootstrapStoredKeys()
  {
    m_memory->PrepareBootstrap();
    for (const u8 read_all : {0, 1})
    {
      SCOPED_TRACE(read_all);
      m_memory->m_ohci1_hci_events.clear();
      m_memory->ExecuteOHCI1HCICommand({0x0d, 0x0c, 7, 0, 0, 0, 0, 0, 0, read_all});
      ASSERT_EQ(m_memory->m_ohci1_hci_events.size(), 1u);
      // Command Complete: capacity 255, zero stored keys, no Return Link Keys event.
      EXPECT_EQ(m_memory->m_ohci1_hci_events.front(),
                (std::vector<u8>{0x0e, 8, 1, 0x0d, 0x0c, 0, 0xff, 0, 0, 0}));
    }
  }

  void CheckOHCI()
  {
    for (const u32 base : {0x10000U, 0x10010000U, 0x0d404000U})
    {
      // The second physical page can precede the first and need not be adjacent.
      const u32 first = base + 0x3fe3;
      const u32 end = base + 0x0017;
      std::vector<u8> payload(70);
      for (size_t i = 0; i < payload.size(); ++i)
        payload[i] = static_cast<u8>(i + 1);
      WriteBytes(first - 1, std::vector<u8>(100, 0xcc));
      WriteBytes(base - 1, std::vector<u8>(100, 0xdd));
      m_memory->WriteOHCIBuffer(first, end, payload.data(), payload.size());
      payload.resize(29 + 24);
      EXPECT_EQ(m_memory->ReadOHCIBuffer(first, end), payload);
      EXPECT_EQ(ReadBytes(first - 1, 1)[0], 0xcc);
      EXPECT_EQ(ReadBytes(first + 29, 1)[0], 0xcc);
      EXPECT_EQ(ReadBytes(base - 1, 1)[0], 0xdd);
      EXPECT_EQ(ReadBytes(end + 1, 1)[0], 0xdd);
      // Short packet must not fill the remainder of either page.
      m_memory->WriteOHCIBuffer(first, end, payload.data(), 3);
      EXPECT_EQ(m_memory->ReadOHCIBuffer(first, first + 2), (std::vector<u8>{1, 2, 3}));
      m_memory->WriteOHCIMemory32(base + 1, 0x12345678);
      EXPECT_EQ(ReadBytes(base + 1, 4), (std::vector<u8>{0x78, 0x56, 0x34, 0x12}));
      EXPECT_EQ(m_memory->ReadOHCIMemory32(base + 1), 0x12345678U);
    }
    EXPECT_TRUE(m_memory->ReadOHCIBuffer(0, 0x1234).empty());
    EXPECT_TRUE(m_memory->ReadOHCIBuffer(0x1234, 0x1233).empty());
  }

  void CheckNAND()
  {
    StarletMemory::NANDPage raw;
    for (size_t i = 0; i < raw.size(); ++i)
      raw[i] = static_cast<u8>((i * 37) ^ (i >> 8));
    m_memory->m_nand_overlay[7] = raw;
    m_memory->m_nand_read_pending = true;
    m_memory->m_nand_read_page = 7;
    for (const u32 base : {0x10000U, 0x10010000U, 0x0d404000U})
    {
      const u32 spare = base + 0x2000;
      m_memory->WriteRegister(0x0d010010, base);
      m_memory->WriteRegister(0x0d010014, spare);
      m_memory->WriteRegister(0x0d01000c, 7);
      for (const u32 column : {0U, 1U, 0x7f0U, 0x800U, 0x830U})
      {
        for (const u32 length : {1U, 0x840U - column})
        {
          SCOPED_TRACE(testing::Message()
                       << "base=" << base << " column=" << column << " length=" << length);
          WriteBytes(base, std::vector<u8>(0x900, 0xcc));
          WriteBytes(spare, std::vector<u8>(0x100, 0xdd));
          m_memory->m_nand_read_column = column;
          ASSERT_TRUE(m_memory->ReadNANDPage(0x2000 | length));
          // Independent old byte-address formula, including spare-only reads via NAND_DATA.
          for (u32 i = 0; i < length; ++i)
          {
            const u32 address =
                column < 0x800 && column + i >= 0x800 ? spare + column + i - 0x800 : base + i;
            ASSERT_EQ(m_memory->ReadDMA8(address), raw[column + i]);
          }
          m_memory->WriteRegister(0x0d010008, column);
          ASSERT_TRUE(m_memory->StageNANDProgram(0x4000 | length, false));
          auto expected = raw;
          std::fill(expected.begin(), expected.begin() + column, 0xff);
          std::fill(expected.begin() + column + length, expected.end(), 0xff);
          EXPECT_EQ(m_memory->m_nand_program_data, expected);
        }
      }
      // ECC output remains separate from the transferred spare bytes.
      m_memory->m_nand_read_column = 0;
      ASSERT_TRUE(m_memory->ReadNANDPage(0x3000 | 0x840));
      const auto ecc = StarletMemory::CalculateNANDECC(raw.data());
      EXPECT_EQ(ReadBytes(spare ^ 0x40, ecc.size()), (std::vector<u8>(ecc.begin(), ecc.end())));
      EXPECT_EQ(ReadBytes(spare, 0x40), (std::vector<u8>(raw.begin() + 0x800, raw.end())));
      m_memory->WriteRegister(0x0d010008, 0);
      ASSERT_TRUE(m_memory->StageNANDProgram(0x5800, false));
      EXPECT_TRUE(m_memory->m_nand_program_ecc_enabled);
      EXPECT_EQ(ReadBytes(spare ^ 0x40, ecc.size()), (std::vector<u8>(ecc.begin(), ecc.end())));
      // Random spare input uses NAND_DATA and preserves the already staged page data.
      WriteBytes(base, std::span(raw).subspan(0x800));
      m_memory->WriteRegister(0x0d010008, 0x800);
      ASSERT_TRUE(m_memory->StageNANDProgram(0x4040, true));
      EXPECT_EQ(m_memory->m_nand_program_data, raw);
    }
  }

  void CheckOverlappingAES()
  {
    for (const u32 base : {0x10000U, 0x10010000U, 0x0d404000U})
    {
      for (const u32 offset : {0U, 16U})
      {
        std::vector<u8> input(2048);
        for (size_t i = 0; i < input.size(); ++i)
          input[i] = static_cast<u8>((i * 37) ^ (i >> 8));
        WriteBytes(base, input);
        m_memory->Write32(0x0d020004, base);
        m_memory->Write32(0x0d020008, base + offset);
        m_memory->Write32(0x0d020000, 0x8000007f);
        EXPECT_EQ(ReadBytes(base + offset, input.size()), input);
        EXPECT_EQ(m_memory->Read32(0x0d020004), base + input.size());
        EXPECT_EQ(m_memory->Read32(0x0d020008), base + offset + input.size());
      }
    }
  }

  void CheckSHA()
  {
    std::vector<u8> input(257);
    for (size_t i = 0; i < input.size(); ++i)
      input[i] = static_cast<u8>(i * 37);
    const auto digest = Common::SHA1::CalculateDigest(input);
    input.resize(320, 0);
    input[257] = 0x80;
    input[318] = 0x08;
    input[319] = 0x08;  // 257 * 8 bits, big endian.
    constexpr std::array<u32, 5> initial = {0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476,
                                            0xc3d2e1f0};
    for (const u32 base : {0x10000U, 0x10010000U, 0x0d404000U})
    {
      WriteBytes(base, input);
      for (u32 i = 0; i < initial.size(); ++i)
        m_memory->Write32(0x0d030008 + i * 4, initial[i]);
      m_memory->Write32(0x0d030004, base);
      m_memory->Write32(0x0d030000, 0x80000004);  // Five blocks.
      EXPECT_EQ(ReadBytes(0x0d030008, digest.size()),
                (std::vector<u8>(digest.begin(), digest.end())));
      EXPECT_EQ(m_memory->Read32(0x0d030004), base + input.size());
    }
  }

  void CheckSD()
  {
    m_temp_directory = File::CreateTempDir();
    ASSERT_FALSE(m_temp_directory.empty());
    ASSERT_TRUE(m_memory->m_sd_card.Open(m_temp_directory + "/dma-test.raw", "w+b"));
    std::vector<u8> image(4096);
    for (size_t i = 0; i < image.size(); ++i)
      image[i] = static_cast<u8>((i * 37) ^ (i >> 8));
    ASSERT_TRUE(m_memory->m_sd_card.WriteBytes(image.data(), image.size()));
    ASSERT_TRUE(m_memory->m_sd_card.Flush());
    m_memory->m_sd_card_inserted = true;
    for (const bool high_capacity : {false, true})
    {
      m_memory->m_sd_card_high_capacity = high_capacity;
      const u32 argument = high_capacity ? 1 : 512;
      for (const u32 base : {0x10000U, 0x10010000U, 0x0d404000U})
      {
        ASSERT_TRUE(m_memory->TransferSDCardBlocks(true, argument, 512, 3, base));
        EXPECT_EQ(ReadBytes(base, 1536),
                  (std::vector<u8>(image.begin() + 512, image.begin() + 2048)));
        EXPECT_EQ(m_memory->ReadRegister(0x0d070000), base + 1536);
        Config::SetCurrent(Config::MAIN_ALLOW_SD_WRITES, false);
        EXPECT_FALSE(m_memory->TransferSDCardBlocks(false, argument, 512, 3, base));
        Config::SetCurrent(Config::MAIN_ALLOW_SD_WRITES, true);
        auto changed = ReadBytes(base, 1536);
        for (auto& byte : changed)
          byte ^= 0x5a;
        WriteBytes(base, changed);
        ASSERT_TRUE(m_memory->TransferSDCardBlocks(false, argument, 512, 3, base));
        std::copy(changed.begin(), changed.end(), image.begin() + 512);
        ASSERT_TRUE(m_memory->m_sd_card.Seek(0, File::SeekOrigin::Begin));
        std::vector<u8> actual(image.size());
        ASSERT_TRUE(m_memory->m_sd_card.ReadBytes(actual.data(), actual.size()));
        EXPECT_EQ(actual, image);
        EXPECT_FALSE(m_memory->TransferSDCardBlocks(true, argument, 4096, 1, base));
      }
    }
  }

  void CheckCPUSRAMReads()
  {
    for (size_t i = 0; i < m_memory->m_boot_rom.size(); ++i)
      m_memory->m_boot_rom[i] = static_cast<u8>((i * 29) ^ (i >> 5) ^ 0xa5);
    for (size_t i = 0; i < m_memory->m_sram.size(); ++i)
      m_memory->m_sram[i] = static_cast<u8>((i * 13) ^ (i >> 8) ^ 0x5a);
    // The byte CPU path is unchanged by the aligned wide-read optimization. Use it as the
    // reference, including nonzero boot ROM, all mirrors, live mapping changes and wrap to MEM1.
    for (const bool split : {false, true, false})
    {
      for (const bool boot0 : {false, true, false})
      {
        m_memory->Write32(0x0d800060, split ? 0x20 : 0);
        m_memory->Write32(0x0d80018c, boot0 ? 0 : 0x1000);
        for (const u32 base : {0x0d400000U, 0xfff00000U, 0xfff20000U, 0xfff40000U, 0xfff60000U,
                               0xfff80000U, 0xfffa0000U, 0xfffc0000U, 0xfffe0000U})
        {
          for (u32 page = 0; page < 32; ++page)
          {
            for (const u32 offset : {0U, 1U, 2U, 3U, 0xffcU, 0xffdU, 0xffeU, 0xfffU})
            {
              const u32 address = base + page * 0x1000 + offset;
              const u16 expected16 =
                  static_cast<u16>((m_memory->Read8(address) << 8) | m_memory->Read8(address + 1));
              const u32 expected32 = (static_cast<u32>(expected16) << 16) |
                                     (m_memory->Read8(address + 2) << 8) |
                                     m_memory->Read8(address + 3);
              ASSERT_EQ(m_memory->Read16(address), expected16)
                  << "address=" << address << " split=" << split << " boot0=" << boot0;
              ASSERT_EQ(m_memory->Read32(address), expected32)
                  << "address=" << address << " split=" << split << " boot0=" << boot0;
            }
          }
        }
      }
    }
  }

  void CheckVersionReads()
  {
    constexpr u32 version = 0x0d800214;
    m_memory->Reset();
    ASSERT_EQ(m_memory->Read32(version), 0x11U);
    for (const u32 value : {0x12345678U, 0xdeadbeefU, 0U, 0xffffffffU})
    {
      m_memory->Write32(version, value);
      for (u32 i = 0; i < 4; ++i)
        EXPECT_EQ(m_memory->Read8(version + i), static_cast<u8>(value >> (24 - i * 8)));
      EXPECT_EQ(m_memory->Read32(version), value);
      m_memory->Write8(version + 1, 0x5a);
      EXPECT_EQ(m_memory->Read32(version), (value & 0xff00ffffU) | 0x005a0000U);
    }
    m_memory->Reset();
    EXPECT_EQ(m_memory->Read32(version), 0x11U);
  }

  std::unique_ptr<StarletMemory> m_memory;
  std::string m_temp_directory;
};

TEST_F(StarletDMATest, BuffersMatchByteBusAcrossMappingsBoundariesAndWrap)
{
  CompareRanges();
}

TEST_F(StarletDMATest, DirectRangesExcludeHolesRegistersMirrorsAndOverflow)
{
  CheckPointerGuards();
}

TEST_F(StarletDMATest, OHCIUsesTwoPhysicalPagesAndLittleEndianDescriptors)
{
  CheckOHCI();
}

TEST_F(StarletDMATest, BootstrapReadStoredKeysHasNoWiimoteDevices)
{
  CheckBootstrapStoredKeys();
}

TEST_F(StarletDMATest, NANDSeparatesDataSpareAndECCIncludingPartialTransfers)
{
  CheckNAND();
}

TEST_F(StarletDMATest, AESStagesAllInputBeforeWritingOverlappingOutput)
{
  CheckOverlappingAES();
}

TEST_F(StarletDMATest, SHAMultipleBlocksMatchDigestInEveryPhysicalBank)
{
  CheckSHA();
}

TEST_F(StarletDMATest, SDTransfersPreserveOffsetsContentsAndWriteProtection)
{
  CheckSD();
}

TEST_F(StarletDMATest, CPUWideSRAMReadsMatchBytesWithBootROMAndEveryMirror)
{
  CheckCPUSRAMReads();
}

TEST_F(StarletDMATest, VersionWordReadsPreserveLiveByteBackingAndReset)
{
  CheckVersionReads();
}
}  // namespace IOS::LLE