// 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>

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

#ifdef ARCHITECTURE_arm64
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,
  };
}
#else
TSOEmulationFacts GetTSOEmulationFacts() {
  return {};
}
#endif
} // namespace

#ifdef ARCHITECTURE_arm64
namespace SIGBUSTest {
static bool* FaultArray {};

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

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

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

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;
}

void TestSIGBUS() {
  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 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());
  };

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

  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);

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

  fprintf(stdout, "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);
}
} // namespace SIGBUSTest
#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");

  Parser.add_option("--tso-emulation-info").action("store_true").help("Print how FEX is emulating the x86-TSO memory model.");

#ifdef ARCHITECTURE_arm64
  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");
#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();

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

#ifdef ARCHITECTURE_arm64
  if (Options.is_set_by_user("test_fault_granularity")) {
    SIGBUSTest::TestSIGBUS();
  }
#endif

  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());
      }
    }
  }

  if (Options.is_set_by_user("tso_emulation_info")) {
    auto TSOFacts = GetTSOEmulationFacts();
    const char* GPRMemoryTSOEmulation {};
    const char* MemcpyMemoryTSOEmulation {};
    const char* VectorMemoryTSOEmulation {};
    const char* UnalignedMemoryLoadStoreTSOEmulation {};

    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";
    }

    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", TSOFacts.LSE2 ? "\e[32m16-byte\e[0m" : "\e[31mNatural alignment\e[0m");
    ///< TODO: Once TME is supported by hardware this can change.
    fprintf(stdout, "\tUnaligned memory loadstore emulation: %s\n", UnalignedMemoryLoadStoreTSOEmulation);
    fprintf(stdout, "\t16-Byte split-lock atomic emulation:  %s\n", TSOFacts.LSE ? "\e[31mTearing CAS loops\e[0m" : "\e[31mTearing LL/SC loops\e[0m");
    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);

    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);
    fprintf(stderr, "Strict: %d\n", StrictInProcessSplitLocks());

    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", StrictInProcessSplitLocks() ? "In-process mutex" : "Tearing");
    fprintf(stdout, "\t64-Byte strict split-lock emulation:  %s\n", StrictInProcessSplitLocks() ? "In-process mutex" : "Tearing");
  }

#ifdef ARCHITECTURE_arm64
  if (Options.is_set_by_user("identification_reg_info")) {
    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("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

  return 0;
}