#include "main.h"

#include <contrib/libs/grpc/include/grpc++/server.h>
#include <contrib/libs/grpc/include/grpc++/server_builder.h>

std::atomic_bool RECEIVE_ENABLED{true};

static void ExitHandler(int signum);
static void SetExitHandler();

void RunMainServer(const TCaptureOptions &options) {
    SetExitHandler();

    TPcap pcap(options.linkType, options.snapSize);
    TPcapDump pcapDump(pcap.Session(), options.output);

    TPacketsWriterQueue packetsQueue;
    TNetworkCaptureRouterImpl captureRouter(&packetsQueue);

    grpc::ServerBuilder builder;
    builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::NUM_CQS, options.threadCount);
    builder.SetSyncServerOption(grpc::ServerBuilder::SyncServerOption::MIN_POLLERS, options.threadCount);
    builder.AddListeningPort(options.bind, grpc::InsecureServerCredentials());
    builder.RegisterService(&captureRouter);

    std::unique_ptr<grpc::Server> server(builder.BuildAndStart());

    auto serverWriter = std::make_unique<TPacketsServerWriter>(&packetsQueue, &pcapDump);
    std::thread movebleThread([&serverWriter, &server]() {
        serverWriter->Start(server.get());
    });
    TSafeThread serverWriterWorker(std::move(movebleThread));

    server->Wait();
}

void TPacketsServerWriter::Start(grpc::Server* grpcServer) {
    TPacketHeader pkthdr;

    while(RECEIVE_ENABLED.load()) {
        auto msg = std::make_unique<NNetworkCaptureFrame::TNetworkCaptureFrame>();
        while (!PacketsQueue_->Dequeue(msg) && RECEIVE_ENABLED.load()) {};

        pkthdr.ts = (struct timeval) {(long) msg->Gettssec(), (int) msg->Gettsusec()};
        pkthdr.caplen = msg->Getcaplen();
        pkthdr.len = msg->Getlen();

        DumpFile_->Dump(&pkthdr, (const u_char *) msg->Getpacket().c_str());
    }

    if (grpcServer != nullptr) {
        grpcServer->Shutdown();
    }
}

grpc::Status TNetworkCaptureRouterImpl::CaptureStream(
    grpc::ServerContext* context,
    grpc::ServerReader<NNetworkCaptureFrame::TNetworkCaptureFrame>* stream,
    NNetworkCaptureFrame::TEmpty* empty)
{
    Y_UNUSED(context);
    Y_UNUSED(empty);

    bool cont = false;
    std::unique_ptr<NNetworkCaptureFrame::TNetworkCaptureFrame> frame;
    do {
        frame = std::make_unique<NNetworkCaptureFrame::TNetworkCaptureFrame>();

        cont = stream->Read(frame.get());
        PacketsQueue_->Enqueue(std::move(frame));
    } while(cont);

    return grpc::Status::OK;
}

static void ExitHandler(int signum) {
    Y_UNUSED(signum);

    RECEIVE_ENABLED.store(false, std::memory_order_relaxed);
}

static void SetExitHandler() {
    struct sigaction act {};

    act.sa_handler = ExitHandler;
    act.sa_flags = 0;
    sigemptyset(&act.sa_mask);
    sigaddset(&act.sa_mask, SIGINT);
    sigaddset(&act.sa_mask, SIGTERM);

    sigaction(SIGINT, &act, NULL);
    sigaction(SIGTERM, &act, NULL);
}
