#pragma once

#include <balancer/production/x/okula/apps/common/pcap.h>
#include <balancer/production/x/okula/apps/tcp_capture/proto/network_frame.pb.h>
#include <balancer/production/x/okula/apps/tcp_capture/proto/network_frame.grpc.pb.h>
#include <library/cpp/threading/chunk_queue/queue.h>

#include <contrib/libs/grpc/include/grpc++/channel.h>
#include <contrib/libs/grpc/include/grpc++/create_channel.h>
#include <contrib/libs/grpc/include/grpc++/security/server_credentials.h>
#include <contrib/libs/grpc/include/grpc++/server.h>

#include <thread>

#if defined(_linux_)
#   include <sched.h>
#   include <pthread.h>
#endif

struct TCaptureOptions {
    bool clientMode;
    TString iface;
    TString filter;
    TString remoteServer;
    TString bind;
    TString output;
    ui32 threadCount;
    ui64 numberOfPackets;
    ui32 snapSize;
    ui32 bufferSize;
    ui32 timeout;
    int linkType;
};

using TPacketHeader = struct pcap_pkthdr;

void RunMain(int argc, char **argv);
void RunMainClient(const TCaptureOptions &options);
void RunMainServer(const TCaptureOptions &options);
static inline void PrintStats(const TString& str, uint64_t counter, int threadId);

class TClientWorker final {
private:
    const TCaptureOptions& Options_;
private:
    using TNetworkCaputreFrameWriter = std::unique_ptr<grpc::ClientWriter<NNetworkCaptureFrame::TNetworkCaptureFrame>>;
    using TNetworkCaputreFrameWriterPtr = grpc::ClientWriter<NNetworkCaptureFrame::TNetworkCaptureFrame>*;
public:
    explicit TClientWorker(const TCaptureOptions& opts)
        : Options_(opts)
    {}

    void Start(size_t threadId);
    void PktCapture(TNetworkCaputreFrameWriterPtr writer, const struct pcap_pkthdr *pkthdr, const u_char *packet);
};

using TPacketsWriterQueue = NThreading::TManyManyQueue<std::unique_ptr<NNetworkCaptureFrame::TNetworkCaptureFrame>>;

class TPacketsServerWriter final {
private:
    TPacketsWriterQueue* PacketsQueue_ = nullptr;
    TPcapDump* DumpFile_ = nullptr;
public:
    TPacketsServerWriter(TPacketsWriterQueue* packetsQueue, TPcapDump* dumpFile)
        : PacketsQueue_(packetsQueue)
        , DumpFile_(dumpFile)
    {}

    void Start(grpc::Server* grpcServer);
};

class TNetworkCaptureRouterImpl final
    : public NNetworkCaptureFrame::TNetworkCaptureRouter::Service {
private:
    TPacketsWriterQueue* PacketsQueue_ = nullptr;
    std::atomic_flag SpinLock_ = ATOMIC_FLAG_INIT;
private:
    grpc::Status CaptureStream(
        grpc::ServerContext* context,
        grpc::ServerReader<NNetworkCaptureFrame::TNetworkCaptureFrame>* stream,
        NNetworkCaptureFrame::TEmpty* empty) override;
public:
    explicit TNetworkCaptureRouterImpl(TPacketsWriterQueue* packetsQueue)
        : PacketsQueue_(packetsQueue)
    {}
};

class TSafeThread {
private:
    std::thread Thread_;
public:
    TSafeThread(TSafeThread& th) = delete;
    TSafeThread& operator=(TSafeThread& th) = delete;
    TSafeThread(TSafeThread&& th) = default;

    explicit TSafeThread(std::thread&& th)
        : Thread_(std::move(th))
    {}

    ~TSafeThread() {
        if (Thread_.joinable())
            Thread_.join();
    }
};
