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

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

#include <gtest/gtest.h>

#include "Common/Crypto/AES.h"
#include "Common/Crypto/HMAC.h"
#include "Core/Core.h"
#include "Core/HW/WII_IPC.h"
#include "Core/IOS/Starlet/StarletMemory.h"
#include "Core/IOS/USB/Bluetooth/BTBase.h"
#include "Core/System.h"
#include "DiscIO/NANDImporter.h"

namespace IOS::LLE
{
class StarletPairingTest : public testing::Test
{
protected:
  using Superblock = DiscIO::NANDImporter::NANDSuperblock;
  static constexpr u32 SUPERBLOCK_START = 0x1fc00000;
  static constexpr u16 SYSCONF_CLUSTER = 0x20;

  void SetUp() override
  {
    Core::DeclareAsCPUThread();
    auto& system = Core::System::GetInstance();
    system.GetWiiIPC().Reset();
    m_memory = std::make_unique<StarletMemory>(system);
    m_memory->Reset();
    m_superblock = std::make_unique<Superblock>();
    m_superblock->magic = {'S', 'F', 'F', 'S'};
    m_superblock->version = 1;
    for (auto& entry : m_superblock->fst)
    {
      entry.sub = 0xffff;
      entry.sib = 0xffff;
    }
    m_superblock->fst[0].mode = 2;
    m_superblock->fat[SYSCONF_CLUSTER] = 0xfffb;
    // All sixteen raw superblock copies can be read, but only copy zero is valid.
    const std::vector<u8> empty(16 * sizeof(Superblock));
    PutLogicalData(SUPERBLOCK_START, empty);
  }

  void PutLogicalData(u32 offset, std::span<const u8> bytes)
  {
    ASSERT_EQ(offset % 0x800, 0u);
    ASSERT_EQ(bytes.size() % 0x800, 0u);
    for (u32 i = 0; i < bytes.size(); i += 0x800)
    {
      NANDJournal::Page page{};
      std::copy_n(bytes.begin() + i, 0x800, page.begin());
      m_memory->m_nand_overlay.insert_or_assign((offset + i) / 0x800, page);
    }
  }

  void AddPath(u32 depth)
  {
    constexpr std::array<std::string_view, 3> names = {"shared2", "sys", "SYSCONF"};
    for (u32 i = 0; i < depth; ++i)
    {
      m_superblock->fst[i].sub = static_cast<u16>(i + 1);
      auto& child = m_superblock->fst[i + 1];
      std::ranges::copy(names[i], child.name);
      child.mode = i == 2 ? 1 : 2;
    }
  }

  void SetSysconf(bool with_pairings)
  {
    AddPath(3);
    m_superblock->fst[3].sub = SYSCONF_CLUSTER;
    m_superblock->fst[3].size = 0x4000;
    std::vector<u8> plain(0x4000);
    std::ranges::copy(std::array{'S', 'C', 'v', '0'}, plain.begin());
    if (with_pairings)
    {
      // Valid but empty BT.DINF: this takes the existing auto-pairing write path.
      const auto pairing = IOS::HLE::GetEmulatedWiimotePairingData();
      plain[5] = 1;
      plain[7] = 0x10;
      plain[0x10] = 0x26;
      constexpr std::string_view name = "BT.DINF";
      std::ranges::copy(name, plain.begin() + 0x11);
      plain[0x18] = static_cast<u8>((pairing.size() - 1) >> 8);
      plain[0x19] = static_cast<u8>(pairing.size() - 1);
    }
    const std::array<u8, 16> synthetic_key{};
    auto aes = Common::AES::CreateContextEncrypt(synthetic_key.data());
    std::vector<u8> encrypted(plain.size());
    ASSERT_TRUE(aes->CryptIvZero(plain.data(), encrypted.data(), plain.size()));
    PutLogicalData(SYSCONF_CLUSTER * 0x4000, encrypted);
  }

  void StoreSuperblock(u32 copy = 0)
  {
    PutLogicalData(SUPERBLOCK_START + copy * sizeof(Superblock),
                   {reinterpret_cast<const u8*>(m_superblock.get()), sizeof(Superblock)});
  }

  void CheckWithoutWrites(bool expected)
  {
    StoreSuperblock();
    const auto before = m_memory->m_nand_overlay;
    EXPECT_EQ(m_memory->InstallEmulatedWiimotePairings(), expected);
    EXPECT_TRUE(m_memory->m_nand_overlay == before);
  }

  void RemoveFirstSuperblockPage() { m_memory->m_nand_overlay.erase(SUPERBLOCK_START / 0x800); }

  void UseUnwritableJournal() { m_memory->m_nand_journal = std::make_unique<NANDJournal>(); }

  bool Install() { return m_memory->InstallEmulatedWiimotePairings(); }

  NANDJournal::Page ReadSysconfPage(u32 page)
  {
    NANDJournal::Page raw{};
    EXPECT_TRUE(m_memory->ReadRawNANDPage(SYSCONF_CLUSTER * 8 + page, &raw));
    return raw;
  }

  static std::array<u8, 16> CalculateECC(const NANDJournal::Page& page)
  {
    return StarletMemory::CalculateNANDECC(page.data());
  }

