#include "yt_dumper.h"

#include <util/folder/path.h>

void TYtPcapDumper::Write(const TDumpMessage& message) {
    ui64 hash = CalcHash(message.Frame.SrcAdr(), message.Frame.DstAdr(),
                         message.Frame.SrcPort(), message.Frame.DstPort());
    ui64 hashInv = CalcHash(message.Frame.DstAdr(), message.Frame.SrcAdr(),
                            message.Frame.DstPort(), message.Frame.SrcPort());

    Dump(message, hash, hashInv);
}

void TYtPcapDumper::Dump(const TDumpMessage& message, ui64 hash, ui64 hashInv) {
    NYT::TNode node;

    struct pcap_sf_pkthdr sf_hdr;
    sf_hdr.ts.tv_sec = message.PktHdr->ts.tv_sec;
    sf_hdr.ts.tv_usec = message.PktHdr->ts.tv_usec;
    sf_hdr.caplen = message.PktHdr->caplen;
    sf_hdr.len = message.PktHdr->len;

    /* PCAP always stores data in 32bit format */
    size_t structSize = sizeof(sf_hdr);
    std::vector<char> buf(structSize + message.PktHdr->caplen);
    std::memcpy(&buf[0], &sf_hdr, structSize);
    std::memcpy(&buf[0] + structSize, message.Pkt, message.PktHdr->caplen);

    TString data;
    std::copy(buf.begin(), buf.end(), std::back_inserter(data));

    if (Set_.find(hashInv) != Set_.end()) {
        Set_.insert(hashInv);
        node["hash"] = hashInv;
    } else {
        Set_.insert(hash);
        node["hash"] = hash;
    }

    node["data"] = data;
    node["order"] = OrderCounter_++;
    YtWriter_->AddRow(node);
}

void TYtPcapDumper::Finish() {
    YtWriter_->Finish();
}

NYT::TNode TYtPcapDumper::CreateTableSchema() {
    NYT::TNode hash;
    hash["name"] = "hash";
    hash["type"] = "uint64";

    NYT::TNode order;
    order["name"] = "order";
    order["type"] = "uint64";

    NYT::TNode data;
    data["name"] = "data";
    data["type"] = "string";

    NYT::TNode schema;
    schema.Add(hash);
    schema.Add(order);
    schema.Add(data);

    return schema;
}

void TYtPcapDumper::WriteHeader(NYT::ITransactionPtr ytClient, const TString& fileName, const TPcapFileHeader& fileHeader) {
    auto writer = ytClient->CreateFileWriter(fileName);
    writer->Write((void*)&fileHeader, sizeof(fileHeader));
    writer->Finish();
}

TYtPcapDumper::TYtPcapDumper(IPcapOptions* options) {
    NYT::ITransactionPtr YtClient = options->YtClient();
    TString fileName = options->TableName() + "-header";

    if (YtClient->Exists(fileName)) {
        YtClient->Abort();
        ythrow yexception() << "Headers file already exists\n";
    } else {
        if (YtClient->Exists(options->TableName())) {
            YtClient->Abort();

            ythrow yexception() << "Temp table ["
                                << options->TableName()
                                << "] already exists\n";
        } else {
            WriteHeader(YtClient, fileName, options->FileHeader());

            YtClient->Create(options->TableName(), NYT::NT_TABLE,
                                   NYT::TCreateOptions()
                                       .IgnoreExisting(true)
                                       .Recursive(true)
                                       .Attributes(NYT::TNode()("schema", CreateTableSchema())("replication_factor", 1)));

            YtWriter_ = YtClient->CreateTableWriter<NYT::TNode>(options->TableName());
        }
    }
}

TYtPcapDumper::~TYtPcapDumper() {
    if (YtWriter_ != nullptr) {
        try {
            YtWriter_->Finish();
        } catch(const yexception& e) {
            Cerr << e.what() << "\n";
        }
    }
}
