#pragma once

#include "address.h"

#include <balancer/serval/contrib/cone/cold.h>
#include <balancer/serval/contrib/cone/mun.h>

#include <util/generic/maybe.h>
#include <util/generic/ptr.h>
#include <util/generic/strbuf.h>
#include <util/network/socket.h>

#ifdef _linux_
// Whether bind()ing on the same port multiple times is allowed.
#define SV_CAN_MULTIBIND 1
#else
// XXX macOS (BSD) also has SO_REUSEPORT, but it does not balance load; should use SO_REUSEPORT_LB.
#define SV_CAN_MULTIBIND 0
#endif

namespace NSv {
    static constexpr int ConstantOne = 1;

    class IO {
    public:
        virtual ~IO() = default;

        // See `man 2 read`, except errors are reported through mun.
        virtual TMaybe<size_t> ReadInto(TStringBuf) noexcept = 0;

        // See `man 2 write`, except errors are reported through mun.
        virtual TMaybe<size_t> Write(TStringBuf) noexcept = 0;

        // Return the address of the other side, or AF_UNSPEC if not a connected socket.
        virtual NSv::IP Peer() const noexcept { return {}; };

        // Return the protocol negotiated by ALPN, or an empty string.
        virtual TStringBuf SelectedProtocol() const noexcept { return {}; }
    };

    class TFile: public TMoveOnly, public IO {
    public:
        TFile() = default;
        TFile(int fd) noexcept
            : N_(fd)
        {}
        TFile(TFile&& other) noexcept
            : N_(other.N_)
        {
            other.N_ = -1;
        }
        TFile& operator=(TFile other) noexcept
        {
            std::swap(N_, other.N_);
            return *this;
        }

        ~TFile() {
            if (N_ >= 0)
                ::close(N_);
        }

        operator int() const noexcept {
            return N_;
        }

        explicit operator bool() const noexcept {
            return N_ >= 0;
        }

        TMaybe<size_t> ReadInto(TStringBuf buf) noexcept override {
            ssize_t r = cold_read(*this, const_cast<char*>(buf.data()), buf.size());
            return !(r < 0 MUN_RETHROW_OS) ? TMaybe<size_t>{r} : TMaybe<size_t>{};
        }

        TMaybe<size_t> Write(TStringBuf buf) noexcept override {
            ssize_t w = cold_write(*this, buf.data(), buf.size());
            return !(w < 0 MUN_RETHROW_OS) ? TMaybe<size_t>{w} : TMaybe<size_t>{};
        }

        NSv::IP Peer() const noexcept override {
            NSv::IP ret;
            socklen_t len = sizeof(ret.Data);
            Y_UNUSED(getpeername(*this, &ret.Data.Base, &len));
            return ret;
        }

        template <int family = SOCK_STREAM, class SockOpt>
        static TFile Bind(const NSv::IP& net, int backlog, SockOpt&& sockOpt) noexcept {
            TFile f = socket(net.Data.Base.sa_family, family, 0);
            if (!f
             || sockOpt(f)
             || bind(f, &net.Data.Base, sizeof(net.Data)) < 0
             || (family == SOCK_STREAM && cold_listen(f, backlog) < 0) MUN_RETHROW_OS) {
                return {};
            }
            return f;
        }

        template <int family = SOCK_STREAM>
        static TFile Bind(const NSv::IP& net, bool reuseAddr, int backlog = 256) noexcept {
            auto sockOpt = [reuseAddr](int f) {
                if (reuseAddr) {
                    if (setsockopt(f, SOL_SOCKET, SO_REUSEADDR, &ConstantOne, sizeof(int)) < 0) {
                        return true;
                    }
                }
#if SV_CAN_MULTIBIND
                return setsockopt(f, SOL_SOCKET, /*SO_REUSEPORT*/15, &ConstantOne, sizeof(int)) < 0;
#else
                return f < 0;
#endif
            };
            return Bind<family>(net, backlog, sockOpt);
        }

        static TFile Connect(const NSv::IP& ip) noexcept {
            TFile f = socket(ip.Data.Base.sa_family, SOCK_STREAM, 0);
            if (!f || cold_connect(f, &ip.Data.Base, sizeof(ip.Data)) < 0
                   || setsockopt(f, IPPROTO_TCP, TCP_NODELAY, &ConstantOne, sizeof(int)) < 0 MUN_RETHROW_OS) {
                return {};
            }
            return f;
        }

        static TFile Accept(NSv::TFile& sk) noexcept {
            TFile f = cold_accept(sk, nullptr, 0);
            // XXX at least on Linux, can actually set TCP_NODELAY on the listening socket...
            if (!f || setsockopt(f, IPPROTO_TCP, TCP_NODELAY, &ConstantOne, sizeof(int)) < 0
                   || setsockopt(f, SOL_SOCKET, SO_KEEPALIVE, &ConstantOne, sizeof(int)) < 0 MUN_RETHROW_OS) {
                return {};
            }
            return f;
        }

    private:
        int N_ = -1;
    };
}