  void MarkSysconfSpare()
  {
    for (u32 page = 0; page < 8; ++page)
    {
      auto& raw = m_memory->m_nand_overlay.at(SYSCONF_CLUSTER * 8 + page);
      for (size_t i = 0; i < 0x30; ++i)
        raw[0x800 + i] = i == 0 ? 0xff : static_cast<u8>(0x40 + i + page);
    }
  }
  std::unique_ptr<StarletMemory> m_memory;
  std::unique_ptr<Superblock> m_superblock;
};

TEST_F(StarletPairingTest, MissingShared2DoesNotPreventBootOrWriteNAND)
{
  CheckWithoutWrites(true);
}

TEST_F(StarletPairingTest, MissingSysDoesNotPreventBootOrWriteNAND)
{
  AddPath(1);
  CheckWithoutWrites(true);
}

TEST_F(StarletPairingTest, MissingSysconfDoesNotPreventBootOrWriteNAND)
{
  AddPath(2);
  CheckWithoutWrites(true);
}

TEST_F(StarletPairingTest, AbsentBluetoothEntryDoesNotPreventBootOrWriteNAND)
{
  SetSysconf(false);
  CheckWithoutWrites(true);
}

TEST_F(StarletPairingTest, UnreadableNANDRemainsAnError)
{
  StoreSuperblock();
  RemoveFirstSuperblockPage();
  EXPECT_FALSE(Install());
}

TEST_F(StarletPairingTest, MalformedDirectoryRemainsAnError)
{
  m_superblock->fst[0].sub = DiscIO::NANDImporter::FSTEntryCount;
  CheckWithoutWrites(false);
}

TEST_F(StarletPairingTest, CyclicSiblingListFailsInsteadOfHanging)
{
  m_superblock->fst[0].sub = 1;
  m_superblock->fst[1].sib = 1;
  CheckWithoutWrites(false);
}

TEST_F(StarletPairingTest, NewestFormattedSuperblockDoesNotRestoreOlderPairings)
{
  SetSysconf(true);
  StoreSuperblock();
  // The formatter's newest root has no shared2, although an old copy still has it.
  m_superblock->version = 2;
  m_superblock->fst[0].sub = 0xffff;
  StoreSuperblock(1);
  // Keep the older copy and all raw data intact while selecting the newer root.
  EXPECT_TRUE(Install());
}

TEST_F(StarletPairingTest, EmptyExistingSysconfRemainsAnError)
{
  AddPath(3);
  CheckWithoutWrites(false);
}

TEST_F(StarletPairingTest, FailedPairingWriteRemainsAnErrorAndDoesNotAlterOverlay)
{
  SetSysconf(true);
  UseUnwritableJournal();
  CheckWithoutWrites(false);
}

TEST_F(StarletPairingTest, ExistingEmptyPairingRecordIsStillPopulated)
{
  SetSysconf(true);
  StoreSuperblock();
  EXPECT_TRUE(Install());
  // Once installed, the trusted identities are reused without another NAND write.
  CheckWithoutWrites(true);
}

TEST_F(StarletPairingTest, PairingWriteUsesNativeSFFSAuthenticationLayout)
{
  SetSysconf(true);
  m_superblock->fst[3].uid = 0x12345678;
  m_superblock->fst[3].x3 = 0x9abcdef0;
  StoreSuperblock();
  MarkSysconfSpare();
  std::array<NANDJournal::Page, 8> before;
  for (u32 page = 0; page < 8; ++page)
    before[page] = ReadSysconfPage(page);
  ASSERT_TRUE(Install());

  std::array<NANDJournal::Page, 8> after;
  std::vector<u8> ciphertext(0x4000);
  for (u32 page = 0; page < 8; ++page)
  {
    after[page] = ReadSysconfPage(page);
    std::copy_n(after[page].begin(), 0x800, ciphertext.begin() + page * 0x800);
  }
  const std::array<u8, 16> aes_key{};
  auto aes = Common::AES::CreateContextDecrypt(aes_key.data());
  std::vector<u8> authenticated(0x4040);
  ASSERT_TRUE(aes->CryptIvZero(ciphertext.data(), authenticated.data() + 0x40, ciphertext.size()));
  const auto pairing = IOS::HLE::GetEmulatedWiimotePairingData();
  EXPECT_TRUE(
      std::ranges::equal(pairing, std::span(authenticated).subspan(0x40 + 0x1a, pairing.size())));

  // Independent fixture salt: UID, file name, cluster ordinal, FST index and x3.
  const std::array<u8, 0x1c> salt = {0x12, 0x34, 0x56, 0x78, 'S',  'Y',  'S',  'C', 'O', 'N',
                                     'F',  0,    0,    0,    0,    0,    0,    0,   0,   0,
                                     0,    0,    0,    3,    0x9a, 0xbc, 0xde, 0xf0};
  std::ranges::copy(salt, authenticated.begin());
  const std::array<u8, 20> hmac_key{};
  std::array<u8, 20> hmac{};
  ASSERT_TRUE(Common::HMAC::HMACWithSHA1(hmac_key, authenticated, hmac.data()));

  // Native SFFS puts the full HMAC at spare byte 1 in page 6. The second
  // copy is split: 12 bytes at page 6 byte 21, then 8 at page 7 byte 1.
  auto expected = before;
  std::ranges::copy(hmac, expected[6].begin() + 0x801);
  std::copy_n(hmac.begin(), 12, expected[6].begin() + 0x815);
  std::copy_n(hmac.begin() + 12, 8, expected[7].begin() + 0x801);
  for (u32 page = 0; page < 8; ++page)
  {
    SCOPED_TRACE(page);
    EXPECT_TRUE(std::ranges::equal(std::span(after[page]).subspan(0x800, 0x30),
                                   std::span(expected[page]).subspan(0x800, 0x30)));
    EXPECT_TRUE(
        std::ranges::equal(CalculateECC(after[page]), std::span(after[page]).subspan(0x830, 16)));
  }
}
}  // namespace IOS::LLE