#pragma once

#include <util/generic/yexception.h>
#include <util/generic/string.h>
#include <util/stream/str.h>
#include <util/string/cast.h>

#include <contrib/libs/libpcap/pcap.h>
#include <contrib/libs/libpcap/pcap-int.h>

#include <memory>

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

/*
 * Standard libpcap format.
 */
constexpr ui32 TCPDUMP_MAGIC = 0xa1b2c3d4;

/*
 * Normal libpcap format, except for seconds/nanoseconds timestamps,
 * as per a request by Ulf Lamping <ulf.lamping@web.de>
 */
constexpr ui32 NSEC_TCPDUMP_MAGIC = 0xa1b23c4d;

class TPcapFilter {
private:
    using TBPFProgram = struct bpf_program;
    std::unique_ptr<TBPFProgram, void (*)(TBPFProgram*)> Filter_{new TBPFProgram{0, nullptr}, PcapFilterCloser};
private:
    static void PcapFilterCloser(TBPFProgram* fd) {
        if (fd != nullptr) {
            pcap_freecode(fd);
            delete fd;
        }

    }
public:
    TPcapFilter(const TPcapFilter& that) = delete;
    TPcapFilter& operator=(const TPcapFilter& that) = delete;

    TPcapFilter(pcap_t* session, const TString& filter) {
        if (pcap_compile(session, Filter_.get(), filter.c_str(), 0, 0) == -1) {
            ythrow yexception() << "Failed to parse filter ["
                                << filter << "]: "
                                << pcap_geterr(session);
        } else {
            if (pcap_setfilter(session, Filter_.get()) == -1) {
                ythrow yexception() << "Failed to install filter ["
                                    << filter << "]: "
                                    << pcap_geterr(session);
            }
        }
    }
};

class TPcapIfaceList {
private:
    std::unique_ptr<pcap_if_t, void (*)(pcap_if_t*)> List_{nullptr, PcapIfaceListCloser};
    char ErrorBuffer_[PCAP_ERRBUF_SIZE]{0};
private:
    static void PcapIfaceListCloser(pcap_if_t* fd) {
        pcap_freealldevs(fd);
    };
public:
    TPcapIfaceList(const TPcapIfaceList& that) = delete;
    TPcapIfaceList& operator=(const TPcapIfaceList& that) = delete;

    TPcapIfaceList() {
        pcap_if_t *devicesList;

        if (pcap_findalldevs(&devicesList, ErrorBuffer_) == -1) {
            ythrow yexception() << "Failed to get device list: "
                                << ErrorBuffer_;
        }

        List_.reset(devicesList);
    };

    TString AsString() const {
        TStringStream result;
        pcap_if_t *loopList;

        for (loopList = List_.get(); loopList != nullptr; loopList = loopList->next) {
            result << loopList->name << " ";
        }

        return result.Str();
    }
};

class TPcapDump {
private:
    std::unique_ptr<pcap_dumper_t, void (*)(pcap_dumper_t*)> DumpFile_{nullptr, PcapDumpCloser};
private:
    static void PcapDumpCloser(pcap_dumper_t* fd) {
        if (fd != nullptr) {
            pcap_dump_flush(fd);
            pcap_dump_close(fd);
        }
    }
public:
    TPcapDump(const TPcapDump& that) = delete;
    TPcapDump& operator=(const TPcapDump& that) = delete;

    TPcapDump(pcap_t *session, const TString& filePath) {
        pcap_dumper_t* dumpFile = pcap_dump_open(session, filePath.c_str());

        if (dumpFile == nullptr) {
            ythrow yexception() << "Failed to open file [" << filePath << "]: "
                                << pcap_geterr(session);
        }

        DumpFile_.reset(dumpFile);
    }

    void Dump(TPcapPacketHeader* pktHdr, const unsigned char* pkt) {
        pcap_dump((u_char *)DumpFile_.get(), pktHdr, pkt);
        pcap_dump_flush(DumpFile_.get());
    }
};

class TPcap {
private:
    std::unique_ptr<pcap_t, void (*)(pcap_t*)> Session_{nullptr, PcapCloser};
    TStringStream ErrorString_;
    ptrdiff_t Offset_ = 0;

    char ErrorBuffer_[PCAP_ERRBUF_SIZE]{0};
private:
    static void PcapCloser(pcap_t* fd) {
        pcap_close(fd);
    }

