#include "third_party/catch/include/catch.hpp"
#include "xenia/kernel/xboxkrnl/xecrypt_rsa.h"

namespace {

// Helper to build an XECRYPT_RSA-compatible blob in memory.
// Layout: [4 bytes size BE][4 bytes exponent BE][8 bytes pad][modulus limbs...]
std::vector<uint8_t> build_rsa_blob(uint32_t num_qwords, uint32_t exponent,
                                    const uint8_t* modulus_limbs) {
  std::vector<uint8_t> blob(0x10 + num_qwords * 8);
  // size (big-endian)
  blob[0] = (num_qwords >> 24) & 0xFF;
  blob[1] = (num_qwords >> 16) & 0xFF;
  blob[2] = (num_qwords >> 8) & 0xFF;
  blob[3] = num_qwords & 0xFF;
  // exponent (big-endian)
  blob[4] = (exponent >> 24) & 0xFF;
  blob[5] = (exponent >> 16) & 0xFF;
  blob[6] = (exponent >> 8) & 0xFF;
  blob[7] = exponent & 0xFF;
  // pad = 0
  // modulus data
  std::memcpy(blob.data() + 0x10, modulus_limbs, num_qwords * 8);
  return blob;
}

// Pirs public key modulus (32 qwords = 2048 bits), exponent = 3
static const uint8_t pirs_modulus[] = {
    0xE6, 0x3B, 0x32, 0xB2, 0x8D, 0x9E, 0x9E, 0xE7, 0x9D, 0xFC, 0x5C, 0x72,
    0x41, 0x94, 0x58, 0x47, 0xDE, 0x0D, 0x18, 0x40, 0x72, 0xD6, 0xE3, 0x46,
    0x8E, 0xBA, 0x8E, 0xBC, 0x1A, 0x90, 0xAC, 0x20, 0xBA, 0x03, 0x85, 0xB5,
    0x1A, 0x3E, 0x25, 0xF9, 0xA6, 0x58, 0xEB, 0xB6, 0xA3, 0xC4, 0xA3, 0xEE,
    0xB2, 0xB0, 0xAE, 0x97, 0x69, 0xEB, 0xFE, 0x71, 0xFC, 0x02, 0xAB, 0x77,
    0xBA, 0xC8, 0xE6, 0x74, 0xE6, 0x7C, 0x63, 0x0E, 0xAF, 0x4C, 0xF7, 0xE7,
    0x11, 0x4A, 0x80, 0x24, 0x72, 0x05, 0x7A, 0x63, 0xD0, 0xF8, 0x91, 0x02,
    0xA6, 0xE7, 0x7D, 0x77, 0xC5, 0xA7, 0x9B, 0x08, 0x11, 0x2E, 0xA0, 0x64,
    0x45, 0x60, 0x46, 0xBC, 0x36, 0xE1, 0x17, 0x71, 0xBE, 0x66, 0x49, 0x2F,
    0xAE, 0x20, 0xA4, 0x76, 0x9C, 0x27, 0x51, 0xCF, 0x4B, 0x34, 0x7A, 0x35,
    0xBC, 0xA4, 0xAA, 0x1C, 0x47, 0x4B, 0xF4, 0x97, 0x22, 0x4E, 0x13, 0x24,
    0xD3, 0xC1, 0x57, 0xDF, 0x4D, 0x84, 0xB9, 0x18, 0x97, 0x99, 0xAC, 0x00,
    0xB3, 0x3D, 0x03, 0x25, 0x60, 0xC8, 0x7A, 0x59, 0xFE, 0x48, 0xFF, 0x28,
    0x3D, 0x10, 0xBB, 0x9E, 0x09, 0x06, 0x2A, 0x61, 0x20, 0x2C, 0xF8, 0x72,
    0xEB, 0x87, 0xE6, 0xD1, 0xFB, 0xB3, 0x66, 0xFC, 0x4A, 0x02, 0xAE, 0xD4,
    0xD8, 0x37, 0xCF, 0xA6, 0x32, 0x25, 0x79, 0x36, 0x0E, 0xF4, 0xED, 0x19,
    0xA2, 0x10, 0x27, 0x96, 0x2F, 0x9F, 0xA9, 0x3D, 0xA4, 0x37, 0x30, 0x11,
    0x51, 0x83, 0xBD, 0xF7, 0xC7, 0xE5, 0xCE, 0xAA, 0xEC, 0xDE, 0x48, 0xA0,
    0x84, 0xF7, 0xB0, 0xF6, 0x4B, 0x8E, 0xF0, 0x89, 0xBD, 0x47, 0x7C, 0x90,
    0xDD, 0x88, 0x12, 0x17, 0x40, 0xD2, 0x4E, 0xA6, 0xC6, 0x11, 0x04, 0x1B,
    0x57, 0xA8, 0x68, 0xB4, 0x61, 0xF4, 0x1B, 0xC6, 0x8B, 0xE8, 0xD9, 0x20,
    0xF2, 0x05, 0xE0, 0x70,
};

}  // namespace

