#include "probe.h"

#include <array>
#include <cerrno>
#include <cstring>
#include <future>
#include <stop_token>
#include <thread>

#include <arpa/inet.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

#include <gtest/gtest.h>

namespace
{

using namespace std::chrono_literals;

bool waitReadable(int fd, std::stop_token stop)
{
    while(!stop.stop_requested())
    {
        pollfd descriptor{fd, POLLIN, 0};
        const int ready = poll(&descriptor, 1, 20);
        if(ready > 0)
        {
            return true;
        }
        if(ready < 0 && errno != EINTR)
        {
            return false;
        }
    }
    return false;
}

class ProbeTest : public testing::Test
{
protected:
    void TearDown() override
    {
        worker.request_stop();
        if(worker.joinable())
        {
            worker.join();
        }
        if(server >= 0)
        {
            close(server);
        }
    }

    bool bindServer(int type, int family = AF_INET, bool listening = false)
    {
        server = socket(family, type | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
        if(server < 0)
        {
            return false;
        }
        sockaddr_storage address{};
        socklen_t size = 0;
        if(family == AF_INET)
        {
            auto& ipv4 = reinterpret_cast<sockaddr_in&>(address);
            ipv4.sin_family = AF_INET;
            ipv4.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
            size = sizeof(ipv4);
        }
        else
        {
            auto& ipv6 = reinterpret_cast<sockaddr_in6&>(address);
            ipv6.sin6_family = AF_INET6;
            ipv6.sin6_addr = in6addr_loopback;
            size = sizeof(ipv6);
        }
        if(bind(server, reinterpret_cast<sockaddr*>(&address), size) < 0 ||
           getsockname(server, reinterpret_cast<sockaddr*>(&address),
                       &size) < 0)
        {
            return false;
        }
        port = family == AF_INET
            ? ntohs(reinterpret_cast<sockaddr_in&>(address).sin_port)
            : ntohs(reinterpret_cast<sockaddr_in6&>(address).sin6_port);
        return !listening || listen(server, 8) == 0;
    }

    static void serveHttp(std::stop_token stop, ProbeTest* self,
                          int status, bool silent)
    {
        if(!waitReadable(self->server, stop))
        {
            return;
        }
        const int client = accept4(self->server, nullptr, nullptr,
                                    SOCK_NONBLOCK | SOCK_CLOEXEC);
        if(client < 0)
        {
            return;
        }
        std::string request;
        while(request.find("\r\n\r\n") == std::string::npos &&
              request.size() < 8192 && waitReadable(client, stop))
        {
            std::array<char, 1024> buffer;
            const auto count = recv(client, buffer.data(), buffer.size(), 0);
            if(count <= 0)
            {
                break;
            }
            request.append(buffer.data(), count);
        }
        if(!silent)
        {
            const auto response = "HTTP/1.1 " + std::to_string(status) +
                " Test\r\nContent-Length: 2\r\nConnection: close\r\n"
                "Location: /redirect\r\n\r\nok";
            send(client, response.data(), response.size(), MSG_NOSIGNAL);
        }
        else
        {
            // Wait for the probe to time out and close its connection.
            while(waitReadable(client, stop))
            {
                char buffer;
                if(recv(client, &buffer, 1, 0) <= 0)
                {
                    break;
                }
            }
        }
        close(client);
    }

    static void serveUdp(std::stop_token stop, ProbeTest* self,
                         bool empty_reply)
    {
        if(!waitReadable(self->server, stop))
        {
            return;
        }
        std::array<char, 1024> buffer;
        sockaddr_storage source{};
        socklen_t size = sizeof(source);
        const auto count = recvfrom(self->server, buffer.data(), buffer.size(),
            0, reinterpret_cast<sockaddr*>(&source), &size);
        if(count >= 0)
        {
            self->received_payload.assign(buffer.data(), count);
            const std::string reply = empty_reply ? "" : "reply";
            sendto(self->server, reply.data(), reply.size(), MSG_NOSIGNAL,
                    reinterpret_cast<sockaddr*>(&source), size);
        }
    }

    int server = -1;
    std::uint16_t port = 0;
    std::string received_payload;
    std::jthread worker;
};

TEST(Probe, InvalidConfiguration)
{
    EXPECT_FALSE(createProbe(TcpEndpoint{"", 80}));
    EXPECT_FALSE(createProbe(TcpEndpoint{"localhost", 0}));
    EXPECT_FALSE(createProbe(UdpEndpoint{"localhost", 0, ""}));
    EXPECT_FALSE(createProbe(UdpEndpoint{"localhost", 80,
                                         std::string(65508, 'x')}));
    EXPECT_FALSE(createProbe(IcmpEndpoint{""}));
    EXPECT_FALSE(createProbe(HttpEndpoint{"file:///etc/hosts"}));
    EXPECT_FALSE(createProbe(HttpEndpoint{"http://"}));
    EXPECT_FALSE(createProbe(HttpEndpoint{"http://localhost"}, 0s));
    EXPECT_FALSE(createProbe(TcpEndpoint{"localhost", 80}, -1s));
    EXPECT_FALSE(createProbe(TcpEndpoint{
        "localhost", 80}, std::chrono::seconds::max()));
    EXPECT_FALSE(createProbe(UdpEndpoint{"localhost", 80, ""}, 0s));
    EXPECT_FALSE(createProbe(IcmpEndpoint{"localhost"}, 0s));
    EXPECT_FALSE(createProbe(TcpEndpoint{
        std::string("host\0suffix", 11), 80}));
}

TEST(Probe, FactoryCreatesConcreteTypes)
{
    auto http = createProbe(HttpEndpoint{"https://example.invalid"});
    ASSERT_TRUE(http);
    EXPECT_NE(dynamic_cast<HttpProbe*>(http->get()), nullptr);
    auto tcp = createProbe(TcpEndpoint{"localhost", 80});
    ASSERT_TRUE(tcp);
    EXPECT_NE(dynamic_cast<TcpProbe*>(tcp->get()), nullptr);
    auto udp = createProbe(UdpEndpoint{"localhost", 53, ""});
    ASSERT_TRUE(udp);
    EXPECT_NE(dynamic_cast<UdpProbe*>(udp->get()), nullptr);
    auto icmp = createProbe(IcmpEndpoint{"localhost"});
    ASSERT_TRUE(icmp);
    EXPECT_NE(dynamic_cast<IcmpProbe*>(icmp->get()), nullptr);
}

TEST_F(ProbeTest, OwnsConfiguration)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    mw::E<std::unique_ptr<ProbeInterface>> probe;
    {
        TcpEndpoint config{"127.0.0.1", port};
        probe = createProbe(config, 1s);
        ASSERT_TRUE(probe) << probe.error().msg();
        config.host.clear();
        config.port = 0;
    }
    auto result = (*probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

mw::E<ProbeResult> runOwnedProbe(std::unique_ptr<ProbeInterface> probe)
{
    return probe->probe();
}

TEST_F(ProbeTest, TaskOwnsProbe)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    auto probe = createProbe(TcpEndpoint{"127.0.0.1", port}, 1s);
    ASSERT_TRUE(probe) << probe.error().msg();
    auto task = std::async(std::launch::async, runOwnedProbe,
                           std::move(*probe));
    EXPECT_EQ(probe->get(), nullptr);
    auto result = task.get();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

TEST_F(ProbeTest, TcpConnectAndTiming)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    const auto before = std::chrono::duration_cast<std::chrono::seconds>(
        std::chrono::system_clock::now().time_since_epoch()).count();
    auto result_probe = createProbe(TcpEndpoint{"127.0.0.1", port}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
    EXPECT_GE(result->timestamp, before);
    EXPECT_GE(result->duration_microsecond, 0);
}

TEST_F(ProbeTest, TcpIpv6)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET6, true));
    auto result_probe = createProbe(TcpEndpoint{"::1", port}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

TEST_F(ProbeTest, HostnameResolution)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    auto result_probe = createProbe(TcpEndpoint{"localhost", port}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

TEST_F(ProbeTest, RefusedTcpAndHttp)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM));
    auto tcp_probe = createProbe(TcpEndpoint{"127.0.0.1", port}, 1s);
    ASSERT_TRUE(tcp_probe) << tcp_probe.error().msg();
    auto tcp = (*tcp_probe)->probe();
    ASSERT_TRUE(tcp) << tcp.error().msg();
    EXPECT_EQ(tcp->status, ProbeStatus::BAD);
    auto http_probe = createProbe(HttpEndpoint{
        "http://127.0.0.1:" + std::to_string(port)}, 1s);
    ASSERT_TRUE(http_probe) << http_probe.error().msg();
    auto http = (*http_probe)->probe();
    ASSERT_TRUE(http) << http.error().msg();
    EXPECT_EQ(http->status, ProbeStatus::BAD);
}

