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

#include <algorithm>
#include <bit>
#include <concepts>

#include "Common/x64ABI.h"
#include "Common/x64Emitter.h"

#include <gtest/gtest.h>

using namespace Gen;

namespace
{
u32 ZeroParameterFunction()
{
  return 123;
}

template <std::integral T>
u32 OneParameterFunction(T a)
{
  return a + 23;
}

u32 TwoParameterFunction(u64 a, u64 b)
{
  return a * 10 + b + 3;
}

u32 ThreeParameterFunction(u64 a, u64 b, u64 c)
{
  return a * 10 + b + c / 10;
}

u32 FourParameterFunction(u64 a, u64 b, u32 c, u32 d)
{
  return a == 0x0102030405060708 && b == 0x090a0b0c0d0e0f10 && c == 0x11121314 && d == 0x15161718 ?
             123 :
             4;
}

class TestCallFunction : public X64CodeBlock
{
public:
  TestCallFunction() { AllocCodeSpace(4096); }

  template <typename F>
  void Emit(F f)
  {
    ResetCodePtr();

    m_code_pointer = GetCodePtr();
    {
      const Common::ScopedJITPageWriteAndNoExecute enable_jit_page_writes;

      ABI_PushRegistersAndAdjustStack(ABI_ALL_CALLEE_SAVED, 8, 16);
      f();
      ABI_PopRegistersAndAdjustStack(ABI_ALL_CALLEE_SAVED, 8, 16);
      RET();
    }
  }

  void Run()
  {
    const u64 actual = std::bit_cast<u64 (*)()>(m_code_pointer)();
    constexpr u64 expected = 123;
    EXPECT_EQ(expected, actual);
  }

private:
  const u8* m_code_pointer = nullptr;
};

}  // namespace

TEST(x64ABITest, CallFunction_ZeroParameters)
{
  TestCallFunction test;
  test.Emit([&] { test.ABI_CallFunction(&ZeroParameterFunction); });
  test.Run();
}

TEST(x64ABITest, CallFunction_OneConstantParameter)
{
  TestCallFunction test;
  test.Emit([&] { test.ABI_CallFunction(&OneParameterFunction<u64>, 100); });
  test.Run();
}