TEST_CASE("XeCryptBnQwNeRsaPubCrypt_rejects_small_keys", "[crypt]") {
  // Keys below 512 bits (8 qwords) must be rejected
  const uint8_t modulus[] = {
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07,
  };
  const uint8_t input[] = {
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
  };
  uint8_t output[8] = {};

  REQUIRE(XeCryptBnQwNeRsaPubCrypt(input, output, modulus, 1, 3) == 0);

  const uint8_t modulus2[] = {
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x11,
  };
  const uint8_t input2[] = {
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  };
  uint8_t output2[16] = {};

  REQUIRE(XeCryptBnQwNeRsaPubCrypt(input2, output2, modulus2, 2, 65537) == 0);
}

TEST_CASE("XeCryptBnQwNeRsaPubCrypt_2048bit_identity", "[crypt]") {
  // 1^3 mod modulus = 1
  uint8_t input[256] = {};
  input[7] = 0x01;

  uint8_t output[256] = {};
  uint32_t ret = XeCryptBnQwNeRsaPubCrypt(input, output, pirs_modulus, 0x20, 3);

  REQUIRE(ret == 1);
  uint8_t expected[256] = {};
  expected[7] = 0x01;
  REQUIRE(std::memcmp(output, expected, 256) == 0);
}

TEST_CASE("XeCryptBnQwNeRsaPubCrypt_2048bit_small_base", "[crypt]") {
  // 2^3 = 8 < modulus, so result is 8
  uint8_t input[256] = {};
  input[7] = 0x02;

  uint8_t output[256] = {};
  uint32_t ret = XeCryptBnQwNeRsaPubCrypt(input, output, pirs_modulus, 0x20, 3);

  REQUIRE(ret == 1);
  uint8_t expected[256] = {};
  expected[7] = 0x08;
  REQUIRE(std::memcmp(output, expected, 256) == 0);
}

TEST_CASE("XeCryptBnQwNeRsaPubCrypt_2048bit_large_base", "[crypt]") {
  // Large base that exercises modular reduction
  uint8_t input[256];
  for (int i = 0; i < 256; i++) {
    input[i] = static_cast<uint8_t>(i + 1);
  }
  // Ensure top limb is smaller than modulus top limb (0xF2...)
  input[248] = 0x01;
  input[249] = 0x02;
  input[250] = 0x03;
  input[251] = 0x04;
  input[252] = 0x05;
  input[253] = 0x06;
  input[254] = 0x07;
  input[255] = 0x08;

  uint8_t output[256] = {};
  uint32_t ret = XeCryptBnQwNeRsaPubCrypt(input, output, pirs_modulus, 0x20, 3);
  REQUIRE(ret == 1);

  // Deterministic: same input produces same output
  uint8_t output2[256] = {};
  uint32_t ret2 =
      XeCryptBnQwNeRsaPubCrypt(input, output2, pirs_modulus, 0x20, 3);
  REQUIRE(ret2 == 1);
  REQUIRE(std::memcmp(output, output2, 256) == 0);

  // Non-trivial result
  bool all_zero = true;
  for (int i = 0; i < 256; i++) {
    if (output[i] != 0) {
      all_zero = false;
      break;
    }
  }
  REQUIRE_FALSE(all_zero);
}

TEST_CASE("XeCryptBnQwNeRsaPubCrypt_2048bit_zero_base", "[crypt]") {
  // 0^e mod m = 0
  uint8_t input[256] = {};
  uint8_t output[256];
  std::memset(output, 0xFF, 256);

  uint32_t ret =
      XeCryptBnQwNeRsaPubCrypt(input, output, pirs_modulus, 0x20, 65537);

  REQUIRE(ret == 1);
  for (int i = 0; i < 256; i++) {
    REQUIRE(output[i] == 0x00);
  }
}
