#include <unity.h>

#include <string>

#include "ipv4.h"
#include "net_probe.h"

using namespace roro::net;

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

void test_what_was_typed() {
    PingArgs p;
    TEST_ASSERT_EQUAL_STRING("", parsePing("example.org", p).c_str());
    TEST_ASSERT_EQUAL_STRING("example.org", p.host.c_str());
    TEST_ASSERT_EQUAL_INT(4, p.count);
    TEST_ASSERT_EQUAL_INT(56, p.size);
    TEST_ASSERT_EQUAL_STRING("", parsePing("  10.9.0.1   10  1372 ", p).c_str());
    TEST_ASSERT_EQUAL_STRING("10.9.0.1", p.host.c_str());
    TEST_ASSERT_EQUAL_INT(10, p.count);
    TEST_ASSERT_EQUAL_INT(1372, p.size);
    TEST_ASSERT_EQUAL_STRING("ping <host> [count] [size]", parsePing("", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a count from 1 to 100", parsePing("a.example 0", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a count from 1 to 100", parsePing("a.example lots", p).c_str());
    TEST_ASSERT_EQUAL_STRING("a size from 0 to 1400 bytes", parsePing("a.example 4 9000", p).c_str());
    TEST_ASSERT_EQUAL_INT(10, p.count);  // a refused line changes nothing

    PortArgs t;
    TEST_ASSERT_EQUAL_STRING("", parsePort("irc.libera.chat 6697", t).c_str());
    TEST_ASSERT_EQUAL_STRING("irc.libera.chat", t.host.c_str());
    TEST_ASSERT_EQUAL_UINT16(6697, t.port);
    TEST_ASSERT_EQUAL_STRING("", parsePort("10.9.0.1:2323", t).c_str());
    TEST_ASSERT_EQUAL_UINT16(2323, t.port);
    TEST_ASSERT_EQUAL_STRING("port <host> <port>", parsePort("10.9.0.1", t).c_str());
    TEST_ASSERT_EQUAL_STRING("a port from 1 to 65535", parsePort("10.9.0.1 70000", t).c_str());

    LookupArgs l;
    TEST_ASSERT_EQUAL_STRING("", parseLookup("roro9stack.net", l).c_str());
    TEST_ASSERT_EQUAL_STRING("", l.server.c_str());
    TEST_ASSERT_EQUAL_STRING("", parseLookup("roro9stack.net 9.9.9.9", l).c_str());
    TEST_ASSERT_EQUAL_STRING("9.9.9.9", l.server.c_str());
    TEST_ASSERT_EQUAL_STRING("the server as an address: 9.9.9.9", parseLookup("roro9stack.net dns.quad9.net", l).c_str());
}

void test_an_echo_request_and_its_answers() {
    uint8_t echo[64];
    size_t len = buildEcho(echo, sizeof echo, 0xBEEF, 7, 32);
    TEST_ASSERT_EQUAL_size_t(40, len);
    TEST_ASSERT_EQUAL_UINT8(8, echo[0]);
    TEST_ASSERT_EQUAL_UINT16(0, inetChecksum(echo, len));  // a packet with its checksum in sums to nothing
    TEST_ASSERT_EQUAL_size_t(0, buildEcho(echo, 20, 1, 1, 32));

    // The reply, as a raw socket hands it: an IP header, then the same with type 0.
    uint8_t reply[20 + 40] = {0x45, 0, 0, 60, 0, 0, 0, 0, 64, 1};
    std::copy(echo, echo + len, reply + 20);
    reply[20] = 0;
    IcmpAnswer a = parseIcmp(reply, sizeof reply);
    TEST_ASSERT_TRUE(a.kind == IcmpAnswer::Kind::Echo);
    TEST_ASSERT_EQUAL_HEX16(0xBEEF, a.id);
    TEST_ASSERT_EQUAL_UINT16(7, a.seq);

    // A router on the way: "time exceeded", with the start of our packet inside it.
    uint8_t exceeded[20 + 8 + 20 + 8] = {0x45, 0, 0, 56, 0, 0, 0, 0, 250, 1};
    exceeded[20] = 11;
    uint8_t* inner = exceeded + 28;
    inner[0] = 0x45;
    inner[9] = 1;
    std::copy(echo, echo + 8, inner + 20);
    a = parseIcmp(exceeded, sizeof exceeded);
    TEST_ASSERT_TRUE(a.kind == IcmpAnswer::Kind::TimeExceeded);
    TEST_ASSERT_EQUAL_HEX16(0xBEEF, a.id);
    TEST_ASSERT_EQUAL_UINT16(7, a.seq);
    exceeded[20] = 3;
    TEST_ASSERT_TRUE(parseIcmp(exceeded, sizeof exceeded).kind == IcmpAnswer::Kind::Unreachable);

    // Not ours to read: someone else's ping to us, a cut packet, UDP.
    TEST_ASSERT_TRUE(parseIcmp(reply, 22).kind == IcmpAnswer::Kind::Other);
    reply[20] = 8;
    TEST_ASSERT_TRUE(parseIcmp(reply, sizeof reply).kind == IcmpAnswer::Kind::Other);
    reply[20] = 0;
    reply[9] = 17;
    TEST_ASSERT_TRUE(parseIcmp(reply, sizeof reply).kind == IcmpAnswer::Kind::Other);
    TEST_ASSERT_TRUE(parseIcmp(exceeded, 40).kind == IcmpAnswer::Kind::Other);
}

void test_the_summary() {
    PingStats s;
    s.sent = 4;
    TEST_ASSERT_EQUAL_STRING("0/4 back, 100% lost", s.summary().c_str());
    s.add(23);
    s.add(12);
    s.add(41);
    TEST_ASSERT_EQUAL_STRING("3/4 back, 25% lost, 12/25/41 ms", s.summary().c_str());
}

void test_a_dns_query() {
    uint8_t q[64];
    size_t len = buildDnsQuery(q, sizeof q, 0x1234, "roro9stack.net");
    const uint8_t want[] = {0x12, 0x34, 0x01, 0x00, 0, 1, 0, 0, 0, 0, 0, 0, 10, 'r', 'o', 'r', 'o', '9', 's', 't', 'a', 'c', 'k', 3, 'n', 'e', 't', 0, 0, 1, 0, 1};
    TEST_ASSERT_EQUAL_size_t(sizeof want, len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(want, q, sizeof want);
    TEST_ASSERT_EQUAL_size_t(len, buildDnsQuery(q, sizeof q, 0x1234, "roro9stack.net."));  // a final dot is the same name
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, ""));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, "a..b"));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, sizeof q, 1, std::string(64, 'a') + ".net"));
    TEST_ASSERT_EQUAL_size_t(0, buildDnsQuery(q, 20, 1, "roro9stack.net"));
}

