#include "main.h"

#include <util/system/thread.h>

#include <atomic>
#include <deque>

#if defined(_linux_)
#   include <linux/if_packet.h>
#   include <sys/socket.h>
#endif

std::atomic_bool CAPTURE_ENABLED{true};
std::deque<std::atomic_bool> STATISTICS_BITS;

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

void RunMainClient(const TCaptureOptions &options) {
    SetExitHandler();
    SetStatsHandler();

    for (size_t i = 0; i < options.threadCount; ++i) {
        STATISTICS_BITS.emplace_back();
    }

    if (options.iface.empty())
        ythrow yexception() << "Empty interface";


    std::vector<TSafeThread> workers;
    auto clientWorker = std::make_unique<TClientWorker>(options);

    for (size_t i = 0; i < options.threadCount; ++i) {
        std::thread th([&, i] {
            clientWorker->Start(i);
        });
        workers.emplace_back(std::move(th));
    }

    while (CAPTURE_ENABLED.load()) {
        sleep(1);
    }
    /*std::for_each(workers.begin(), workers.end(),
                  [] (std::thread& t) {
                      t.join();
                  });*/

}

void TClientWorker::Start(size_t threadId) {
    TStringStream threadName;
    threadName << "Capture worker thread " << threadId;
    TThread::SetCurrentThreadName(threadName.Str().c_str());

    int ret = 0;
#if defined(_linux_)
    cpu_set_t cpuset;
    pthread_t thread = pthread_self();
    CPU_ZERO(&cpuset);
    CPU_SET(threadId, &cpuset);

    ret = pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset);
    if (ret != 0) {
        ythrow yexception() << "Could not set CPU affinity for thread: " << threadId;
    }
#endif
    std::chrono::system_clock::time_point deadline =
        std::chrono::system_clock::now() + std::chrono::seconds(Options_.timeout);

    grpc::ClientContext context;
    NNetworkCaptureFrame::TEmpty empty;

    auto grpcChannel = grpc::CreateChannel(Options_.remoteServer, grpc::InsecureChannelCredentials());
    if (!grpcChannel->WaitForConnected(deadline)) {
        ythrow yexception() << "GRPC connection timeout: " << Options_.remoteServer;
    }
    auto grpcStub = NNetworkCaptureFrame::TNetworkCaptureRouter::NewStub(grpcChannel);
    TNetworkCaputreFrameWriter grpcWriter(
        grpcStub->CaptureStream(&context, &empty));

    TPcap pcap(Options_.iface, Options_.bufferSize, Options_.snapSize);
    TPcapFilter pcapFilter(pcap.Session(), Options_.filter);

#if defined(_linux_)
    int val = (getpid() & 0xffff) | (PACKET_FANOUT_LB << 16);
    if (setsockopt(pcap_fileno(pcap.Session()), SOL_PACKET, PACKET_FANOUT, &val, sizeof(val)) < 0) {
        Cerr << "Failed to set fanout: "
             << strerror(errno)
             << "\n";
        return;
    }
#endif

    TPacketHeader* pktHdr = nullptr;
    const unsigned char* pkt = nullptr;

    bool unlimited = Options_.numberOfPackets == 0;
    uint64_t counter = 0;

    while ((ret = pcap.NextEx(&pktHdr, &pkt)) >= 0 && CAPTURE_ENABLED.load()
        && (unlimited || counter < Options_.numberOfPackets)) {
        /* ret == 0 is timeout and pktHdr, pkt invalid pointers */
        if (ret == 0)
            usleep(100000);
        else {
            PktCapture(grpcWriter.get(), pktHdr, pkt);
            ++counter;
        }

        if (STATISTICS_BITS[threadId].load()) {
            PrintStats(pcap.Stats(), counter, threadId);
            STATISTICS_BITS[threadId].store(false, std::memory_order_relaxed);
        }
    }

    if (ret == -1) {
        ythrow yexception() << "Failed to capture packets: "
                            << pcap.ErrorString() << "\n";
    } else if (ret == -2) {
        Cout << "Capture loop is ended\n";
    }

    grpcWriter->WritesDone();
    grpc::Status status = grpcWriter->Finish();

    if (!status.ok()) {
        ythrow yexception() << "Error sending data to server: " << status.error_message();
    }

    PrintStats(pcap.Stats(), counter, threadId);
}

void TClientWorker::PktCapture(TNetworkCaputreFrameWriterPtr writer, const TPacketHeader *pkthdr, const u_char *packet) {
    NNetworkCaptureFrame::TNetworkCaptureFrame captureFrame;

    captureFrame.set_tssec(pkthdr->ts.tv_sec);
    captureFrame.set_tsusec(pkthdr->ts.tv_usec);
    captureFrame.set_caplen(pkthdr->caplen);
    captureFrame.set_len(pkthdr->len);
    captureFrame.set_packet(packet, pkthdr->caplen);

    if (!writer->Write(captureFrame)) {
        ythrow yexception() << "Error sending data to server: writer failed";
    }
}

static inline void PrintStats(const TString& str, uint64_t counter, int threadId) {
    TStringStream ss;

    ss << "Stats for thread: " << threadId << "\n";
    ss << "Packets captured: " << counter << "\n";
    ss << str << "\n";

    Cerr << ss.Str();
}

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

    for (auto &i : STATISTICS_BITS) {
        i.store(true, std::memory_order_relaxed);
    }
};

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

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

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

    act.sa_handler = StatsHandler;
    act.sa_flags = 0;
    sigemptyset(&act.sa_mask);
    sigaddset(&act.sa_mask, SIGUSR1);

    sigaction(SIGUSR1, &act, NULL);
}

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);
}
