// SPDX-License-Identifier: MIT
#include "Common/cpp-optparse/OptionParser.h"
#include "Common/Config.h"
#include "Common/FEXServerClient.h"
#include "Common/HostFeatures.h"
#include "git_version.h"
#include <FEXCore/Config/Config.h>
#include <FEXCore/Utils/PrctlUtils.h>
#include <FEXCore/fextl/fmt.h>
#include <FEXCore/fextl/string.h>
#include <FEXHeaderUtils/Filesystem.h>

#include <cstdio>
#include <filesystem>
#include <string>
#include <sys/prctl.h>
#include <signal.h>
#include <ucontext.h>

#ifdef ARCHITECTURE_arm64
namespace {
struct TSOEmulationFacts {
  bool LSE {}, LSE2 {};
  bool HardwareTSO {};
  bool LRCPC1 {}, LRCPC2 {}, LRCPC3 {};
};

bool CheckForHardwareTSO() {
  // Check to see if this is supported.
  auto Result = prctl(PR_GET_MEM_MODEL, 0, 0, 0, 0);
  if (Result == -1) {
    // Unsupported, early exit.
    return false;
  }

  if (Result == PR_SET_MEM_MODEL_DEFAULT) {
    // Try to set the TSO mode if we are currently default.
    Result = prctl(PR_SET_MEM_MODEL, PR_SET_MEM_MODEL_TSO, 0, 0, 0);
    if (Result == 0) {
      Result = prctl(PR_SET_MEM_MODEL, PR_SET_MEM_MODEL_DEFAULT, 0, 0, 0);
      return true;
    }
  }
  return false;
}

enum ISAR0_FIELDS {
  LSE = 20,
};

enum ISAR1_FIELDS {
  LRCPC = 20,
};

enum MMFR2_FIELDS {
  AT = 32,
};

constexpr static uint32_t IDFIELDMASK = 0b1111;
uint64_t GetISAR0() {
  uint64_t Result {};
  asm("mrs %0, ID_AA64ISAR0_EL1;" : "=r"(Result));
  return Result;
}

uint64_t GetISAR1() {
  uint64_t Result {};
  asm("mrs %0, ID_AA64ISAR1_EL1;" : "=r"(Result));
  return Result;
}

uint64_t GetMMFR2() {
  uint64_t Result {};
  asm("mrs %0, ID_AA64MMFR2_EL1;" : "=r"(Result));
  return Result;
}

TSOEmulationFacts GetTSOEmulationFacts() {
  const auto ISAR0 = GetISAR0();
  const auto ISAR1 = GetISAR1();
  const auto MMFR2 = GetMMFR2();

  return {
    .LSE = ((ISAR0 >> ISAR0_FIELDS::LSE) & IDFIELDMASK) >= 0b0010,
    .LSE2 = ((MMFR2 >> MMFR2_FIELDS::AT) & IDFIELDMASK) >= 0b0001,
    .HardwareTSO = CheckForHardwareTSO(),
    .LRCPC1 = ((ISAR1 >> ISAR1_FIELDS::LRCPC) & IDFIELDMASK) >= 0b0001,
    .LRCPC2 = ((ISAR1 >> ISAR1_FIELDS::LRCPC) & IDFIELDMASK) >= 0b0010,
    .LRCPC3 = ((ISAR1 >> ISAR1_FIELDS::LRCPC) & IDFIELDMASK) >= 0b0011,
  };
}
} // namespace

namespace SIGBUSTest {
static bool* FaultArray {};

__attribute__((naked)) static void atomic_load_u16(std::byte* Data) {
  asm volatile(R"(
    ldarh w1, [x0];
    ret;
    )" ::
                 : "x1", "memory");
}

__attribute__((naked)) static void atomic_load_u32(std::byte* Data) {
  asm volatile(R"(
    ldar w1, [x0];
    ret;
    )" ::
                 : "x1", "memory");
}

