#include "unity.h"

#include <array>
#include <cstdint>
#include <vector>

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

#include "auth/auth.hpp"

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

namespace {

using u8 = std::uint8_t;
using u16 = std::uint16_t;

class mock_sink : public xone::auth::auth_sink {
public:
    struct sent {
        std::vector<u8> data;
        bool ack;
    };

    std::vector<sent> packets;
    std::vector<u8> key;

    auto send(std::span<u8 const> pkt, bool ack) -> int override
    {
        packets.push_back({std::vector<u8>(pkt.begin(), pkt.end()), ack});
        return 0;
    }

    auto set_encryption_key(std::span<u8 const> k) -> int override
    {
        key.assign(k.begin(), k.end());
        return 0;
    }
};

// Build an auth data packet: full header + payload.
auto make_pkt(u8 cmd, std::vector<u8> const& payload) -> std::vector<u8>
{
    std::vector<u8> p;
    u16 data_len = static_cast<u16>(payload.size());

    p.push_back(0x00);               // context = handshake
    p.push_back(0x00);               // options (client -> host)
    p.push_back(0x00);               // error
    p.push_back(cmd);
    p.push_back(static_cast<u8>(data_len >> 8));
    p.push_back(static_cast<u8>(data_len & 0xFF));

    p.push_back(cmd);                // data.command
    p.push_back(0x01);               // data.version
    u16 data_body = static_cast<u16>(data_len - 4);
    p.push_back(static_cast<u8>(data_body >> 8));
    p.push_back(static_cast<u8>(data_body & 0xFF));

    p.insert(p.end(), payload.begin(), payload.end());
    return p;
}

auto make_rsa_der() -> std::vector<u8>
{
    CFTypeRef keys[] = {kSecAttrKeyType, kSecAttrKeySizeInBits,
                        kSecAttrIsPermanent};
    CFNumberRef size = CFNumberCreate(kCFAllocatorDefault, kCFNumberIntType,
                                      (int[]){2048});
    CFTypeRef vals[] = {kSecAttrKeyTypeRSA, size, kCFBooleanFalse};
    CFDictionaryRef attrs = CFDictionaryCreate(
        kCFAllocatorDefault, keys, vals, 3, &kCFTypeDictionaryKeyCallBacks,
        &kCFTypeDictionaryValueCallBacks);
    CFErrorRef error = nullptr;
    SecKeyRef priv = SecKeyCreateRandomKey(attrs, &error);
    SecKeyRef pub = SecKeyCopyPublicKey(priv);
    CFDataRef ext = SecKeyCopyExternalRepresentation(pub, &error);
    std::vector<u8> der(CFDataGetLength(ext));
    CFDataGetBytes(ext, CFRangeMake(0, CFDataGetLength(ext)), der.data());
    CFRelease(ext);
    CFRelease(pub);
    CFRelease(priv);
    CFRelease(size);
    CFRelease(attrs);
    return der;
}

auto be16_at(std::vector<u8> const& v, std::size_t i) -> u16
{
    return static_cast<u16>((v[i] << 8) | v[i + 1]);
}

} // namespace

void test_start_sends_host_hello(void)
{
    mock_sink sink;
    xone::auth::auth auth{sink};

    TEST_ASSERT_EQUAL_INT(0, auth.start());
    TEST_ASSERT_EQUAL_INT(1, sink.packets.size());
    TEST_ASSERT_TRUE(sink.packets[0].ack);
    TEST_ASSERT_EQUAL_INT(58, sink.packets[0].data.size());

    auto const& p = sink.packets[0].data;
    TEST_ASSERT_EQUAL_UINT8(0x00, p[0]);  // context
    TEST_ASSERT_EQUAL_UINT8(0x41, p[1]);  // ACK | FROM_HOST
    TEST_ASSERT_EQUAL_UINT8(0x01, p[3]);  // host hello
    TEST_ASSERT_EQUAL_INT(44, be16_at(p, 4));
    TEST_ASSERT_EQUAL_UINT8(0x01, p[6]);  // data.command
    TEST_ASSERT_EQUAL_UINT8(0x01, p[7]);  // data.version
    TEST_ASSERT_EQUAL_INT(40, be16_at(p, 8));
}

