#pragma once

#include <contrib/libs/rapidjson/include/rapidjson/stringbuffer.h>
#include <contrib/libs/rapidjson/include/rapidjson/prettywriter.h>
#include <contrib/libs/rapidjson/include/rapidjson/writer.h>
#include <contrib/libs/sparsehash/src/sparsehash/dense_hash_map>

#include <util/datetime/base.h>

#include <arpa/inet.h>
#include <mutex>

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

using TPacketHeader = struct pcap_pkthdr;

struct TCounterOptions {
    TString iface;
    TString filter;
    size_t threadCount = 0;
    int bufferSize = 0;
    int snapSize = 0;
    ui32 displayTimeout = 0;
    bool prettyPrint = false;
};

void RunMain(int argc, char** argv);
void RunMainCounter(const TCounterOptions& options);

/* Statistics structures */
typedef struct {
    unsigned char srcIp[16];
    unsigned char dstIp[16];
    ui16 srcPort;
    ui16 dstPort;

    ui64 synCount;
    ui64 ackCount;
    ui64 finCount;
    ui64 rstCount;

    TInstant start;
    TInstant end;
} TConnectionHolder;

using TConnectionHolderMap = google::dense_hash_map<ui64, TConnectionHolder>;

class TAtomicLock {
private:
    std::atomic_flag Lock_ = ATOMIC_FLAG_INIT;
public:
    void lock() {
        while (Lock_.test_and_set(std::memory_order_acquire)) {
        };
    }

    void unlock() {
        Lock_.clear(std::memory_order_release);
    }

    TAtomicLock() = default;
    TAtomicLock(const TAtomicLock& lock_) = delete;
    TAtomicLock(const TAtomicLock&& lock_) = delete;
};

class TStatisticsHolder {
private:
    bool PrettyPrint_;
    TConnectionHolderMap Map_;
    mutable TAtomicLock Lock_;
private:
    using TJsonWriterPretty = rapidjson::PrettyWriter<rapidjson::StringBuffer>;
    using TJsonWriter = rapidjson::Writer<rapidjson::StringBuffer>;

    template <typename T>
    TString InnerJsonDump() {
        rapidjson::StringBuffer sb;
        T json(sb);

        json.StartArray();

        char addrBuff[46]{0};
        TString strBuff;
        for (auto pair = Map_.begin(); pair != Map_.end(); ++pair) {
            json.StartObject();

            inet_ntop(AF_INET6, &pair->second.srcIp, addrBuff, sizeof(addrBuff));
            json.String("src_ip");
            json.String(addrBuff);

            inet_ntop(AF_INET6, &pair->second.dstIp, addrBuff, sizeof(addrBuff));
            json.String("dst_ip");
            json.String(addrBuff);

            json.String("src_port");
            json.Uint(ntohs(pair->second.srcPort));

            json.String("dst_port");
            json.Uint(ntohs(pair->second.dstPort));

            json.String("ack_count");
            json.Uint64(pair->second.ackCount);

            json.String("fin_count");
            json.Uint64(pair->second.finCount);

            json.String("rst_count");
            json.Uint64(pair->second.rstCount);

            json.String("syn_count");
            json.Uint64(pair->second.synCount);

            json.String("time");
            json.Uint64((pair->second.end - pair->second.start).MilliSeconds());

            json.EndObject();
        }

        json.EndArray();

        return sb.GetString();
    }
public:
    TStatisticsHolder()
        : PrettyPrint_(true)
    {
        Map_.set_empty_key(0);
    }

    explicit TStatisticsHolder(bool prettyPrint)
        : PrettyPrint_(prettyPrint)
    {
        Map_.set_empty_key(0);
    }

    bool HasKey(ui64 key) const {
        std::lock_guard<TAtomicLock> lock(Lock_);
        return Map_.find(key) != Map_.end();
    }

    void FillData(ui64 key, const unsigned char* srcIp, ui16 srcPort, const unsigned char* dstIp, ui16 dstPort) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        Map_[key].srcPort = srcPort;
        Map_[key].dstPort = dstPort;
        std::memcpy(Map_[key].srcIp, srcIp, 16);
        std::memcpy(Map_[key].dstIp, dstIp, 16);
    }

    void IncAck(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        ++Map_[key].ackCount;
    }
    void IncFin(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        ++Map_[key].finCount;
    }
    void IncRst(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        ++Map_[key].rstCount;
    }
    void IncSyn(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        ++Map_[key].synCount;
    }

    void SetTimeStart(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        Map_[key].start = TInstant::Now();
    }

    void SetTimeEnd(ui64 key) {
        std::lock_guard<TAtomicLock> lock(Lock_);
        Map_[key].end = TInstant::Now();
    }

    TString JsonDump() {
        std::lock_guard<TAtomicLock> lock(Lock_);

        if (PrettyPrint_)
            return InnerJsonDump<TJsonWriterPretty>();
        else
            return InnerJsonDump<TJsonWriter>();
    }
};

class TClientWorker final {
private:
    const TString& Iface_;
    const TString& Filter_;
    int SnapSize_;
    int BufferSize_;
    TStatisticsHolder* StatsHolder_;
public:
    TClientWorker(const TCounterOptions& options,
                  TStatisticsHolder* statsHolder)
        : Iface_(options.iface)
        , Filter_(options.filter)
        , SnapSize_(options.snapSize)
        , BufferSize_(options.bufferSize)
        , StatsHolder_(statsHolder){};

    void Start(size_t threadId);
    void StatsObserver(ui32 timeout);

private:
    void ProcessPacket(const TPacketHeader* pktHdr, const unsigned char* packet, ptrdiff_t etherOffset);
};