__attribute__((naked)) static void atomic_load_u64(std::byte* Data) {
  asm volatile(R"(
    ldar x1, [x0];
    ret;
    )" ::
                 : "x1", "memory");
}
__attribute__((naked)) static void atomic_load_u128(std::byte* Data) {
  asm volatile(R"(
    ldaxp x1, x2, [x0];
    ret;
    )" ::
                 : "x1", "x2", "x3", "memory");
}

__attribute__((naked)) static void atomic_set_u16(std::byte* Data, uint16_t value) {
  asm volatile(R"(
    .word 0x78e13002; // ldsetalh w1, w2, [x0];
    ret;
    )" ::
                 : "memory");
}

__attribute__((naked)) static void atomic_set_u32(std::byte* Data, uint32_t value) {
  asm volatile(R"(
    .word 0xb8e13002; // ldsetal w1, w2, [x0];
    ret;
    )" ::
                 : "memory");
}

__attribute__((naked)) static void atomic_set_u64(std::byte* Data, uint64_t value) {
  asm volatile(R"(
    .word 0xf8e13002; // ldsetal x1, x2, [x0];
    ret;
    )" ::
                 : "memory");
}
__attribute__((naked)) static void atomic_set_u128_impl(__uint128_t expected, __uint128_t desired, std::byte* Data) {
  asm volatile(R"(
    .word 0x4860fc82; // caspal x0, x1, x2, x3, [x4];
    ret;
    )" ::
                 : "memory");
}

static inline void atomic_set_u128(std::byte* Data, __uint128_t value) {
  atomic_set_u128_impl(*reinterpret_cast<__uint128_t*>(Data), value, Data);
}

static void HandleSIGBUS(int, siginfo_t* info, void* context) {
  FaultArray[reinterpret_cast<uintptr_t>(info->si_addr) & 63] = true;

  ucontext_t* ucontext = (ucontext_t*)context;
  mcontext_t* mcontext = &ucontext->uc_mcontext;
  // Skip the stlr.
  mcontext->pc += 4;
}

static bool CalculatedFaultOffsets {};
static bool FaultOffset_16bit[64] {};
static bool FaultOffset_32bit[64] {};
static bool FaultOffset_64bit[64] {};
static bool FaultOffset_128bit[64] {};

static bool FaultOffset_RMW_16bit[64] {};
static bool FaultOffset_RMW_32bit[64] {};
static bool FaultOffset_RMW_64bit[64] {};
static bool FaultOffset_RMW_128bit[64] {};

static void RunFaultTests() {
  if (CalculatedFaultOffsets) {
    return;
  }

  struct sigaction act {};
  act.sa_sigaction = HandleSIGBUS;
  act.sa_flags = SA_SIGINFO;
  sigaction(SIGBUS, &act, &act);
  auto ptr = reinterpret_cast<std::byte*>(mmap(nullptr, 4096, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));

  auto test_fault = [](bool* FaultOffsets, auto AccessFunction, std::byte* AccessArray) {
    FaultArray = FaultOffsets;
    for (size_t i = 0; i < 64; ++i) {
      AccessFunction(AccessArray + i);
    }
  };

  auto test_rmw_fault = [](bool* FaultOffsets, auto AccessFunction, std::byte* AccessArray) {
    FaultArray = FaultOffsets;
    for (size_t i = 0; i < 64; ++i) {
      AccessFunction(AccessArray + i, 1);
    }
  };

  test_fault(FaultOffset_16bit, atomic_load_u16, ptr);
  test_fault(FaultOffset_32bit, atomic_load_u32, ptr);
  test_fault(FaultOffset_64bit, atomic_load_u64, ptr);
  test_fault(FaultOffset_128bit, atomic_load_u128, ptr);

  auto TSOFacts = GetTSOEmulationFacts();

  if (TSOFacts.LSE) {
    test_rmw_fault(FaultOffset_RMW_16bit, atomic_set_u16, ptr);
    test_rmw_fault(FaultOffset_RMW_32bit, atomic_set_u32, ptr);
    test_rmw_fault(FaultOffset_RMW_64bit, atomic_set_u64, ptr);
    test_rmw_fault(FaultOffset_RMW_128bit, atomic_set_u128, ptr);
  }

  munmap(ptr, 4096);
  sigaction(SIGBUS, &act, nullptr);

  CalculatedFaultOffsets = true;
}