void test_client_hello_requests_certificate(void)
{
    mock_sink sink;
    xone::auth::auth auth{sink};
    auth.start();
    std::size_t base = sink.packets.size();

    std::vector<u8> payload(80, 0x11);  // client hello: random + unknown
    TEST_ASSERT_EQUAL_INT(0, auth.process_pkt(make_pkt(0x02, payload)));

    TEST_ASSERT_EQUAL_INT(base + 1, sink.packets.size());
    auto const& p = sink.packets.back().data;
    TEST_ASSERT_EQUAL_INT(14, p.size());  // request packet
    TEST_ASSERT_TRUE(sink.packets.back().ack);
    TEST_ASSERT_EQUAL_UINT8(0x42, p[1]);  // REQUEST | FROM_HOST
    TEST_ASSERT_EQUAL_UINT8(0x03, p[3]);  // client certificate
    TEST_ASSERT_EQUAL_INT(1028, be16_at(p, 4));  // cert max len + data hdr
}

void test_certificate_triggers_rsa_exchange(void)
{
    mock_sink sink;
    xone::auth::auth auth{sink};
    auth.start();

    // client hello -> cert request
    auth.process_pkt(make_pkt(0x02, std::vector<u8>(80, 0x22)));
    std::size_t base = sink.packets.size();

    // certificate: some prefix bytes + ASN.1 SEQUENCE + 270-byte DER
    std::vector<u8> der = make_rsa_der();
    TEST_ASSERT_EQUAL_INT(270, der.size());
    TEST_ASSERT_EQUAL_UINT8(0x30, der[0]);
    TEST_ASSERT_EQUAL_UINT8(0x82, der[1]);
    TEST_ASSERT_EQUAL_UINT8(0x01, der[2]);
    TEST_ASSERT_EQUAL_UINT8(0x0a, der[3]);

    std::vector<u8> cert(64, 0x33);
    cert.insert(cert.end(), der.begin(), der.end());
    cert.insert(cert.end(), 32, 0x44);

    TEST_ASSERT_EQUAL_INT(0, auth.process_pkt(make_pkt(0x03, cert)));

    // host secret packet
    TEST_ASSERT_EQUAL_INT(base + 1, sink.packets.size());
    auto const& p = sink.packets.back().data;
    TEST_ASSERT_EQUAL_INT(274, p.size());  // header_full + 256 + trailer
    TEST_ASSERT_TRUE(sink.packets.back().ack);
    TEST_ASSERT_EQUAL_UINT8(0x05, p[3]);  // host secret
}

void test_acknowledge_sends_host_finish(void)
{
    mock_sink sink;
    xone::auth::auth auth{sink};
    auth.start();
    auth.process_pkt(make_pkt(0x02, std::vector<u8>(80, 0x22)));

    std::vector<u8> der = make_rsa_der();
    std::vector<u8> cert(der.begin(), der.end());
    auth.process_pkt(make_pkt(0x03, cert));
    std::size_t base = sink.packets.size();

    // ACK: handshake header with options=ACK, command=0x01
    std::vector<u8> ack{0x00, 0x01, 0x00, 0x01, 0x00, 0x00};
    TEST_ASSERT_EQUAL_INT(0, auth.process_pkt(ack));

    TEST_ASSERT_EQUAL_INT(base + 1, sink.packets.size());
    auto const& p = sink.packets.back().data;
    TEST_ASSERT_EQUAL_INT(50, p.size());  // header_full + 32 + trailer
    TEST_ASSERT_EQUAL_UINT8(0x07, p[3]);  // host finish
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_start_sends_host_hello);
    RUN_TEST(test_client_hello_requests_certificate);
    RUN_TEST(test_certificate_triggers_rsa_exchange);
    RUN_TEST(test_acknowledge_sends_host_finish);
    return UNITY_END();
}