// Axel '0vercl0k' Souchet - December 29 2022
#pragma once
#include "oua.h"
#include <fmt/printf.h>
#include <optional>
#include <vector>
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>

#pragma comment(lib, "ws2_32.lib")

namespace sock {

    constexpr bool DumpSends = false;
    constexpr size_t DumpSendsMaxBytes = 100;

    constexpr bool DumpReceives = false;
    constexpr size_t DumpReceivesMaxBytes = 100;

    [[nodiscard]] std::optional<SOCKET> Connect(const std::string& Endpoint,
        const uint32_t Port) {
        const int SockFamily = AF_INET;
        const int SockType = SOCK_STREAM;
        const int SockProto = IPPROTO_TCP;

        struct addrinfo Hints = {};
        Hints.ai_family = SockFamily;
        Hints.ai_socktype = SockType;
        Hints.ai_protocol = SockProto;

        struct addrinfo* Addresses = nullptr;
        if (getaddrinfo(Endpoint.data(), nullptr, &Hints, &Addresses) != 0) {
            fmt::print("getaddrinfo failed w/ {}\n", WSAGetLastError());
            return {};
        }

        if (Addresses == nullptr) {
            fmt::print("Failed to find an address for {}:{}\n", Endpoint, Port);
            return {};
        }

        SOCKET Socket = socket(SockFamily, SockType, SockProto);
        if (Socket == INVALID_SOCKET) {
            fmt::print("socket failed w/ {}\n", WSAGetLastError());
            return {};
        }

        sockaddr_in ConnectAddr = {};
        ConnectAddr.sin_family = Addresses->ai_family;
        memcpy(&ConnectAddr, Addresses->ai_addr, sizeof(ConnectAddr));
        ConnectAddr.sin_port = htons(Port);
        freeaddrinfo(Addresses);

        if (connect(Socket, (sockaddr*)&ConnectAddr, sizeof(ConnectAddr)) != 0) {
            fmt::print("connect failed to {}:{} w/ {}\n", Endpoint, Port,
                WSAGetLastError());
            closesocket(Socket);
            return {};
        }

        return Socket;
    }

    [[nodiscard]] bool Send(SOCKET Socket, const std::vector<uint8_t>& Packet) {
        if constexpr (DumpSends) {
            size_t Counter = 0;
            fmt::print("  > Sending {} bytes: ", Packet.size());
            for (const auto& Byte : Packet) {
                fmt::print("{:02x} ", Byte);
                if (Counter++ > DumpSendsMaxBytes) {
                    fmt::print("...");
                    break;
                }
            }
            fmt::print("\n");
        }
        const int Status =
            send(Socket, (const char*)Packet.data(), int(Packet.size()), 0);
        if (Status != Packet.size()) {
            fmt::print("send failed w/ {} {}\n", WSAGetLastError(), Status);
            return false;
        }

        return true;
    }

    template <typename Message_t>
    [[nodiscard]] bool Send(SOCKET Socket, const Message_t& Msg) {
        std::vector<uint8_t> Buffer;
        Msg.SerializeInto(Buffer);
        const size_t MessageMaxLength = 0x2'000;
        if (Buffer.size() < MessageMaxLength) {
            return Send(Socket, Buffer);
        }

        size_t FragmentN = 0;
        for (const auto& Fragment : oua::helpers::Fragment(Buffer)) {
            if (!Send(Socket, Fragment)) {
                fmt::print("Send fragment {} failed\n", FragmentN);
                return false;
            }

            FragmentN++;
        }

        return true;
    }

    [[nodiscard]] bool Send(SOCKET Socket,
        const std::vector<std::vector<uint8_t>>& Fragments) {
        size_t FragmentN = 0;
        for (const auto& Fragment : Fragments) {
            if (!Send(Socket, Fragment)) {
                fmt::print("Send fragment {} failed\n", FragmentN);
                return false;
            }

            FragmentN++;
        }

        return true;
    }

    [[nodiscard]] std::optional<size_t> RecvInto(SOCKET Socket,
        std::vector<uint8_t>& Buffer) {
        uint8_t StackBuffer[128];
        size_t Received = 0;
        while (1337) {
            const int Result =
                recv(Socket, (char*)StackBuffer, sizeof(StackBuffer), 0);
            if (Result < 0) {
                fmt::print("recv failed w/ {}\n", WSAGetLastError());
                return {};
            }

            Received += Result;
            Buffer.insert(Buffer.end(), StackBuffer, StackBuffer + Result);
            if (Result == sizeof(StackBuffer)) {
                continue;
            }

            break;
        }

        if constexpr (DumpReceives) {
            fmt::print("  < Receiving {} bytes: ", Received);
            size_t Counter = 0;
            for (const auto& Byte : Buffer) {
                fmt::print("{:02x} ", Byte);
                if (Counter++ > DumpReceivesMaxBytes) {
                    fmt::print("...");
                    break;
                }
            }
            fmt::print("\n");
        }

        return Received;
    }

    struct WSAInitializer_t {
        WSADATA Wsa;
        WSAInitializer_t() {
            if (WSAStartup(MAKEWORD(2, 2), &Wsa) != 0) {
                std::abort();
            }
        }

        ~WSAInitializer_t() { WSACleanup(); }
    };
}; // namespace sock