static void PrintSIGBUSInfo() {
  RunFaultTests();

  auto print_granule = [](const char* size, bool* FaultArray) {
    std::string output {};
    for (size_t i = 0; i < 64; ++i) {
      if (i && (i % 16 == 0)) {
        output += " ";
      }

      if (FaultArray[i]) {
        output += "\e[31m■\e[0m";
      } else {
        output += "\e[32m■\e[0m";
      }
    }

    fprintf(stdout, "%s: %s\n", size, output.c_str());
  };

  auto TSOFacts = GetTSOEmulationFacts();
  const bool RMWIsDifferent = TSOFacts.LSE && (memcmp(FaultOffset_16bit, FaultOffset_RMW_16bit, sizeof(FaultOffset_16bit)) != 0 ||
                                               memcmp(FaultOffset_32bit, FaultOffset_RMW_32bit, sizeof(FaultOffset_32bit)) != 0 ||
                                               memcmp(FaultOffset_64bit, FaultOffset_RMW_64bit, sizeof(FaultOffset_64bit)) != 0 ||
                                               memcmp(FaultOffset_128bit, FaultOffset_RMW_128bit, sizeof(FaultOffset_128bit)) != 0);

  if (!RMWIsDifferent) {
    fprintf(stdout, "Fault Granularity: Split every 16 bytes\n");
  } else {
    fprintf(stdout, "Load/Store Fault Granularity: Split every 16 bytes\n");
  }
  print_granule(" 16-bit", FaultOffset_16bit);
  print_granule(" 32-bit", FaultOffset_32bit);
  print_granule(" 64-bit", FaultOffset_64bit);
  print_granule("128-bit", FaultOffset_128bit);

  if (RMWIsDifferent) {
    fprintf(stdout, "RMW Atomic Fault Granularity: Split every 16 bytes\n");
    print_granule(" 16-bit", FaultOffset_RMW_16bit);
    print_granule(" 32-bit", FaultOffset_RMW_32bit);
    print_granule(" 64-bit", FaultOffset_RMW_64bit);
    print_granule("128-bit", FaultOffset_RMW_128bit);
  }
}

struct FirstFaultInformation {
  int32_t LoadStoreFaultAlignment {};
  int32_t RMWFaultAlignment {};
};

static FirstFaultInformation CalculateFirstFaultInformation() {
  RunFaultTests();
  FirstFaultInformation Info {};
  auto FindFirstFaultOffset = [](bool FaultOffsets[64]) -> int32_t {
    for (int32_t i = 0; i < 64; ++i) {
      if (FaultOffsets[i]) {
        return i;
      }
    }

    return -1;
  };

  Info.LoadStoreFaultAlignment = FindFirstFaultOffset(FaultOffset_16bit) + 1;
  Info.RMWFaultAlignment = FindFirstFaultOffset(FaultOffset_RMW_16bit) + 1;
  return Info;
}

} // namespace SIGBUSTest

