#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>
#include <cerrno>
#include <chrono>
#include <iostream>
#include <stdexcept>
#include <string>

struct Observation {
    int closed{0};
    bool close_ok{true};
};

class Fd {
    int value_;
    const char* name_;
    Observation& observation_;
public:
    Fd(int value, const char* name, Observation& observation)
        : value_{value}, name_{name}, observation_{observation} {
        if (value_ < 0) throw std::runtime_error(name_);
        std::cout << "open=" << name_ << '\n';
    }
    Fd(const Fd&) = delete;
    Fd& operator=(const Fd&) = delete;
    ~Fd() {
        const int result{::close(value_)};
        ++observation_.closed;
        observation_.close_ok = observation_.close_ok && result == 0;
        std::cout << "close=" << name_ << " result=" << result << '\n';
    }
    int get() const { return value_; }
};

using Clock = std::chrono::steady_clock;

void nonblocking(int fd) {
    const int flags{::fcntl(fd, F_GETFL, 0)};
    if (flags < 0 || ::fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0)
        throw std::runtime_error("fcntl");
}

void wait_ready(int fd, short events, Clock::time_point deadline) {
    for (;;) {
        const auto now{Clock::now()};
        if (now >= deadline) throw std::runtime_error("deadline");
        auto left{std::chrono::duration_cast<std::chrono::milliseconds>(deadline - now).count()};
        if (left == 0) left = 1; // Wait at least 1 ms when a positive fraction remains.
        // This demonstration only supplies deadlines at most 2000 ms ahead.
        pollfd event{fd, events, 0};
        const int result{::poll(&event, 1, static_cast<int>(left))};
        if (result < 0 && errno == EINTR) continue;
        if (result < 0) throw std::runtime_error("poll");
        if (Clock::now() >= deadline) throw std::runtime_error("deadline");
        if (result == 0) continue; // An early, rounded timeout does not reset the deadline.
        if ((event.revents & (POLLERR | POLLNVAL)) != 0)
            throw std::runtime_error("poll-event");
        if ((event.revents & (events | POLLHUP)) != 0) return;
    }
}

int main(int argc, char* argv[]) {
    const std::string mode{argc == 2 ? argv[1] : "normal"};
    if (argc > 2 || (mode != "normal" && mode != "one-byte" && mode != "timeout")) return 2;
    Observation observation;
    bool caught_timeout{false};
    bool timeout_armed{false};
    try {
        const auto deadline{Clock::now() + std::chrono::milliseconds{2000}};
        Fd listener{::socket(AF_INET, SOCK_STREAM, 0), "listener", observation};
        sockaddr_in address{};
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        address.sin_port = htons(0);
        if (::bind(listener.get(), reinterpret_cast<sockaddr*>(&address), sizeof address) != 0 ||
            ::listen(listener.get(), 1) != 0) throw std::runtime_error("bind/listen");
        socklen_t length{sizeof address};
        if (::getsockname(listener.get(), reinterpret_cast<sockaddr*>(&address), &length) != 0)
            throw std::runtime_error("getsockname");
        std::cout << "port=" << ntohs(address.sin_port) << '\n';
        nonblocking(listener.get());
        Fd client{::socket(AF_INET, SOCK_STREAM, 0), "client", observation};
        nonblocking(client.get());
        if (::connect(client.get(), reinterpret_cast<sockaddr*>(&address), length) != 0 &&
            errno != EINPROGRESS) throw std::runtime_error("connect");
        wait_ready(client.get(), POLLOUT, deadline);
        int error{0};
        length = sizeof error;
        if (::getsockopt(client.get(), SOL_SOCKET, SO_ERROR, &error, &length) != 0 || error != 0)
            throw std::runtime_error("SO_ERROR");
        std::cout << "connect-error=" << error << '\n';
        int accepted{-1};
        while (accepted < 0) {
            wait_ready(listener.get(), POLLIN, deadline);
            accepted = ::accept(listener.get(), nullptr, nullptr);
            if (accepted < 0 && errno != EINTR && errno != EAGAIN && errno != EWOULDBLOCK)
                throw std::runtime_error("accept");
        }
        Fd server{accepted, "server", observation};
        nonblocking(server.get());
        if (mode == "timeout") {
            timeout_armed = true;
            std::cout << "waiting=no-data\n";
            wait_ready(server.get(), POLLIN, Clock::now() + std::chrono::milliseconds{30});
            throw std::runtime_error("unexpected-readiness");
        }
        const std::string input{"abcdef"};
        std::size_t sent{0};
        // The peer remains alive. This file does not test premature close or suppress SIGPIPE.
        while (sent < input.size()) {
            wait_ready(client.get(), POLLOUT, deadline);
            const auto n{::send(client.get(), input.data() + sent, input.size() - sent, 0)};
            if (n < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)) continue;
            if (n <= 0) throw std::runtime_error("send");
            sent += static_cast<std::size_t>(n);
            std::cout << "send=" << n << " total=" << sent << '\n';
        }
        if (::shutdown(client.get(), SHUT_WR) != 0) throw std::runtime_error("shutdown");
        std::cout << "shutdown=write\n";
        char buffer[2];
        const std::size_t capacity{mode == "one-byte" ? 1U : 2U};
        std::string output;
        for (;;) {
            wait_ready(server.get(), POLLIN, deadline);
            const auto n{::recv(server.get(), buffer, capacity, 0)};
            if (n < 0 && (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)) continue;
            if (n < 0) throw std::runtime_error("recv");
            if (n == 0) break;
            output.append(buffer, static_cast<std::size_t>(n));
            std::cout << "recv=" << n << " total=" << output.size() << '\n';
        }
        if (output != input) throw std::runtime_error("data");
        std::cout << "data=" << output << " eof=1 capacity=" << capacity << '\n';
    } catch (const std::exception& error) {
        if (mode != "timeout" || !timeout_armed || std::string{error.what()} != "deadline") {
            std::cerr << "unexpected=" << error.what() << '\n';
            return 1;
        }
        caught_timeout = true;
        std::cout << "caught=deadline\n";
    }
    std::cout << "closed=" << observation.closed << " close-ok=" << observation.close_ok
              << " timeout=" << caught_timeout << '\n';
    return observation.closed == 3 && observation.close_ok &&
           caught_timeout == (mode == "timeout") ? 0 : 1;
}