#include "unity.h"

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <span>
#include <string>
#include <string_view>
#include <vector>

#include <CoreFoundation/CoreFoundation.h>
#include <Security/Security.h>

#include "auth/crypto.hpp"

void setUp() {}
void tearDown() {}

namespace {

using u8 = std::uint8_t;

auto hex(std::string_view s) -> std::vector<u8>
{
    std::vector<u8> out;
    for (std::size_t i = 0; i + 1 < s.size(); i += 2) {
        char buf[3] = {s[i], s[i + 1], 0};
        out.push_back(static_cast<u8>(std::strtoul(buf, nullptr, 16)));
    }
    return out;
}

template <std::size_t N>
auto to_array(std::vector<u8> const& v) -> std::array<u8, N>
{
    std::array<u8, N> a{};
    TEST_ASSERT(v.size() == N);
    std::copy(v.begin(), v.end(), a.begin());
    return a;
}

auto span(std::vector<u8> const& v) -> std::span<u8 const>
{
    return {v.data(), v.size()};
}

} // namespace

// --------------------------------------------------------------------------
// SHA-256
// --------------------------------------------------------------------------
void test_sha256_abc(void)
{
    std::array<u8, 32> out{};
    xone::auth::sha256 h;
    std::string data = "abc";
    h.update({reinterpret_cast<u8 const*>(data.data()), data.size()});
    h.finalize(out);
    auto expect = hex("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect.data(), out.data(), 32);
}

void test_sha256_snapshot(void)
{
    // snapshot must not disturb continued hashing
    std::array<u8, 32> out{};
    xone::auth::sha256 h;
    std::string a = "abc";
    h.update({reinterpret_cast<u8 const*>(a.data()), a.size()});

    xone::auth::sha256 snap = h;  // snapshot
    std::string b = "def";
    h.update({reinterpret_cast<u8 const*>(b.data()), b.size()});
    h.finalize(out);
    auto expect = hex("bef57ec7f53a6d40beb640a780a639c83bc29ac8a9816f1fc6c5c6dcd93c4721");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect.data(), out.data(), 32);

    // snapshot still hashes only "abc"
    std::array<u8, 32> snap_out{};
    snap.finalize(snap_out);
    auto expect_snap = hex("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect_snap.data(), snap_out.data(), 32);
}

// --------------------------------------------------------------------------
// HMAC-SHA256 (RFC 4231)
// --------------------------------------------------------------------------
void test_hmac_rfc4231_case2(void)
{
    // key = "Jefe", data = "what do ya want for nothing?"
    std::string key = "Jefe";
    std::string data = "what do ya want for nothing?";
    xone::auth::hmac_sha256 h{{reinterpret_cast<u8 const*>(key.data()), key.size()}};
    h.update({reinterpret_cast<u8 const*>(data.data()), data.size()});
    std::array<u8, 32> out{};
    h.finalize(out);
    auto expect = hex("5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect.data(), out.data(), 32);
}

void test_hmac_rfc4231_case1(void)
{
    // key = 0x0b * 20, data = "Hi There"
    std::vector<u8> key(20, 0x0b);
    std::string data = "Hi There";
    xone::auth::hmac_sha256 h{span(key)};
    h.update({reinterpret_cast<u8 const*>(data.data()), data.size()});
    std::array<u8, 32> out{};
    h.finalize(out);
    auto expect = hex("b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect.data(), out.data(), 32);
}

// --------------------------------------------------------------------------
// PRF (TLS P_SHA256)
// --------------------------------------------------------------------------
void test_prf_expansion(void)
{
    // Long output forces multiple blocks; verify against Python-computed value.
    std::vector<u8> key = hex("00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff");
    std::vector<u8> seed = hex("deadbeefcafebabe");
    std::array<u8, 48> out{};
    xone::auth::prf_sha256(span(key), "Master Secret", span(seed), out);
    auto expect = hex(
        "78b92e5b14f8b98dc9f5ddd33668d3dd19e74467f38992009d71bdfd69b5dfa2"
        "8ff26eaad240df9fe705798d6af784d8");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect.data(), out.data(), 48);
}

// --------------------------------------------------------------------------
// ECDH P-256
// --------------------------------------------------------------------------
void test_ec_small_multiples(void)
{
    auto two = hex("0000000000000000000000000000000000000000000000000000000000000002");
    auto three = hex("0000000000000000000000000000000000000000000000000000000000000003");

    xone::auth::ec_point q{};
    xone::auth::ec_base_point_multiply(to_array<32>(two), q);
    auto expect2g = hex(
        "7cf27b188d034f7e8a52380304b51ac3c08969e277f21b35a60b48fc47669978"
        "07775510db8ed040293d9ac69f7430dbba7dade63ce982299e04b79d227873d1");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect2g.data(), q.data(), 64);

    xone::auth::ec_base_point_multiply(to_array<32>(three), q);
    auto expect3g = hex(
        "5ecbe4d1a6330a44c8f7ef951d4bf165e6c6b721efada985fb41661bc6e7fd6c"
        "8734640c4998ff7e374b06ce1a64a2ecd82ab036384fb83d9a79b127a27d5032");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expect3g.data(), q.data(), 64);
}

void test_ec_base_multiply(void)
{
    // vectors generated with OpenSSL (see test comments in commit)
    auto d = hex("ca9aca63d14014baf4e37bcaeb90317557728b479ca42e126fcb9ae9a4f7f765");
    auto puba = hex(
        "dca2dafb3baa46602e7512de4690ede383e390148b3c041b34e1beae2b9ec997"
        "abf83c388bb7b53cd086a8d20810b22ac1f5d92484abf761f53b6725988604cd");

    xone::auth::ec_point q{};
    xone::auth::ec_base_point_multiply(to_array<32>(d), q);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(puba.data(), q.data(), 64);
}