static void PrintTSOInfo() {
  auto TSOFacts = GetTSOEmulationFacts();
  FEX_CONFIG_OPT(TSOEnabled, TSOENABLED);
  FEX_CONFIG_OPT(MemcpySetTSOEnabled, MEMCPYSETTSOENABLED);
  FEX_CONFIG_OPT(VectorTSOEnabled, VECTORTSOENABLED);
  FEX_CONFIG_OPT(HalfBarrierTSOEnabled, HALFBARRIERTSOENABLED);
  FEX_CONFIG_OPT(StrictInProcessSplitLocks, STRICTINPROCESSSPLITLOCKS);

  const char* GPRMemoryTSOEmulation {};
  const char* MemcpyMemoryTSOEmulation {};
  const char* VectorMemoryTSOEmulation {};
  const char* UnalignedMemoryLoadStoreTSOEmulation {};
  const char* SplitLock16BEmulationType {};
  const char* SplitLock16BConfigurationType {};
  std::string UnalignedMemoryLoadStoreAlignmentGranularity {};
  std::string UnalignedRMWAlignmentGranularity {};

  if (TSOFacts.HardwareTSO) {
    GPRMemoryTSOEmulation = "\e[32mHardware TSO\e[0m";
  } else if (TSOFacts.LRCPC3) {
    GPRMemoryTSOEmulation = "\e[32mLRCPC3\e[0m";
  } else if (TSOFacts.LRCPC2) {
    GPRMemoryTSOEmulation = "\e[32mLRCPC2\e[0m";
  } else if (TSOFacts.LRCPC1) {
    GPRMemoryTSOEmulation = "\e[32mLRCPC\e[0m";
  } else {
    GPRMemoryTSOEmulation = "\e[31mAtomics\e[0m";
  }

  // Memcpy only uses Hardware TSO, LRCPC, and Atomics.
  if (TSOFacts.HardwareTSO) {
    MemcpyMemoryTSOEmulation = "\e[32mHardware TSO\e[0m";
  } else if (TSOFacts.LRCPC1) {
    MemcpyMemoryTSOEmulation = "\e[32mLRCPC\e[0m";
  } else {
    MemcpyMemoryTSOEmulation = "\e[31mAtomics\e[0m";
  }

  if (TSOFacts.HardwareTSO) {
    VectorMemoryTSOEmulation = "\e[32mHardware TSO\e[0m";
  } else if (TSOFacts.LRCPC3) {
    VectorMemoryTSOEmulation = "\e[32mLRCPC3\e[0m";
  } else {
    VectorMemoryTSOEmulation = "\e[31mHalf-Barriers\e[0m";
  }

  if (TSOFacts.HardwareTSO) {
    UnalignedMemoryLoadStoreTSOEmulation = "\e[32mHardware TSO\e[0m";
  } else {
    UnalignedMemoryLoadStoreTSOEmulation = "\e[31mHalf-Barriers\e[0m";
  }

  const auto FFInfo = SIGBUSTest::CalculateFirstFaultInformation();

  if (FFInfo.RMWFaultAlignment >= 64) {
    SplitLock16BEmulationType = "\e[32mHardware cacheline unaligned atomics\e[0m";
    SplitLock16BConfigurationType = "\e[32mTear-free\e[0m";
  } else {
    SplitLock16BEmulationType = TSOFacts.LSE ? "\e[31mTearing CAS loops\e[0m" : "\e[31mTearing LL/SC loops\e[0m";
    SplitLock16BConfigurationType = StrictInProcessSplitLocks() ? "In-process mutex" : "Tearing";
  }

  if (FFInfo.LoadStoreFaultAlignment != 1) {
    UnalignedMemoryLoadStoreAlignmentGranularity = fmt::format("\e[32m{}-byte\e[0m", FFInfo.LoadStoreFaultAlignment);
  } else {
    UnalignedMemoryLoadStoreAlignmentGranularity = TSOFacts.LSE2 ? "\e[32m16-byte\e[0m" : "\e[31mNatural alignment\e[0m";
  }

  if (FFInfo.LoadStoreFaultAlignment != FFInfo.RMWFaultAlignment) {
    if (FFInfo.LoadStoreFaultAlignment != 1) {
      UnalignedRMWAlignmentGranularity = fmt::format("\e[32m{}-byte\e[0m", FFInfo.RMWFaultAlignment);
    } else {
      UnalignedRMWAlignmentGranularity = TSOFacts.LSE2 ? "\e[32m16-byte\e[0m" : "\e[31mNatural alignment\e[0m";
    }
  }

  fprintf(stdout, "Hardware Features:\n");
  fprintf(stdout, "\tMemory atomics emulation method:      %s\n", TSOFacts.LSE ? "\e[32mLSE\e[0m" : "\e[31mLL/SC\e[0m");
  fprintf(stdout, "\tUnaligned atomic memory granularity:  %s\n", UnalignedMemoryLoadStoreAlignmentGranularity.c_str());
  if (FFInfo.LoadStoreFaultAlignment != FFInfo.RMWFaultAlignment) {
    fprintf(stdout, "\tUnaligned atomic RMW granularity:     %s\n", UnalignedRMWAlignmentGranularity.c_str());
  }
  fprintf(stdout, "\tUnaligned memory loadstore emulation: %s\n", UnalignedMemoryLoadStoreTSOEmulation);
  fprintf(stdout, "\t16-Byte split-lock atomic emulation:  %s\n", SplitLock16BEmulationType);
  fprintf(stdout, "\t64-Byte split-lock atomic emulation:  %s\n", TSOFacts.LSE ? "\e[31mTearing CAS loops\e[0m" : "\e[31mTearing LL/SC loops\e[0m");
  fprintf(stdout, "\tGPR memory model emulation:           %s\n", GPRMemoryTSOEmulation);
  fprintf(stdout, "\tMemcpy memory model emulation:        %s\n", MemcpyMemoryTSOEmulation);
  fprintf(stdout, "\tVector memory model emulation:        %s\n", VectorMemoryTSOEmulation);

  fprintf(stdout, "\nConfiguration:\n");
  fprintf(stdout, "\tTSO Emulation:                        %s\n", TSOEnabled() ? "Enabled" : "Disabled");
  fprintf(stdout, "\tMemcpy TSO Emulation:                 %s\n", TSOEnabled() && MemcpySetTSOEnabled() ? "Enabled" : "Disabled");
  fprintf(stdout, "\tVector TSO Emulation:                 %s\n", TSOEnabled() && VectorTSOEnabled() ? "Enabled" : "Disabled");
  fprintf(stdout, "\tHalf-barrier unaligned TSO emulation: %s\n", TSOEnabled() && HalfBarrierTSOEnabled() ? "Enabled" : "Disabled");
  fprintf(stdout, "\t16-Byte strict split-lock emulation:  %s\n", SplitLock16BConfigurationType);
  fprintf(stdout, "\t64-Byte strict split-lock emulation:  %s\n", StrictInProcessSplitLocks() ? "In-process mutex" : "Tearing");
}