    void UpdateOffSet() {
        switch (LinkType()) {
            case DLT_NULL: // loopback bsd interface
                Offset_ = 4;
                break;
            case DLT_EN10MB: // eth interface
                Offset_ = 14;
                break;
            case DLT_LINUX_SLL: // Linux "any" interface
                Offset_ = 16;
                break;
            default:
                ythrow yexception() << "Offset for link layer ["
                                    << LinkType()
                                    << "] is not defined";
        }
    }
public:
    TPcap(const TPcap& that) = delete;
    TPcap& operator=(const TPcap& that) = delete;

    TPcap(int linkType, int snapSize) {
        pcap_t* session = pcap_open_dead(linkType, snapSize);

        if (session == nullptr) {
            ythrow yexception() << "Failed to create dead interface: "
                                << ErrorBuffer_;
        } else {
            Session_.reset(session);

            UpdateOffSet();
        }
    }

    TPcap(const TString& filePath) {
        pcap_t* session = pcap_open_offline(filePath.c_str(), ErrorBuffer_);

        if (session == nullptr) {
            ythrow yexception() << "Failed to open file ["
                                << filePath << "]: "
                                << ErrorBuffer_;
        } else {
            Session_.reset(session);

            UpdateOffSet();
        }
    }

    TPcap(const TString& iface, int bufferSize, int snapSize) {
        bpf_u_int32 ifaceMask, ifaceAddr;

        if (pcap_lookupnet(iface.c_str(), &ifaceAddr, &ifaceMask, ErrorBuffer_) == -1) {
            Cerr << "Failed to get netmask for device ["
                                << iface << "]: "
                                << ErrorBuffer_
                                << "\n";
        }

        pcap_t* session = pcap_create(iface.c_str(), ErrorBuffer_);
        if (session == nullptr) {
            TPcapIfaceList list;

            ythrow yexception() << "Failed to open device ["
                  << iface << "]: "
                  << ErrorBuffer_
                  << "\nAvailable interfaces: "
                  << list.AsString();
        }

        Session_.reset(session);

        if (pcap_set_snaplen(Session_.get(), snapSize) < 0) {
            ythrow yexception() << "Could not set snap size to: " << snapSize << "\n";
        }
        if (pcap_set_promisc(Session_.get(), 0) < 0) {
            ythrow yexception() << "Could not start promisc mode";
        }
        if (pcap_set_timeout(Session_.get(), 1000) < 0) {
            ythrow yexception() << "Could not set timeout";
        }
        if (pcap_set_immediate_mode(Session_.get(), 1) < 0)
        {
            ythrow yexception() << "Could not set immediate mode";
        }
        if (pcap_set_buffer_size(Session_.get(), bufferSize) < 0) {
            ythrow yexception() << "Could not set buffer size to: " << bufferSize << "\n";
        }
        if (pcap_activate(Session_.get()) < 0) {
            ythrow yexception() << "Could not start capturing";
        }

        UpdateOffSet();
    }

    pcap_t* Session() {
        return Session_.get();
    }

    TPcapFileHeader FileHeader() const {
        TPcapFileHeader hdr;

        hdr.magic = Session_->opt.tstamp_precision == PCAP_TSTAMP_PRECISION_NANO ? NSEC_TCPDUMP_MAGIC : TCPDUMP_MAGIC;
        hdr.version_major = PCAP_VERSION_MAJOR;
        hdr.version_minor = PCAP_VERSION_MINOR;

        hdr.thiszone = Session_->tzoff;
        hdr.snaplen = Session_->snapshot;
        hdr.sigfigs = 0;
        hdr.linktype = Session_->linktype;
        hdr.linktype |= Session_->linktype_ext;

        return hdr;
    }

    ptrdiff_t Offset() const {
        return Offset_;
    }

    int NextEx(TPcapPacketHeader** pktHdr, const unsigned char** pkt) {
        int ret = pcap_next_ex(Session_.get(), pktHdr, pkt);

        if (ret < 0) {
            ErrorString_.Clear();
            ErrorString_ << ToString(pcap_geterr(Session_.get()));
        }

        return ret;
    }

    TString Stats() {
        TStringStream ss;
        struct pcap_stat stats;

        stats.ps_ifdrop = 0;
        if (pcap_stats(Session_.get(), &stats) < 0) {
            ythrow yexception() << "Could not get statistics";
        }

        ss << "Packets recv: " << stats.ps_recv << "\n";
        ss << "Packets dropped: " << stats.ps_drop << "\n";
        if (stats.ps_ifdrop != 0) {
            ss << "Packets dropped by interface: " << stats.ps_ifdrop;
        }

        return ss.Str();
    }

    TStringBuf ErrorString() {
        return ErrorString_.Str();
    }

    int LinkType() const {
        return Session_->linktype;
    }

    int SnapSize() const {
        return Session_->snapshot;
    }
};
