#include <cstdint>
#include <future>
#include <thread>
#include <utility>
#include <vector>

#include <Poco/Net/PollSet.h>
#include <Poco/Net/SocketAddress.h>
#include <Poco/Net/StreamSocket.h>

#include <Daemon/BaseDaemon.h>
#include <Interpreters/Context.h>


int mainEntryDatastoreServer(int argc, char ** argv);

namespace
{

std::string datastore("datastore-server");
std::vector<char *> args{datastore.data()};
std::future<int> main_app;

std::string s_host("0.0.0.0");
char * host = s_host.data();
int64_t port = 9000;

using namespace std::chrono_literals;

bool isMerge(int argc, const char * const * argv)
{
    for (int i = 1; i < argc; ++i)
    {
        std::string_view arg{argv[i]};
        if (std::string_view{arg.begin(), std::ranges::find(arg, '=')} == "-ignore_remaining_args")
            break;
        if (std::string_view{arg.begin(), std::ranges::find(arg, '=')} == "-merge")
            return true;
    }
    return false;
}

void on_exit()
{
    BaseDaemon::terminate();
    main_app.wait();
}

}

extern "C" int LLVMFuzzerInitialize(int * argc, char ***argv);
extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size);

extern "C"
int LLVMFuzzerInitialize(int * argc, char ***argv)
{
        // If it's a merge coordinator don't initialize anything
    if (isMerge(*argc, *argv))
        return 0;

    for (int i = 1; i < *argc; ++i)
    {
        if ((*argv)[i][0] == '-')
        {
            if ((*argv)[i][1] == '-')
                args.push_back((*argv)[i]);
            else
            {
                if (strncmp((*argv)[i], "-host=", 6) == 0)
                {
                    host = (*argv)[i] + 6;
                }
                else if (strncmp((*argv)[i], "-port=", 6) == 0)
                {
                    char * p_end = nullptr;
                    port = strtol((*argv)[i] + 6, &p_end, 10);
                }
            }
        }
    }

    args.push_back(nullptr);

    main_app = std::async(std::launch::async, mainEntryDatastoreServer, args.size() - 1, args.data());

    while (!DB::Context::getGlobalContextInstance() || !DB::Context::getGlobalContextInstance()->isServerCompletelyStarted())
    {
        std::this_thread::sleep_for(100ms);
        if (main_app.wait_for(0s) == std::future_status::ready)
            exit(-1);
    }

    atexit(on_exit);

    return 0;
}

extern "C"
int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
{
    try
    {
        if (main_app.wait_for(0s) == std::future_status::ready)
            return -1;

        if (size == 0)
            return -1;

        Poco::Net::SocketAddress address(host, static_cast<uint16_t>(port));
        Poco::Net::StreamSocket socket;

        socket.connectNB(address);

        Poco::Net::PollSet ps;
        ps.add(socket, Poco::Net::PollSet::POLL_READ | Poco::Net::PollSet::POLL_WRITE);

        std::vector<char> buf(1048576);
        size_t sent = 0;
        while (true)
        {
            auto m = ps.poll(Poco::Timespan(1000000));
            if (m.empty())
                continue;
            if (m.begin()->second & Poco::Net::PollSet::POLL_READ)
            {
                if (int n = socket.receiveBytes(buf.data(), static_cast<int>(buf.size())); n == 0)
                {
                    socket.close();
                    break;
                }

                continue;
            }

            if (sent < size && m.begin()->second & Poco::Net::PollSet::POLL_WRITE)
            {
                sent += socket.sendBytes(data + sent, static_cast<int>(size - sent));
                if (sent == size)
                {
                    socket.shutdownSend();
                    continue;
                }
            }
        }
    }
    catch (...)
    {
        // Ok
    }

    return 0;
}