static void PrintIDRegInfo() {
  auto Features = FEX::GetCPUFeaturesFromIDRegisters();
  fextl::string features {};
  features += fmt::format("isar0=0x{:x},", Features.ISAR0.Get());
  features += fmt::format("isar1=0x{:x},", Features.ISAR1.Get());
  features += fmt::format("isar2=0x{:x},", Features.ISAR2.Get());
  features += fmt::format("pfr0=0x{:x},", Features.PFR0.Get());
  features += fmt::format("pfr1=0x{:x},", Features.PFR1.Get());
  features += fmt::format("midr=0x{:x},", Features.MIDR.Get());
  features += fmt::format("mmfr0=0x{:x},", Features.MMFR0.Get());
  features += fmt::format("mmfr1=0x{:x},", Features.MMFR1.Get());
  features += fmt::format("mmfr2=0x{:x},", Features.MMFR2.Get());
  features += fmt::format("mmfr3=0x{:x},", Features.MMFR3.Get());
  features += fmt::format("zfr0=0x{:x},", Features.ZFR0.Get());
  features += fmt::format("dczid=0x{:x},", Features.DCZID.Get());
  features += fmt::format("svevl=0x{:x}", Features.SVEVL.Get());
  fprintf(stderr, "Features: '%s'\n", features.c_str());
}
#endif

