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

#include <array>
#include <string>
#include <unordered_map>
#include <utility>

#include <gtest/gtest.h>

#include "Core/IOS/Starlet/ARMCore.h"
#include "Core/IOS/Starlet/BootstrapELF.h"

namespace IOS::LLE
{
namespace
{
constexpr u32 TEXT_BASE = 0xffff0000;
constexpr u32 DATA_BASE = 0xfffe0000;

class RecordingBus final : public ARMBus
{
public:
  u8 Read8(u32 address) override
  {
    const auto it = bytes.find(address);
    return it == bytes.end() ? 0 : it->second;
  }
  void Write8(u32 address, u8 value) override { bytes[address] = value; }

  std::unordered_map<u32, u8> bytes;
};

void Put16(std::array<u8, 0x300>& image, size_t offset, u16 value)
{
  image[offset] = static_cast<u8>(value >> 8);
  image[offset + 1] = static_cast<u8>(value);
}

void Put32(std::array<u8, 0x300>& image, size_t offset, u32 value)
{
  image[offset] = static_cast<u8>(value >> 24);
  image[offset + 1] = static_cast<u8>(value >> 16);
  image[offset + 2] = static_cast<u8>(value >> 8);
  image[offset + 3] = static_cast<u8>(value);
}

std::array<u8, 0x300> MakeImage()
{
  std::array<u8, 0x300> image{};
  image[0] = 0x7f;
  image[1] = 'E';
  image[2] = 'L';
  image[3] = 'F';
  image[4] = 1;  // ELF32
  image[5] = 2;  // big endian
  image[6] = 1;
  Put16(image, 16, 2);   // ET_EXEC
  Put16(image, 18, 40);  // ARM
  Put32(image, 20, 1);   // ELF version
  Put32(image, 24, TEXT_BASE);
  Put32(image, 28, 52);     // program headers
  Put32(image, 32, 0x200);  // section headers
  Put16(image, 40, 52);
  Put16(image, 42, 32);
  Put16(image, 44, 2);
  Put16(image, 46, 40);
  Put16(image, 48, 3);

  Put32(image, 52, 1);
  Put32(image, 56, 0x100);
  Put32(image, 60, TEXT_BASE);
  Put32(image, 64, TEXT_BASE);
  Put32(image, 68, 4);
  Put32(image, 72, 8);
  Put32(image, 76, 5);
  image[0x100] = 0xea;  // ARM branch instruction

  Put32(image, 84, 1);
  Put32(image, 88, 0x110);
  Put32(image, 92, DATA_BASE);
  Put32(image, 96, DATA_BASE);
  Put32(image, 100, 4);
  Put32(image, 104, 0x420);
  Put32(image, 108, 6);
  image[0x110] = 0x12;
  image[0x111] = 0x34;
  image[0x112] = 0x56;
  image[0x113] = 0x78;

  constexpr char names[] = "\0REV_TRACE\0__fault_code\0__fault_lr\0__fault_spsr";
  for (size_t index = 0; index < sizeof(names); ++index)
    image[0x180 + index] = static_cast<u8>(names[index]);
  Put32(image, 0x228 + 4, 3);  // SHT_STRTAB
  Put32(image, 0x228 + 16, 0x180);
  Put32(image, 0x228 + 20, sizeof(names));
  Put32(image, 0x250 + 4, 2);  // SHT_SYMTAB
  Put32(image, 0x250 + 16, 0x120);
  Put32(image, 0x250 + 20, 5 * 16);
  Put32(image, 0x250 + 24, 1);
  Put32(image, 0x250 + 36, 16);

  const auto symbol = [&image](size_t index, u32 name, u32 address, u32 size) {
    const size_t offset = 0x120 + index * 16;
    Put32(image, offset, name);
    Put32(image, offset + 4, address);
    Put32(image, offset + 8, size);
  };
  symbol(1, 1, DATA_BASE + 8, 1032);
  symbol(2, 11, DATA_BASE + 0x410, 4);
  symbol(3, 24, DATA_BASE + 0x414, 4);
  symbol(4, 35, DATA_BASE + 0x418, 4);
  return image;
}

TEST(BootstrapELFTest, LoadsSegmentsAndZerosUninitializedData)
{
  const auto image = MakeImage();
  RecordingBus bus;
  BootstrapELFInfo info;
  std::string error;
  ASSERT_TRUE(LoadBootstrapELF(image, bus, &info, &error)) << error;
  EXPECT_EQ(info.entry, TEXT_BASE);
  EXPECT_EQ(info.trace, DATA_BASE + 8);
  EXPECT_EQ(info.fault_spsr, DATA_BASE + 0x418);
  EXPECT_EQ(bus.Read8(TEXT_BASE), 0xea);
  EXPECT_EQ(bus.Read8(TEXT_BASE + 4), 0);
  EXPECT_EQ(bus.Read32(DATA_BASE), 0x12345678u);
  EXPECT_EQ(bus.Read32(DATA_BASE + 0x41c), 0u);
  EXPECT_EQ(bus.bytes.size(), size_t{8 + 0x420});
}

TEST(BootstrapELFTest, RejectsMalformedImageBeforeWritingMemory)
{
  auto image = MakeImage();
  RecordingBus bus;
  BootstrapELFInfo info;
  std::string error;
  image[5] = 1;
  EXPECT_FALSE(LoadBootstrapELF(image, bus, &info, &error));
  EXPECT_TRUE(bus.bytes.empty());

  image = MakeImage();
  Put32(image, 0x120 + 2 * 16 + 4, DATA_BASE + 0x420);
  EXPECT_FALSE(LoadBootstrapELF(image, bus, &info, &error));
  EXPECT_TRUE(bus.bytes.empty());

  image = MakeImage();
  Put32(image, 88, 0x300);
  EXPECT_FALSE(LoadBootstrapELF(image, bus, &info, &error));
  EXPECT_TRUE(bus.bytes.empty());
}

// Revolution IOS Phase 06H: the optional MEM2 text segment of the emulator-only
// disposable-image kernel, accepted only when the caller allows it.
constexpr u32 MEM2_TEXT = 0x12600000;

std::array<u8, 0x300> MakeMem2Image(u32 address, u32 flags)
{
  auto image = MakeImage();
  Put16(image, 44, 3);
  Put32(image, 116, 1);
  Put32(image, 120, 0x114);
  Put32(image, 124, address);
  Put32(image, 128, address);
  Put32(image, 132, 4);
  Put32(image, 136, 0x10);
  Put32(image, 140, flags);
  Put32(image, 0x114, 0xe12fff1e);  // BX LR
  return image;
}

TEST(BootstrapELFTest, Mem2TextSegmentOnlyWhenAllowed)
{
  RecordingBus bus;
  BootstrapELFInfo info;
  std::string error;
  auto image = MakeMem2Image(MEM2_TEXT, 5);
  EXPECT_FALSE(LoadBootstrapELF(image, bus, &info, &error));
  EXPECT_TRUE(bus.bytes.empty());

  ASSERT_TRUE(LoadBootstrapELF(image, bus, &info, &error, true)) << error;
  EXPECT_EQ(bus.Read32(MEM2_TEXT), 0xe12fff1eu);
  EXPECT_EQ(bus.Read32(MEM2_TEXT + 0xc), 0u);
  EXPECT_EQ(bus.Read8(TEXT_BASE), 0xea);
  EXPECT_EQ(bus.Read32(DATA_BASE), 0x12345678u);
  EXPECT_EQ(bus.bytes.size(), size_t{8 + 0x420 + 0x10});

  RecordingBus plain;
  const auto two_segments = MakeImage();
  EXPECT_TRUE(LoadBootstrapELF(two_segments, plain, &info, &error, true)) << error;
}

TEST(BootstrapELFTest, RejectsInvalidMem2TextSegmentBeforeWritingMemory)
{
  RecordingBus bus;
  BootstrapELFInfo info;
  std::string error;
  // Writable, outside MEM2, crossing its end, or a second SRAM text segment.
  for (const auto& image : {MakeMem2Image(MEM2_TEXT, 6), MakeMem2Image(0x0fff0000, 5),
                            MakeMem2Image(0x13fffff8, 5), MakeMem2Image(TEXT_BASE, 5)})
  {
    EXPECT_FALSE(LoadBootstrapELF(image, bus, &info, &error, true));
    EXPECT_TRUE(bus.bytes.empty());
  }
}

TEST(ARMCoreTest, Armv5DoublewordTransfersAndAlignmentAbort)
{
  RecordingBus bus;
  bus.Write32(TEXT_BASE, 0xe1c020d0);      // LDRD r2, r3, [r0]
  bus.Write32(TEXT_BASE + 4, 0xe1c120f0);  // STRD r2, r3, [r1]
  bus.Write32(DATA_BASE, 0x11223344);
  bus.Write32(DATA_BASE + 4, 0x55667788);

  ARMCore core(bus);
  core.Reset(TEXT_BASE);
  core.SetBigEndian(true);
  core.SetRegister(0, DATA_BASE);
  core.SetRegister(1, DATA_BASE + 8);
  core.Step();
  EXPECT_EQ(core.GetRegister(2), 0x11223344u);
  EXPECT_EQ(core.GetRegister(3), 0x55667788u);
  core.Step();
  EXPECT_EQ(bus.Read32(DATA_BASE + 8), 0x11223344u);
  EXPECT_EQ(bus.Read32(DATA_BASE + 12), 0x55667788u);

  core.Reset(TEXT_BASE);
  core.SetBigEndian(true);
  core.SetRegister(0, DATA_BASE + 4);
  core.Step();
  EXPECT_EQ(core.GetRegister(15), TEXT_BASE + 0x10);
  EXPECT_EQ(core.GetRegister(14), TEXT_BASE + 8);
  EXPECT_EQ(core.GetCP15State().fault_address, DATA_BASE + 4);
}

// StarOS ADR-0014 R6: the optional panic words. The string table moves past the section headers
// to make room for two more symbols.
std::array<u8, 0x300> MakePanicImage(u32 state, u32 code)
{
  auto image = MakeImage();
  constexpr char names[] =
      "\0REV_TRACE\0__fault_code\0__fault_lr\0__fault_spsr\0__panic_state\0__panic_code";
  for (size_t index = 0; index < sizeof(names); ++index)
    image[0x278 + index] = static_cast<u8>(names[index]);
  Put32(image, 0x228 + 16, 0x278);
  Put32(image, 0x228 + 20, sizeof(names));
  Put32(image, 0x250 + 20, 7 * 16);
  const auto symbol = [&image](size_t index, u32 name, u32 address) {
    const size_t offset = 0x120 + index * 16;
    Put32(image, offset, name);
    Put32(image, offset + 4, address);
    Put32(image, offset + 8, 4);
  };
  // An omitted word keeps its slot, with the empty name.
  symbol(5, state != 0 ? 48 : 0, state);
  symbol(6, code != 0 ? 62 : 0, code);
  return image;
}

TEST(BootstrapELFTest, PanicWordsAreOptionalButBothInData)
{
  RecordingBus bus;
  BootstrapELFInfo info;
  std::string error;
  ASSERT_TRUE(LoadBootstrapELF(MakeImage(), bus, &info, &error)) << error;
  EXPECT_EQ(info.panic_state, 0u);
  EXPECT_EQ(info.panic_code, 0u);

  ASSERT_TRUE(LoadBootstrapELF(MakePanicImage(DATA_BASE + 0x41c, DATA_BASE + 4), bus, &info, &error))
      << error;
  EXPECT_EQ(info.panic_state, DATA_BASE + 0x41c);
  EXPECT_EQ(info.panic_code, DATA_BASE + 4);

  for (const auto [state, code] : {std::pair{DATA_BASE + 0x41c, 0u}, std::pair{0u, DATA_BASE + 4},
                                   std::pair{DATA_BASE + 0x41c, DATA_BASE + 0x420}})
  {
    RecordingBus rejected;
    EXPECT_FALSE(LoadBootstrapELF(MakePanicImage(state, code), rejected, &info, &error));
    EXPECT_EQ(error, "bootstrap ELF has an invalid panic symbol");
    EXPECT_TRUE(rejected.bytes.empty());
  }
}

TEST(BootstrapELFTest, PanicWordsMustNotAlias)
{
  BootstrapELFInfo info;
  std::string error;
  // The same word, or two words sharing bytes, in either order; all inside the 0x420-byte data
  // segment, so only the aliasing rule rejects them.
  for (const auto [state, code] :
       {std::pair{DATA_BASE + 4, DATA_BASE + 4}, std::pair{DATA_BASE + 4, DATA_BASE + 6},
        std::pair{DATA_BASE + 7, DATA_BASE + 4}, std::pair{DATA_BASE + 0x41c, DATA_BASE + 0x41a}})
  {
    RecordingBus rejected;
    EXPECT_FALSE(LoadBootstrapELF(MakePanicImage(state, code), rejected, &info, &error))
        << std::hex << state << ' ' << code;
    EXPECT_EQ(error, "bootstrap ELF has an invalid panic symbol");
    EXPECT_TRUE(rejected.bytes.empty());
  }
  // Adjacent words do not alias.
  for (const auto [state, code] :
       {std::pair{DATA_BASE + 4, DATA_BASE + 8}, std::pair{DATA_BASE + 8, DATA_BASE + 4}})
  {
    RecordingBus bus;
    EXPECT_TRUE(LoadBootstrapELF(MakePanicImage(state, code), bus, &info, &error)) << error;
    EXPECT_EQ(info.panic_state, state);
    EXPECT_EQ(info.panic_code, code);
  }
}
}  // namespace
}  // namespace IOS::LLE