// Copyright (c) 2014-2026 Savoir-faire Linux Inc.
// SPDX-License-Identifier: MIT

#include "network_utils.h"

#ifdef _WIN32
#include "utils.h"
#include <io.h>
#include <string>
#include <cstring>
#define close(x)                  closesocket(x)
#define write(s, b, f)            send(s, b, (int) strlen(b), 0)
#define _poll(fds, nfds, timeout) WSAPoll(fds, nfds, timeout)
#else
#include <fcntl.h>
#include <poll.h>
#define _poll(fds, nfds, timeout) poll(fds, nfds, timeout)
#endif

// number of file descriptors to poll in openSockets()
#define NUM_FDS 3

#include <iostream>

namespace dht {
namespace net {

int
bindSocket(const SockAddr& addr, SockAddr& bound)
{
    bool is_ipv6 = addr.getFamily() == AF_INET6;
    int sock = socket(is_ipv6 ? PF_INET6 : PF_INET, SOCK_DGRAM, 0);
    if (sock < 0)
        throw DhtException(std::string("Can't open socket: ") + strerror(sock));
    int set = 1;
#ifdef SO_NOSIGPIPE
    setsockopt(sock, SOL_SOCKET, SO_NOSIGPIPE, (const char*) &set, sizeof(set));
#endif
    if (is_ipv6)
        setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, (const char*) &set, sizeof(set));
    net::setNonblocking(sock);
    int rc = bind(sock, addr.get(), addr.getLength());
    if (rc < 0) {
        rc = errno;
        close(sock);
        throw DhtException("Can't bind socket on " + addr.toString() + " " + strerror(rc));
    }
    sockaddr_storage ss;
    socklen_t ss_len = sizeof(ss);
    getsockname(sock, (sockaddr*) &ss, &ss_len);
    bound = {ss, ss_len};
    return sock;
}

bool
setNonblocking(int fd, bool nonblocking)
{
#ifdef _WIN32
    unsigned long mode = !!nonblocking;
    int rc = ioctlsocket(fd, FIONBIO, &mode);
    return rc == 0;
#else
    int rc = fcntl(fd, F_GETFL, 0);
    if (rc < 0)
        return false;
    rc = fcntl(fd, F_SETFL, nonblocking ? (rc | O_NONBLOCK) : (rc & ~O_NONBLOCK));
    return rc >= 0;
#endif
}

#ifdef _WIN32
void
udpPipe(int fds[2])
{
    int lst = socket(AF_INET, SOCK_DGRAM, 0);
    if (lst < 0)
        throw DhtException(std::string("Can't open socket: ") + strerror(WSAGetLastError()));
    sockaddr_in inaddr;
    sockaddr addr;
    memset(&inaddr, 0, sizeof(inaddr));
    memset(&addr, 0, sizeof(addr));
    inaddr.sin_family = AF_INET;
    inaddr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    inaddr.sin_port = 0;
    int yes = 1;
    setsockopt(lst, SOL_SOCKET, SO_REUSEADDR, (char*) &yes, sizeof(yes));
    int rc = bind(lst, (sockaddr*) &inaddr, sizeof(inaddr));
    if (rc < 0) {
        close(lst);
        throw DhtException("Can't bind socket on " + print_addr((sockaddr*) &inaddr, sizeof(inaddr)) + " "
                           + strerror(rc));
    }
    socklen_t len = sizeof(addr);
    getsockname(lst, &addr, &len);
    fds[0] = lst;
    fds[1] = socket(AF_INET, SOCK_DGRAM, 0);
    connect(fds[1], &addr, len);
}
#endif

UdpSocket::UdpSocket(in_port_t port, const std::shared_ptr<Logger>& l)
    : logger(l)
{
    SockAddr bind4;
    bind4.setFamily(AF_INET);
    bind4.setPort(port);
    SockAddr bind6;
    bind6.setFamily(AF_INET6);
    bind6.setPort(port);
    std::lock_guard lk(lock);
    openSockets(bind4, bind6);
}

UdpSocket::UdpSocket(const SockAddr& bind4, const SockAddr& bind6, const std::shared_ptr<Logger>& l)
    : logger(l)
{
    std::lock_guard lk(lock);
    openSockets(bind4, bind6);
}

UdpSocket::~UdpSocket()
{
    stop();
    if (rcv_thread.joinable())
        rcv_thread.join();
}

int
UdpSocket::sendTo(const SockAddr& dest, const uint8_t* data, size_t size, [[maybe_unused]] bool replied)
{
    if (not dest)
        return EFAULT;

    int s;
    switch (dest.getFamily()) {
    case AF_INET:
        s = s4;
        break;
    case AF_INET6:
        s = s6;
        break;
    default:
        s = -1;
        break;
    }

    if (s < 0)
        return EAFNOSUPPORT;

    int flags = 0;
#ifdef MSG_CONFIRM
    if (replied)
        flags |= MSG_CONFIRM;
#endif
#ifdef MSG_NOSIGNAL
    flags |= MSG_NOSIGNAL;
#endif

    if (sendto(s, (const char*) data, size, flags, dest.get(), dest.getLength()) == -1) {
        int err = errno;
        if (logger)
            logger->d("Can't send message to %s: %s", dest.toString().c_str(), strerror(err));
        if (err == EPIPE || err == ENOTCONN || err == ECONNRESET) {
            std::lock_guard lk(lock);
            auto bind4 = std::move(bound4), bind6 = std::move(bound6);
            openSockets(bind4, bind6);
            return sendTo(dest, data, size, false);
        }
        return err;
    }
    return 0;
}

void
UdpSocket::openSockets(const SockAddr& bind4, const SockAddr& bind6)
{
    stop();
    if (rcv_thread.joinable())
        rcv_thread.join();

    int stopfds[2];
#ifndef _WIN32
    auto status = pipe(stopfds);
    if (status == -1) {
        throw DhtException(std::string("Can't open pipe: ") + strerror(errno));
    }
#else
    udpPipe(stopfds);
#endif
    int stop_readfd = stopfds[0];

    stopfd = stopfds[1];
    s4 = -1;
    s6 = -1;

    bound4 = {};
    if (bind4) {
        try {
            s4 = bindSocket(bind4, bound4);
        } catch (const DhtException& e) {
            if (logger)
                logger->e("Can't bind inet socket: %s", e.what());
        }
    }

#if 1
    bound6 = {};
    if (bind6) {
        if (bind6.getPort() == 0) {
            // Attempt to use the same port as IPv4 with IPv6
            if (auto p4 = bound4.getPort()) {
                auto b6 = bind6;
                b6.setPort(p4);
                try {
                    s6 = bindSocket(b6, bound6);
                } catch (const DhtException& e) {
                    if (logger)
                        logger->e("Can't bind inet6 socket: %s", e.what());
                }
            }
        }
        if (s6 == -1) {
            try {
                s6 = bindSocket(bind6, bound6);
            } catch (const DhtException& e) {
                if (logger)
                    logger->e("Can't bind inet6 socket: %s", e.what());
            }
        }
    }
#endif

    if (s4 == -1 && s6 == -1) {
        throw DhtException("Can't bind socket");
    }

    running = true;
    rcv_thread = std::thread([this, stop_readfd, ls4 = s4, ls6 = s6]() mutable {
        struct pollfd fds[NUM_FDS];
        for (int i = 0; i < NUM_FDS; i++)
            fds[i].events = POLLIN;
        constexpr size_t stop_readfd_index = 0, ls4_index = 1, ls6_index = 2;
        try {
            while (running) {
                // ls4 and ls6 can be negative, but this doesn't require any special handling
                // because poll() will simply ignore them (and set revents to 0) in that case
                fds[stop_readfd_index].fd = stop_readfd;
                fds[ls4_index].fd = ls4;
                fds[ls6_index].fd = ls6;

                int rc = _poll(fds, NUM_FDS, -1);
                if (rc < 0) {
                    if (errno != EINTR) {
                        if (logger)
                            logger->e("Poll error: %s", strerror(errno));
                        std::this_thread::sleep_for(std::chrono::seconds(1));
                    }
                }

                if (not running)
                    break;

                if (rc > 0) {
                    std::array<uint8_t, 1024 * 64> buf;
                    sockaddr_storage from;
                    socklen_t from_len = sizeof(from);

                    if (fds[stop_readfd_index].revents & POLLIN) {
                        if (recv(stop_readfd, (char*) buf.data(), buf.size(), 0) < 0) {
                            if (logger)
                                logger->e("Got stop packet error: %s", strerror(errno));
                            break;
                        }
                    } else if (fds[ls4_index].revents & POLLIN)
                        rc = recvfrom(ls4, (char*) buf.data(), buf.size(), 0, (sockaddr*) &from, &from_len);
                    else if (fds[ls6_index].revents & POLLIN)
                        rc = recvfrom(ls6, (char*) buf.data(), buf.size(), 0, (sockaddr*) &from, &from_len);
                    else
                        continue;

                    if (rc > 0) {
                        auto pkts = getNewPacket();
                        auto& pkt = pkts.front();
                        pkt.data.insert(pkt.data.end(), buf.begin(), buf.begin() + rc);
                        pkt.from = {from, from_len};
                        pkt.received = clock::now();
                        onReceived(std::move(pkts));
                    } else if (rc == -1) {
                        int err = errno;
                        if (logger)
                            logger->e("Error receiving packet: %s", strerror(err));
                        if (err == EPIPE || err == ENOTCONN || err == ECONNRESET) {
                            if (not running)
                                break;
                            std::unique_lock lk(lock, std::try_to_lock);
                            if (lk.owns_lock()) {
                                if (not running)
                                    break;
                                if (ls4 >= 0) {
                                    close(ls4);
                                    try {
                                        ls4 = bindSocket(bound4, bound4);
                                    } catch (const DhtException& e) {
                                        if (logger)
                                            logger->e("Can't bind inet socket: %s", e.what());
                                    }
                                }
                                if (ls6 >= 0) {
                                    close(ls6);
                                    try {
                                        ls6 = bindSocket(bound6, bound6);
                                    } catch (const DhtException& e) {
                                        if (logger)
                                            logger->e("Can't bind inet6 socket: %s", e.what());
                                    }
                                }
                                if (ls4 < 0 && ls6 < 0)
                                    break;
                                s4 = ls4;
                                s6 = ls6;
                            } else {
                                break;
                            }
                        }
                    }
                }
            }
        } catch (const std::exception& e) {
            if (logger)
                logger->e("Error in UdpSocket rx thread: %s", e.what());
        }
        if (ls4 >= 0)
            close(ls4);
        if (ls6 >= 0)
            close(ls6);
        if (stop_readfd != -1)
            close(stop_readfd);
        if (stopfd != -1)
            close(stopfd);
        std::unique_lock lk(lock, std::try_to_lock);
        if (lk.owns_lock()) {
            s4 = -1;
            s6 = -1;
            bound4 = {};
            bound6 = {};
            stopfd = -1;
        }
    });
}

void
UdpSocket::stop()
{
    if (running.exchange(false)) {
        auto sfd = stopfd;
        if (sfd != -1 && write(sfd, "\0", 1) == -1) {
            if (logger)
                logger->e("Can't write to stop fd");
        }
    }
}

} // namespace net
} // namespace dht