int main(int argc, char** argv, char** envp) {
  FEX::Config::InitializeConfigs(FEX::Config::PortableInformation {});
  FEXCore::Config::Initialize();
  FEXCore::Config::AddLayer(FEX::Config::CreateGlobalMainLayer());
  FEXCore::Config::AddLayer(FEX::Config::CreateMainLayer());
  // No FEX arguments passed through command line
  FEXCore::Config::AddLayer(FEX::Config::CreateEnvironmentLayer(envp));

  // Load the arguments
  optparse::OptionParser Parser = optparse::OptionParser().description("Simple application to get a couple of FEX options");

  Parser.add_option("--install-prefix").action("store_true").help("Print the FEX install prefix");

  Parser.add_option("--app").help("Load an application profile for this application if it exists");

  Parser.add_option("--current-rootfs").action("store_true").help("Print the directory that contains the FEX rootfs. Mounted in the case of squashfs");

#ifdef ARCHITECTURE_arm64
  Parser.add_option("--tso-emulation-info").action("store_true").help("Print how FEX is emulating the x86-TSO memory model.");
  Parser.add_option("--test-fault-granularity").action("store_true").help("Show SIGBUS fault granularity");
  Parser.add_option("--identification-reg-info").action("store_true").help("Print identification registers");
  Parser.add_option("-e", "--all-emu-info").action("store_true").help("Prints all relevant emulation related information");
#endif

  Parser.add_option("--version").action("store_true").help("Print the installed FEX-Emu version");

  optparse::Values Options = Parser.parse_args(argc, argv);

  if (Options.is_set_by_user("app")) {
    // Load the application config if one was provided
    const auto ProgramName = FHU::Filesystem::GetFilename(Options["app"]);
    FEXCore::Config::AddLayer(FEX::Config::CreateAppLayer(ProgramName, FEXCore::Config::LayerType::LAYER_GLOBAL_APP));
    FEXCore::Config::AddLayer(FEX::Config::CreateAppLayer(ProgramName, FEXCore::Config::LayerType::LAYER_LOCAL_APP));

    auto SteamID = getenv("SteamAppId");
    if (SteamID) {
      // If a SteamID exists then let's search for Steam application configs as well.
      // We want to key off both the SteamAppId number /and/ the executable since we may not want to thunk all binaries.
      const auto SteamAppName = fextl::fmt::format("Steam_{}_{}", SteamID, ProgramName);
      FEXCore::Config::AddLayer(FEX::Config::CreateAppLayer(SteamAppName, FEXCore::Config::LayerType::LAYER_GLOBAL_STEAM_APP));
      FEXCore::Config::AddLayer(FEX::Config::CreateAppLayer(SteamAppName, FEXCore::Config::LayerType::LAYER_LOCAL_STEAM_APP));
    }
  }

  FEXCore::Config::Load();

  // Reload the meta layer
  FEXCore::Config::ReloadMetaLayer();

  const bool IsAllEmuInfo = Options.is_set_by_user("all_emu_info");

  if (IsAllEmuInfo || Options.is_set_by_user("version")) {
    fprintf(stdout, GIT_DESCRIBE_STRING "\n");
  }

  if (Options.is_set_by_user("install_prefix")) {
    char SelfPath[PATH_MAX];
    auto Result = readlink("/proc/self/exe", SelfPath, PATH_MAX);
    if (Result == -1) {
      Result = 0;
    }
    auto InstallPrefix = std::filesystem::path(&SelfPath[0], &SelfPath[Result]).parent_path().parent_path().string();
    fprintf(stdout, "%s\n", InstallPrefix.c_str());
  }

  if (Options.is_set_by_user("current_rootfs")) {
    int ServerFD = FEXServerClient::ConnectToServer();
    if (ServerFD != -1) {
      auto RootFS = FEXServerClient::RequestRootFSPath(ServerFD);
      if (!RootFS.empty()) {
        fprintf(stdout, "%s\n", RootFS.c_str());
      }
    }
  }

#ifdef ARCHITECTURE_arm64
  if (IsAllEmuInfo || Options.is_set_by_user("tso_emulation_info")) {
    PrintTSOInfo();
  }

  if (IsAllEmuInfo || Options.is_set_by_user("identification_reg_info")) {
    PrintIDRegInfo();
  }

  if (IsAllEmuInfo || Options.is_set_by_user("test_fault_granularity")) {
    SIGBUSTest::PrintSIGBUSInfo();
  }
#endif

  return 0;
}