void test_a_dns_answer() {
    // www.example.org is an alias of example.org, which has two addresses. Names after the first
    // are pointers back into the message, as servers send them.
    const uint8_t msg[] = {
        0x12, 0x34, 0x81, 0x80, 0, 1, 0, 3, 0, 0, 0, 0,
        3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'o', 'r', 'g', 0, 0, 1, 0, 1,  // the question, at 12
        0xC0, 12, 0, 5, 0, 1, 0, 0, 1, 0, 0, 2, 0xC0, 16,                                       // CNAME -> example.org (pointer to 16)
        0xC0, 16, 0, 1, 0, 1, 0, 0, 1, 0, 0, 4, 93, 184, 216, 34,
        0xC0, 16, 0, 1, 0, 1, 0, 0, 1, 0, 0, 4, 93, 184, 216, 35,
    };
    DnsAnswer a;
    TEST_ASSERT_TRUE(parseDnsAnswer(msg, sizeof msg, 0x1234, a));
    TEST_ASSERT_EQUAL_INT(0, a.rcode);
    TEST_ASSERT_FALSE(a.truncated);
    TEST_ASSERT_EQUAL_size_t(2, a.addresses.size());
    TEST_ASSERT_EQUAL_STRING("93.184.216.34", formatIpv4(a.addresses[0]).c_str());
    TEST_ASSERT_EQUAL_STRING("93.184.216.35", formatIpv4(a.addresses[1]).c_str());
    TEST_ASSERT_EQUAL_STRING("example.org", a.alias.c_str());

    TEST_ASSERT_FALSE(parseDnsAnswer(msg, sizeof msg, 0x9999, a));  // someone else's answer
    for (size_t cut = 0; cut < sizeof msg; cut++) parseDnsAnswer(msg, cut, 0x1234, a);  // cut anywhere: no crash
    TEST_ASSERT_FALSE(parseDnsAnswer(msg, sizeof msg - 3, 0x1234, a));

    uint8_t none[sizeof msg];
    std::copy(msg, msg + sizeof msg, none);
    none[3] = 0x83;  // no such name
    none[7] = 0;
    TEST_ASSERT_TRUE(parseDnsAnswer(none, 33, 0x1234, a));
    TEST_ASSERT_EQUAL_INT(3, a.rcode);
    TEST_ASSERT_EQUAL_size_t(0, a.addresses.size());

    // A pointer that points at itself must not hang the reader.
    uint8_t loop[] = {0x12, 0x34, 0x81, 0x80, 0, 1, 0, 0, 0, 0, 0, 0, 0xC0, 12, 0, 1, 0, 1};
    TEST_ASSERT_FALSE(parseDnsAnswer(loop, sizeof loop, 0x1234, a));
}

