#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
using socket_type = SOCKET;
static void close_socket(socket_type socket) { closesocket(socket); }
#else
#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
using socket_type = int;
static void close_socket(socket_type socket) { close(socket); }
#endif

#include "crypto/crypto.h"
#include "net/voice_frame.h"
#include <filesystem>
#include <fstream>
#include <iostream>

static bool transfer(voicecat::crypto::TlsContext& tls, uint8_t* data, size_t size, bool writing) {
    while (size != 0) {
        int count = writing ? tls.write(data, size) : tls.read(data, size);
        if (count <= 0) return false;
        data += count;
        size -= count;
    }
    return true;
}

int main(int argc, char** argv) {
    if (argc != 2 || sodium_init() < 0) return 1;
#ifdef _WIN32
    WSADATA data{};
    if (WSAStartup(MAKEWORD(2, 2), &data) != 0) return 1;
#endif
    try {
        auto certificate = voicecat::crypto::ServerCert::generate("dotnet-tls-oracle");
        auto directory = std::filesystem::path(argv[1]);
        socket_type listener = socket(AF_INET, SOCK_STREAM, 0);
        sockaddr_in address{};
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        if (bind(listener, reinterpret_cast<sockaddr*>(&address), sizeof(address)) != 0 || listen(listener, 1) != 0) return 1;
        socklen_t length = sizeof(address);
        if (getsockname(listener, reinterpret_cast<sockaddr*>(&address), &length) != 0) return 1;
        certificate.save(directory / "server.crt", directory / "server.key");
        voicecat::crypto::ServerIdentity::generate().save(directory / "identity.key");
        std::ofstream(directory / "port.txt") << ntohs(address.sin_port);
        socket_type peer = accept(listener, nullptr, nullptr);
        close_socket(listener);
        if (peer == static_cast<socket_type>(-1)) return 1;
        voicecat::crypto::TlsContext tls(voicecat::crypto::TlsContext::Role::Server, &certificate);
        tls.set_read_timeout(10000);
        std::string error;
        if (!tls.handshake(static_cast<int>(peer), error)) { std::cerr << error; return 1; }
        auto sender = voicecat::crypto::SodiumMediaCrypto::derive_send(tls, false);
        auto receiver = voicecat::crypto::SodiumMediaCrypto::derive_recv(tls, false);
        if (!sender || !receiver) return 1;
        voicecat::net::VoiceFrame header;
        header.ssrc = 42;
        header.seq = sender->peek_send_counter();
        std::array<uint8_t, 41> packet{};
        voicecat::net::serialize_header(header, packet.data());
        const std::array<uint8_t, 5> message{ 'h', 'e', 'l', 'l', 'o' };
        if (sender->seal(message.data(), message.size(), packet.data(), 20, packet.data() + 20, 21) != 21) return 1;
        if (!transfer(tls, packet.data(), packet.size(), true) || !transfer(tls, packet.data(), packet.size(), false)) return 1;
        std::array<uint8_t, 5> recovered{};
        if (receiver->open(packet.data() + 20, 21, packet.data(), 20, recovered.data(), recovered.size()) != 5 || recovered != message) return 1;
        uint8_t acknowledgement = 1;
        if (!transfer(tls, &acknowledgement, 1, true)) return 1;
        return 0;
    } catch (const std::exception& error) {
        std::cerr << error.what();
        return 1;
    }
}