#pragma once

#include "main.h"

#include <library/cpp/digest/murmur/murmur.h>
#include <mapreduce/yt/interface/client.h>

#include <unordered_set>
#include <unordered_map>

class TMultiFrame;

using TPcapFileHeader = struct pcap_file_header;
using TPcapPacketHeader = struct pcap_pkthdr;

struct IPcapOptions {
    virtual void SetHeader(TPcapFileHeader header) {
        Y_UNUSED(header);
    }

    virtual void SetLinkType(int type) {
        Y_UNUSED(type);
    }

    virtual void SetSnapSize(int size) {
        Y_UNUSED(size);
    }

    virtual int LinkType() const {
        ythrow yexception() << "Not implemented LinkType";
    }

    virtual int SnapSize() const {
        ythrow yexception() << "Not implemented SnapSize";
    }

    virtual const TString& OutputDir() const {
        ythrow yexception() << "Not implemented OutputDir";

    }

    virtual int Depth() const {
        ythrow yexception() << "Not implemented Depth";
    }

    virtual const TString& TableName() const {
        ythrow yexception() << "Not implemented TableName";
    }

    virtual TPcapFileHeader FileHeader() const {
        ythrow yexception() << "Not implemented FileHeader";
    }

    virtual NYT::ITransactionPtr YtClient() {
        ythrow yexception() << "Not implemented YtClient";
    }

    virtual ~IPcapOptions() {}
};

struct TFilePcapOptions: public IPcapOptions {
private:
    const TString& OutpitDir;
    int Depth_ = 0;
    int LinkType_ = 0;
    int SnapSize_ = 0;
public:
    TFilePcapOptions(const TString& outputDir, int depth)
        : OutpitDir(outputDir)
        , Depth_(depth){};

    void SetLinkType(int type) override {
        LinkType_ = type;
    }

    void SetSnapSize(int size) override {
        SnapSize_ = size;
    }

    int LinkType() const override {
        return LinkType_;
    }

    int SnapSize() const override {
        return SnapSize_;
    }

    const TString& OutputDir() const override {
        return OutpitDir;
    }

    int Depth() const override {
        return Depth_;
    }
};

struct TYtPcapOptions: public IPcapOptions {
private:
    NYT::ITransactionPtr YtClient_;
    const TString& TableName_;
    TPcapFileHeader FileHeader_{};
public:
    TYtPcapOptions(NYT::ITransactionPtr ytClient, const TString& tableName)
        : YtClient_(ytClient)
        , TableName_(tableName){};

    void SetHeader(TPcapFileHeader header) override {
        FileHeader_ = header;
    }

    const TString& TableName() const override {
        return TableName_;
    }

    TPcapFileHeader FileHeader() const override {
        return FileHeader_;
    }

    NYT::ITransactionPtr YtClient() override {
        return YtClient_;
    }
};

struct TDumpMessage {
public:
    const TPcapPacketHeader* PktHdr;
    const unsigned char* Pkt;
    const TMultiFrame& Frame;
public:
    TDumpMessage(const TPcapPacketHeader* pktHdr,
                 const unsigned char* packet,
                 const TMultiFrame& frame)
        : PktHdr(pktHdr)
        , Pkt(packet)
        , Frame(frame)
    {}
};

class IPcapDumper {
public:
    virtual ~IPcapDumper() {}
    virtual void Finish() {}
    virtual void Write(const TDumpMessage& message) = 0;

    ui64 CalcHash(const unsigned char* srcIp, const unsigned char* dstIp, ui16 srcPort, ui16 dstPort) {
        TMurmurHash2A<ui64> hasher;
        hasher.Update(srcIp, 16);
        hasher.Update(&srcPort, 2);
        hasher.Update(dstIp, 16);
        hasher.Update(&dstPort, 2);

        return hasher.Value();
    }
};