TEST_F(ProbeTest, HttpStatusCodes)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    for(int status : {200, 204, 299, 301, 404, 503})
    {
        worker = std::jthread(&ProbeTest::serveHttp, this, status, false);
        auto result_probe = createProbe(HttpEndpoint{
            "http://127.0.0.1:" + std::to_string(port)}, 1s);
        ASSERT_TRUE(result_probe) << result_probe.error().msg();
        auto result = (*result_probe)->probe();
        ASSERT_TRUE(result) << result.error().msg();
        EXPECT_EQ(result->status, status >= 200 && status < 300
            ? ProbeStatus::GOOD : ProbeStatus::BAD) << status;
        worker.join();
    }
}

TEST_F(ProbeTest, HttpTimeout)
{
    ASSERT_TRUE(bindServer(SOCK_STREAM, AF_INET, true));
    worker = std::jthread(&ProbeTest::serveHttp, this, 200, true);
    auto result_probe = createProbe(HttpEndpoint{
        "http://127.0.0.1:" + std::to_string(port)}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::BAD);
    EXPECT_GE(result->duration_microsecond, 900000);
    EXPECT_LT(result->duration_microsecond, 3000000);
}

TEST_F(ProbeTest, UdpReplyAndBinaryPayload)
{
    ASSERT_TRUE(bindServer(SOCK_DGRAM));
    worker = std::jthread(&ProbeTest::serveUdp, this, false);
    const std::string payload("hello\0world", 11);
    auto result_probe = createProbe(
        UdpEndpoint{"127.0.0.1", port, payload}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
    worker.join();
    EXPECT_EQ(received_payload, payload);
}

TEST_F(ProbeTest, UdpEmptyReplyAndIpv6)
{
    ASSERT_TRUE(bindServer(SOCK_DGRAM, AF_INET6));
    worker = std::jthread(&ProbeTest::serveUdp, this, true);
    auto result_probe = createProbe(UdpEndpoint{"::1", port, ""}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
    worker.join();
    EXPECT_TRUE(received_payload.empty());
}

TEST_F(ProbeTest, UdpSilenceIsOther)
{
    ASSERT_TRUE(bindServer(SOCK_DGRAM));
    auto result_probe = createProbe(
        UdpEndpoint{"127.0.0.1", port, "probe"}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::OTHER);
    EXPECT_GE(result->duration_microsecond, 900000);
    EXPECT_LT(result->duration_microsecond, 3000000);
}

TEST_F(ProbeTest, UdpRefused)
{
    ASSERT_TRUE(bindServer(SOCK_DGRAM));
    close(server);
    server = -1;
    auto result_probe = createProbe(
        UdpEndpoint{"127.0.0.1", port, "probe"}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::BAD);
}

TEST(Probe, DnsFailureIsBad)
{
    auto result_probe = createProbe(
        TcpEndpoint{"status-tracker.invalid", 80}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::BAD);
    EXPECT_LT(result->duration_microsecond, 3000000);
}

TEST(Probe, IcmpIpv4)
{
    const int allowed = socket(AF_INET, SOCK_DGRAM, IPPROTO_ICMP);
    auto result_probe = createProbe(IcmpEndpoint{"127.0.0.1"}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    if(allowed < 0)
    {
        EXPECT_FALSE(result);
        GTEST_SKIP() << "ICMP ping sockets unavailable for this process";
    }
    close(allowed);
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

TEST(Probe, IcmpIpv6)
{
    const int allowed = socket(AF_INET6, SOCK_DGRAM, IPPROTO_ICMPV6);
    auto result_probe = createProbe(IcmpEndpoint{"::1"}, 1s);
    ASSERT_TRUE(result_probe) << result_probe.error().msg();
    auto result = (*result_probe)->probe();
    if(allowed < 0)
    {
        EXPECT_FALSE(result);
        GTEST_SKIP() << "ICMPv6 ping sockets unavailable for this process";
    }
    close(allowed);
    ASSERT_TRUE(result) << result.error().msg();
    EXPECT_EQ(result->status, ProbeStatus::GOOD);
}

}