void test_ec_shared_secret(void)
{
    auto d = hex("ca9aca63d14014baf4e37bcaeb90317557728b479ca42e126fcb9ae9a4f7f765");
    auto pubb = hex(
        "5baeca87ad9b623f7d0cd33af316e5f57d17b711e821acc449a12e5427c95aa3"
        "a0b86826cf4f028cb90bfd284e99db36cc19c640a914b97ef35dd965e2af19da");
    auto secret = hex("df01a8c296d70b09a91bc74f4bd25d877cffa0fe43287d24939e44cdc2e9ec25");

    xone::auth::ec_scalar shared{};
    TEST_ASSERT_TRUE(
        xone::auth::ec_compute_shared(to_array<32>(d), to_array<64>(pubb), shared));
    TEST_ASSERT_EQUAL_UINT8_ARRAY(secret.data(), shared.data(), 32);
}

void test_ec_rejects_off_curve_point(void)
{
    auto d = hex("ca9aca63d14014baf4e37bcaeb90317557728b479ca42e126fcb9ae9a4f7f765");
    // point with x=1, y=1 is not on the curve
    xone::auth::ec_point bad{};
    bad[31] = 1;
    bad[63] = 1;

    xone::auth::ec_scalar shared{};
    TEST_ASSERT_FALSE(
        xone::auth::ec_compute_shared(to_array<32>(d), bad, shared));
}

// --------------------------------------------------------------------------
// RSA (PKCS#1 v1.5) round-trip against Security.framework
// --------------------------------------------------------------------------
void test_rsa_roundtrip(void)
{
    CFErrorRef error = nullptr;
    CFTypeRef attr_keys[] = {kSecAttrKeyType, kSecAttrKeySizeInBits,
                             kSecAttrIsPermanent};
    CFNumberRef size = CFNumberCreate(kCFAllocatorDefault, kCFNumberIntType,
                                      (int[]){2048});
    CFTypeRef attr_vals[] = {kSecAttrKeyTypeRSA, size, kCFBooleanFalse};
    CFDictionaryRef attrs = CFDictionaryCreate(
        kCFAllocatorDefault, attr_keys, attr_vals, 3,
        &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
    TEST_ASSERT_NOT_NULL(attrs);

    SecKeyRef priv = SecKeyCreateRandomKey(attrs, &error);
    CFRelease(size);
    TEST_ASSERT_NOT_NULL(priv);
    if (!priv) {
        CFRelease(attrs);
        return;
    }

    SecKeyRef pub = SecKeyCopyPublicKey(priv);
    TEST_ASSERT_NOT_NULL(pub);
    if (!pub) {
        CFRelease(priv);
        CFRelease(attrs);
        return;
    }

    CFDataRef ext = SecKeyCopyExternalRepresentation(pub, &error);
    TEST_ASSERT_NOT_NULL(ext);
    if (!ext) {
        CFRelease(pub);
        CFRelease(priv);
        CFRelease(attrs);
        return;
    }

    CFIndex key_len = CFDataGetLength(ext);
    std::vector<u8> der(key_len);
    CFDataGetBytes(ext, CFRangeMake(0, key_len), der.data());
    CFRelease(ext);

    // 2048-bit RSA PKCS#1 RSAPublicKey: 4-byte ASN.1 header + 266-byte body.
    TEST_ASSERT_EQUAL_INT(270, key_len);
    TEST_ASSERT_EQUAL_UINT8(0x30, der[0]);
    TEST_ASSERT_EQUAL_UINT8(0x82, der[1]);

    std::vector<u8> plaintext(48, 0x5a);
    std::vector<u8> ciphertext(256);
    TEST_ASSERT_TRUE(xone::auth::rsa_encrypt_pkcs1(span(der), span(plaintext),
                                                  ciphertext));

    // Decrypt to prove the round-trip.
    CFDataRef ct = CFDataCreate(kCFAllocatorDefault, ciphertext.data(),
                                static_cast<CFIndex>(ciphertext.size()));
    CFDataRef dec = SecKeyCreateDecryptedData(
        priv, kSecKeyAlgorithmRSAEncryptionPKCS1, ct, &error);
    CFRelease(ct);
    TEST_ASSERT_NOT_NULL(dec);
    if (dec) {
        CFIndex dec_len = CFDataGetLength(dec);
        TEST_ASSERT_EQUAL_INT(48, dec_len);
        std::vector<u8> dec_bytes(dec_len);
        CFDataGetBytes(dec, CFRangeMake(0, dec_len), dec_bytes.data());
        TEST_ASSERT_EQUAL_UINT8_ARRAY(plaintext.data(), dec_bytes.data(), 48);
        CFRelease(dec);
    }

    CFRelease(priv);
    CFRelease(pub);
    CFRelease(attrs);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_sha256_abc);
    RUN_TEST(test_sha256_snapshot);
    RUN_TEST(test_hmac_rfc4231_case2);
    RUN_TEST(test_hmac_rfc4231_case1);
    RUN_TEST(test_prf_expansion);
    RUN_TEST(test_ec_small_multiples);
    RUN_TEST(test_ec_base_multiply);
    RUN_TEST(test_ec_shared_secret);
    RUN_TEST(test_ec_rejects_off_curve_point);
    RUN_TEST(test_rsa_roundtrip);
    return UNITY_END();
}