void test_ntp() {
    uint8_t req[kNtpPacket];
    buildNtpRequest(req);
    TEST_ASSERT_EQUAL_HEX8(0x23, req[0]);
    TEST_ASSERT_EQUAL_UINT8(0, req[47]);
    // A server's answer: stratum 2, transmit time 2026-10-08 00:45:12.5 UTC.
    uint8_t ans[kNtpPacket] = {0x24, 2};
    uint32_t secs = 1791420312u + 2208988800u;
    ans[40] = static_cast<uint8_t>(secs >> 24);
    ans[41] = static_cast<uint8_t>(secs >> 16);
    ans[42] = static_cast<uint8_t>(secs >> 8);
    ans[43] = static_cast<uint8_t>(secs);
    ans[44] = 0x80;  // half a second
    NtpAnswer a;
    TEST_ASSERT_TRUE(parseNtpAnswer(ans, sizeof ans, a));
    TEST_ASSERT_EQUAL_INT(2, a.stratum);
    TEST_ASSERT_TRUE(a.seconds == 1791420312);
    TEST_ASSERT_EQUAL_UINT32(500, a.millis);
    TEST_ASSERT_FALSE(parseNtpAnswer(ans, 40, a));  // cut short
    ans[1] = 0;
    TEST_ASSERT_FALSE(parseNtpAnswer(ans, sizeof ans, a));  // "go away"
    ans[1] = 2;
    ans[0] = 0x23;
    TEST_ASSERT_FALSE(parseNtpAnswer(ans, sizeof ans, a));  // a client's packet, not a server's
    // After 2036 the count starts again from zero.
    ans[0] = 0x24;
    ans[40] = ans[41] = ans[42] = 0;
    ans[43] = 10;
    TEST_ASSERT_TRUE(parseNtpAnswer(ans, sizeof ans, a));
    TEST_ASSERT_TRUE(a.seconds == 4294967296LL + 10 - 2208988800LL);

    TEST_ASSERT_EQUAL_STRING("right, to 0.1 s", clockOffset(1000000, 1000040).c_str());
    TEST_ASSERT_EQUAL_STRING("0.3 s ahead", clockOffset(1000300, 1000000).c_str());
    TEST_ASSERT_EQUAL_STRING("2.5 s behind", clockOffset(1000000, 1002500).c_str());
    TEST_ASSERT_EQUAL_STRING("45 s behind", clockOffset(0, 45000).c_str());
    TEST_ASSERT_EQUAL_STRING("3 min ahead", clockOffset(180000, 0).c_str());
    TEST_ASSERT_EQUAL_STRING("5 h behind", clockOffset(0, 18000000).c_str());
    TEST_ASSERT_EQUAL_STRING("20552 days behind", clockOffset(0, 1775700000000LL).c_str());  // a clock still in 1970
}

void test_certificates_and_ports() {
    TEST_ASSERT_EQUAL_STRING("R11", certName("C=US, O=Let's Encrypt, CN=R11").c_str());
    TEST_ASSERT_EQUAL_STRING("git.twis.la", certName("CN=git.twis.la").c_str());
    TEST_ASSERT_EQUAL_STRING("Example Org", certName("C=BE, O=Example Org").c_str());
    TEST_ASSERT_EQUAL_STRING("C=BE", certName("C=BE").c_str());
    TEST_ASSERT_EQUAL_STRING("x", certName("O=Not this, DCN=nor this, CN=x").c_str());
    TEST_ASSERT_EQUAL_STRING("*.badssl.com", certName("OU=PositiveSSL Wildcard, CN=#140C2A2E62616473736C2E636F6D").c_str());  // an old kind of string
    TEST_ASSERT_EQUAL_STRING("#14zz", certName("CN=#14zz").c_str());
    TEST_ASSERT_EQUAL_INT(1, daysBetween(2026, 10, 7, 2026, 10, 8));
    TEST_ASSERT_EQUAL_INT(53, daysBetween(2026, 10, 8, 2026, 11, 30));
    TEST_ASSERT_EQUAL_INT(-8, daysBetween(2026, 10, 8, 2026, 9, 30));
    TEST_ASSERT_EQUAL_INT(366, daysBetween(2028, 1, 1, 2029, 1, 1));  // a leap year
    TEST_ASSERT_EQUAL_INT(365, daysBetween(2100, 1, 1, 2101, 1, 1));  // and a year that isn't one
    TEST_ASSERT_EQUAL_STRING("updates", portLabel(3232, true));
    TEST_ASSERT_EQUAL_STRING("Debug Console", portLabel(2323, true));
    TEST_ASSERT_EQUAL_STRING("sharing", portLabel(80, true));
    TEST_ASSERT_EQUAL_STRING("", portLabel(80, false));
    TEST_ASSERT_EQUAL_STRING("DHCP", portLabel(68, false));
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_what_was_typed);
    RUN_TEST(test_an_echo_request_and_its_answers);
    RUN_TEST(test_the_summary);
    RUN_TEST(test_a_dns_query);
    RUN_TEST(test_a_dns_answer);
    RUN_TEST(test_ntp);
    RUN_TEST(test_certificates_and_ports);
    return UNITY_END();
}