TEST(x64ABITest, CallFunction_OneImm64Parameter)
{
  TestCallFunction test;
  test.Emit([&] {
    test.ABI_CallFunction(&OneParameterFunction<u64>, XEmitter::CallFunctionArg(64, Imm64(100)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_OneImm32Parameter)
{
  TestCallFunction test;
  test.Emit([&] {
    test.ABI_CallFunction(&OneParameterFunction<u32>, XEmitter::CallFunctionArg(32, Imm32(100)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_OneImm16Parameter)
{
  TestCallFunction test;
  test.Emit([&] {
    test.ABI_CallFunction(&OneParameterFunction<u16>, XEmitter::CallFunctionArg(16, Imm16(100)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_OneImm8Parameter)
{
  TestCallFunction test;
  test.Emit([&] {
    test.ABI_CallFunction(&OneParameterFunction<u8>, XEmitter::CallFunctionArg(8, Imm8(100)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One64BitRegisterParameterNoMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), Imm64(100));
    test.ABI_CallFunction(&OneParameterFunction<u64>, XEmitter::CallFunctionArg(64, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One64BitRegisterParameterMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM2), Imm64(100));
    test.ABI_CallFunction(&OneParameterFunction<u64>, XEmitter::CallFunctionArg(64, ABI_PARAM2));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One32BitRegisterParameterNoMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), Imm64(0x0101010100000064));
    test.ABI_CallFunction(&OneParameterFunction<u32>, XEmitter::CallFunctionArg(32, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One32BitRegisterParameterMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM2), Imm64(0x0101010100000064));
    test.ABI_CallFunction(&OneParameterFunction<u32>, XEmitter::CallFunctionArg(32, ABI_PARAM2));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One16BitRegisterParameterNoMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), Imm64(0x0101010101010064));
    test.ABI_CallFunction(&OneParameterFunction<u16>, XEmitter::CallFunctionArg(16, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One16BitRegisterParameterMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM2), Imm64(0x0101010101010064));
    test.ABI_CallFunction(&OneParameterFunction<u16>, XEmitter::CallFunctionArg(16, ABI_PARAM2));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One8BitRegisterParameterNoMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), Imm64(0x0101010101010164));
    test.ABI_CallFunction(&OneParameterFunction<u8>, XEmitter::CallFunctionArg(8, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One8BitRegisterParameterMov)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM2), Imm64(0x0101010101010164));
    test.ABI_CallFunction(&OneParameterFunction<u8>, XEmitter::CallFunctionArg(8, ABI_PARAM2));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One64BitMemoryParameter)
{
  constexpr u64 value = 100;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(RAX), ImmPtr(&value));
    test.ABI_CallFunction(&OneParameterFunction<u64>, XEmitter::CallFunctionArg(64, MatR(RAX)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One32BitMemoryParameter)
{
  constexpr u32 value = 100;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM3), ImmPtr(&value));
    test.ABI_CallFunction(&OneParameterFunction<u32>,
                          XEmitter::CallFunctionArg(32, MatR(ABI_PARAM3)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One16BitMemoryParameter)
{
  constexpr u16 value = 100;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM2), ImmPtr(&value));
    test.ABI_CallFunction(&OneParameterFunction<u16>,
                          XEmitter::CallFunctionArg(16, MatR(ABI_PARAM2)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_One8BitMemoryParameter)
{
  constexpr u8 value = 100;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), ImmPtr(&value));
    test.ABI_CallFunction(&OneParameterFunction<u8>,
                          XEmitter::CallFunctionArg(8, MatR(ABI_PARAM1)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_TwoRegistersMixed)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(32, R(ABI_PARAM1), Imm32(20));
    test.ABI_CallFunction(&TwoParameterFunction, 10, XEmitter::CallFunctionArg(32, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_TwoRegistersCycle)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(32, R(ABI_PARAM1), Imm32(20));
    test.MOV(32, R(ABI_PARAM2), Imm32(10));
    test.ABI_CallFunction(&TwoParameterFunction, XEmitter::CallFunctionArg(32, ABI_PARAM2),
                          XEmitter::CallFunctionArg(32, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_ThreeRegistersMixed)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(32, R(ABI_PARAM2), Imm32(10));
    test.MOV(32, R(ABI_PARAM3), Imm32(20));
    test.ABI_CallFunction(&ThreeParameterFunction, XEmitter::CallFunctionArg(32, ABI_PARAM2),
                          XEmitter::CallFunctionArg(32, ABI_PARAM3), 30);
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_ThreeRegistersCycle1)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(32, R(ABI_PARAM1), Imm32(30));
    test.MOV(32, R(ABI_PARAM2), Imm32(10));
    test.MOV(32, R(ABI_PARAM3), Imm32(20));
    test.ABI_CallFunction(&ThreeParameterFunction, XEmitter::CallFunctionArg(32, ABI_PARAM2),
                          XEmitter::CallFunctionArg(32, ABI_PARAM3),
                          XEmitter::CallFunctionArg(32, ABI_PARAM1));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_ThreeRegistersCycle2)
{
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(32, R(ABI_PARAM1), Imm32(20));
    test.MOV(32, R(ABI_PARAM2), Imm32(30));
    test.MOV(32, R(ABI_PARAM3), Imm32(10));
    test.ABI_CallFunction(&ThreeParameterFunction, XEmitter::CallFunctionArg(32, ABI_PARAM3),
                          XEmitter::CallFunctionArg(32, ABI_PARAM1),
                          XEmitter::CallFunctionArg(32, ABI_PARAM2));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_ThreeMemoryRegistersCycle1)
{
  constexpr u32 value1 = 10;
  constexpr u32 value2 = 20;
  constexpr u32 value3 = 30;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), ImmPtr(&value3));
    test.MOV(64, R(ABI_PARAM2), ImmPtr(&value1));
    test.MOV(64, R(ABI_PARAM3), ImmPtr(&value2));
    test.ABI_CallFunction(&ThreeParameterFunction, XEmitter::CallFunctionArg(32, MatR(ABI_PARAM2)),
                          XEmitter::CallFunctionArg(32, MatR(ABI_PARAM3)),
                          XEmitter::CallFunctionArg(32, MatR(ABI_PARAM1)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_ThreeMemoryRegistersCycle2)
{
  constexpr u32 value1 = 10;
  constexpr u32 value2 = 20;
  constexpr u32 value3 = 30;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), ImmPtr(&value2));
    test.MOV(64, R(ABI_PARAM2), ImmPtr(&value3));
    test.MOV(64, R(ABI_PARAM3), ImmPtr(&value1));
    test.ABI_CallFunction(&ThreeParameterFunction, XEmitter::CallFunctionArg(32, MatR(ABI_PARAM3)),
                          XEmitter::CallFunctionArg(32, MatR(ABI_PARAM1)),
                          XEmitter::CallFunctionArg(32, MatR(ABI_PARAM2)));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_FourRegistersMixed)
{
  // Loosely based on an actual piece of code where we call Core::BranchWatch::HitVirtualTrue_fk_n
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(RAX), Imm64(0x0102030405060708));
    test.MOV(32, R(RDX), Imm32(0x15161718));
    test.ABI_CallFunction(&FourParameterFunction, XEmitter::CallFunctionArg(64, RAX),
                          0x090a0b0c0d0e0f10, 0x11121314, XEmitter::CallFunctionArg(32, RDX));
  });
  test.Run();
}

TEST(x64ABITest, CallFunction_FourRegistersComplexAddressing)
{
  constexpr u64 value1 = 0x0102030405060708;
  constexpr u64 value2 = 0x090a0b0c0d0e0f10;
  constexpr u32 value3 = 0x11121314;
  constexpr u32 value4 = 0x15161718;
  TestCallFunction test;
  test.Emit([&] {
    test.MOV(64, R(ABI_PARAM1), Imm32(3));
    test.MOV(64, R(ABI_PARAM2), Imm64(reinterpret_cast<uintptr_t>(&value1) - 3));
    test.MOV(64, R(ABI_PARAM4), Imm64(reinterpret_cast<uintptr_t>(&value2) / 4 + 3));
    test.MOV(64, R(ABI_PARAM3), Imm64(reinterpret_cast<uintptr_t>(&value3) - 30));
    test.MOV(64, R(RAX), Imm32(15));
    static_assert(std::count(ABI_PARAMS.begin(), ABI_PARAMS.end(), RAX) == 0);
    test.ABI_CallFunction(
        &FourParameterFunction, XEmitter::CallFunctionArg(64, MRegSum(ABI_PARAM1, ABI_PARAM2)),
        XEmitter::CallFunctionArg(64, MScaled(ABI_PARAM4, SCALE_4, -12)),
        XEmitter::CallFunctionArg(32, MComplex(ABI_PARAM3, ABI_PARAM1, SCALE_8, 6)),
        XEmitter::CallFunctionArg(32, MComplex(ABI_PARAM3, RAX, SCALE_2,
                                               reinterpret_cast<uintptr_t>(&value4) -
                                                   reinterpret_cast<uintptr_t>(&value3))));
  });
  test.Run();
}