#include <balancer/serval/contrib/cone/mun.h>
#include <balancer/serval/core/log.h>

#include <contrib/libs/lz4/lz4.h>
#include <library/cpp/getopt/small/last_getopt.h>
#include <library/cpp/eventlog/events_extension.h>

#include <util/datetime/base.h>
#include <util/generic/deque.h>
#include <util/generic/hash.h>
#include <util/generic/hash_set.h>
#include <util/generic/scope.h>
#include <util/memory/blob.h>
#include <util/stream/str.h>
#include <util/system/unaligned_mem.h>

namespace {
    bool Read(google::protobuf::io::CodedInputStream* in, google::protobuf::MessageLite* m) {
        int size;
        if (!in->ReadVarintSizeAsInt(&size))
            return false;
        auto limit = in->PushLimit(size);
        Y_DEFER { in->PopLimit(limit); };
        return m->ParseFromCodedStream(in);
    }

    bool Skip(google::protobuf::io::CodedInputStream* in) {
        int size;
        return in->ReadVarintSizeAsInt(&size) && in->Skip(size);
    }

    auto Stream(TStringBuf buf) {
        return google::protobuf::io::CodedInputStream(reinterpret_cast<const ui8*>(buf.data()), buf.size());
    }

    // If the provided pointer is in the middle of a valid sync block, return a pointer
    // to the first event; otherwise, locate the first valid sync block after it, or
    // return a pointer to the end of the file if there is none.
    const char* HeadAt(TStringBuf buf, const char* p) {
        auto unwindToHead = [buf](size_t i) -> const char* {
            return i > 0 && i >= (ui8)buf[i - 1] ? buf.data() + i - (ui8)buf[i - 1] : nullptr;
        };
        auto prev = buf.rfind(NSv::TLog::SyncField, p - buf.data());
        auto next = buf.find(NSv::TLog::SyncField, p - buf.data() + 1);
        while (next != TStringBuf::npos && !unwindToHead(next))
            next = buf.find(NSv::TLog::SyncField, next + 1);
        auto prevHead = prev != TStringBuf::npos ? unwindToHead(prev) : nullptr;
        auto nextHead = next != TStringBuf::npos ? unwindToHead(next) : buf.end();
        // `nextHead <= p` if `p` points to the sync field of a valid block or before it.
        // `prevHead` is non-null if `p` is in a valid block, but after the sync field.
        return nextHead > p && prevHead ? prevHead : nextHead;
    }

    // Given a pointer to an arbitrary byte of the log, return the time of the first
    // event in the sync block it belongs to, or max TInstant if at the end.
    TInstant TimeAt(TStringBuf buf, const char* p) {
        auto s = Stream({HeadAt(buf, p), buf.end()});
        for (NSv::TLogEvent ev; Read(&s, &ev); )
            if (ev.TimeSize() || ev.GetLZ4Time())
                return ev.TimeSize() ? TInstant::MicroSeconds(ev.GetTime(0)) : TInstant::MicroSeconds(ev.GetLZ4Time());
        return TInstant::Max();
    }

    // Locate a region in which (probably nearly all of) the events between the two
    // timestamps are located.
    TStringBuf TimeBounded(TStringBuf buf, TInstant a, TInstant b) {
        auto s = HeadAt(buf, buf.begin());
        struct mun_vec(const char) vec = mun_vec_init_borrow(s, static_cast<size_t>(buf.end() - s));
        auto p = s + mun_vec_bisect(&vec, a <= TimeAt(buf, _));
        auto q = s + mun_vec_bisect(&vec, b < TimeAt(buf, _));
        return {p != s ? HeadAt(buf, p - 1) : p, q};
    }

    struct TTree {
        TString Head;
        TDeque<THolder<TTree>> Children;
        ui32 Level = 0;
        bool Finished = false;

        bool PrintReady() {
            if (Head) {
                Cout << Head << Endl;
                Head.clear();
            }
            for (; Children; Children.pop_front())
                if (!Children.front()->PrintReady())
                    return false;
            return Finished;
        }
    };

    template <typename F>
    bool ReadEvents(google::protobuf::io::CodedInputStream* in, F&& f) {
        NSv::TLogEvent ev;
        while (true) {
            while (!Read(in, &ev)) {
                const char* data; int size;
                in->GetDirectBufferPointerInline((const void**)&data, &size);
                // Skip 1 byte to avoid an infinite loop. We're not at a sync block boundary anyway.
                if (!size || !in->Skip(HeadAt({data + 1, (size_t)size - 1}, data + 1) - data))
                    return in->ExpectAtEnd();
            }
            Y_ENSURE(ev.TimeSize() == ev.TypeSize() && ev.TimeSize() == ev.ContextSize() &&
                     ev.TimeSize() == ev.CurrentSize(), "broken log: TLogEvent field length mismatch");
            for (size_t i = 0; i < ev.TimeSize(); i++)
                f(in, TInstant::MicroSeconds(ev.GetTime(i)), ev.GetType(i), ev.GetContext(i), ev.GetCurrent(i));
            if (ev.GetLZ4()) {
                TString s(ev.GetLZ4DecompressedSize(), '\0');
                Y_ENSURE(LZ4_decompress_fast(ev.GetLZ4().begin(), s.begin(), s.size()) >= 0, "malformed LZ4 data");
                auto substream = Stream(s);
                Y_ENSURE(ReadEvents(&substream, f), "truncated LZ4 data");
            }
        }
    }
}

template <>
TInstant FromStringImpl<TInstant, char>(const char* data, size_t size) {
    ui64 ui;
    double fp;
    return TryFromString(TStringBuf{data, size}, ui) ? ui > 1000000000000ull ? TInstant::MicroSeconds(ui) : TInstant::Seconds(ui)
         : TryFromString(TStringBuf{data, size}, fp) ? TInstant::MicroSeconds(fp * 1000000)
         : TInstant::ParseIso8601({data, size});
}

void ParseEventIds(TStringBuf b, THashSet<ui32>* out) {
    ui32 asInt;
    for (TStringBuf tok : StringSplitter(b).Split(',').SkipEmpty())
        out->emplace(TryFromString(tok, asInt) ? asInt : NProtoBuf::TEventFactory::Instance()->IdByName(tok));
}

int main(int argc, const char **argv) {
    TInstant start = TInstant::Zero();
    TInstant end = TInstant::Max();
    THashSet<ui32> includeSet;
    THashSet<ui32> excludeSet;
    THashSet<ui32> trees;
    bool printDeps = false;

    NLastGetopt::TOpts opts = NLastGetopt::TOpts::Default();
    opts.SetTitle(
        "Streamed-tree event log parser. Dumps data in a TSV format: time, parent id,\n"
        "subtree id, event type, arguments. For leaf events, subtree id is 0; for roots,\n"
        "parent id is 0. Unknown event types are denoted as `?N' with no arguments.");
    opts.AddLongOption('s', "start-time",
        "Left bound (inclusive) of the interval the events in which to output. Formats:\n"
        "  * ISO 8601 (%Y-%m-%dT%H:%M:%S.%uZ, %Y-%m-%dT%H:%M:%SZ, or %Y-%m-%d);\n"
        "  * Unix (floating point seconds, integer seconds, integer microseconds).\n"
        "Note: if events are severely out-of-order, some of them may be dropped\n"
        "(e.g. if `PushAt` was used to insert a frame far in the past/future).")
        .RequiredArgument("TIME").StoreResult(&start);
    opts.AddLongOption('e', "end-time",
        "Right bound (exclusive) of the interval; see above.")
        .RequiredArgument("TIME").StoreResult(&end);
    opts.AddLongOption('n', "frame-id",
        "Only output the event with the specified id and its descendants.\n"
        "Note: when combined with -s, if the event itself is not in the interval,\n"
        "but some of its descendants are, they may not be printed.")
        .RequiredArgument("TREE").Handler1T<TStringBuf>([&](TStringBuf b) { StringSplitter(b).Split(',').ParseInto(&trees); });
    opts.AddLongOption('i', "event-list",
        "Comma-separated list of event types to output, as names or IDs.")
        .RequiredArgument("TYPES").Handler1T<TStringBuf>([&](TStringBuf b) { ParseEventIds(b, &includeSet); });
    opts.AddLongOption('x', "exclude-list",
        "Comma-separated list of event types to ignore.")
        .RequiredArgument("TYPES").Handler1T<TStringBuf>([&](TStringBuf b) { ParseEventIds(b, &excludeSet); });
    opts.AddLongOption('t', "tree",
        "Print the children of each event directly after it, and visualize the tree.\n"
        "Note: when combined with -s, -i, or -x, events that are children of those\n"
        "that are not shown are displayed as roots.")
        .NoArgument().StoreValue(&printDeps, true);
    opts.SetFreeArgsNum(1);
    opts.SetFreeArgTitle(0, "<file>", "A log file written by serval");
    NLastGetopt::TOptsParseResult optParse(&opts, argc, argv);
    bool allExcluded = includeSet && std::all_of(includeSet.begin(), includeSet.end(), [&](ui32 i) {
        return excludeSet.contains(i); });
    Y_ENSURE(!allExcluded, "--event-list/-i fully contained within --exclude-list/-x");
    Y_ENSURE(start < end, "--start-time/-s is not before --end-time/-e");

    TTree root;
    THashSet<ui32> toOutput;
    THashMap<ui32, TTree*> byId;

    TBlob file = TBlob::FromFile(optParse.GetFreeArgs()[0]);
    auto in = Stream(TimeBounded(TStringBuf{file.AsCharPtr(), file.Size()}, start, end));
    auto factory = NProtoBuf::TEventFactory::Instance();
    bool eof = ReadEvents(&in, [&](google::protobuf::io::CodedInputStream* in, TInstant t, ui32 type, ui32 parentId, ui32 id) {
        bool write = start <= t && t < end && (!includeSet || includeSet.contains(type)) && !excludeSet.contains(type);
        if (trees) {
            // TODO if there is nothing to output, skip forward while the next sync block
            //      only contains trees with id less than needed.
            if (!trees.contains(id) && !toOutput.contains(parentId))
                write = false;
            else if (type == 0) // TFrameEnd
                toOutput.erase(parentId);
            else if (id != 0)
                toOutput.emplace(id);
        }
        auto parent = &root;
        if (printDeps && (write || type == 0)) {
            auto it = byId.find(parentId);
            if (it != byId.end()) {
                parent = it->second;
                if (type == 0) {
                    parent->Finished = true;
                    byId.erase(it);
                }
            }
        }
        if (write) {
            auto node = MakeHolder<TTree>();
            TStringOutput out(node->Head);
            out << t << '\t' << parentId << '\t' << id << '\t';
            for (size_t i = parent->Level; i --> 0;)
                out << (i ? "| " : type == 0 ? "\\-" : "+-");
            if (THolder<NProtoBuf::Message> msg{factory->CreateEvent(type)}) {
                Y_ENSURE(Read(in, msg.Get()), "truncated log: could not read payload of " << msg->GetDescriptor()->name());
                factory->PrintEvent(type, msg.Get(), out);
            } else {
                Y_ENSURE(Skip(in), "truncated log: could not skip payload of unknown type " << type);
                out << '?' << type;
            }
            node->Level = parent->Level + 1;
            node->Finished = !(printDeps && id);
            if (!node->Finished)
                byId.emplace(id, node.Get());
            parent->Children.emplace_back(std::move(node));
        } else {
            Y_ENSURE(Skip(in), "truncated log: could not skip filtered event payload");
        }
        if (write || parent != &root)
            root.PrintReady();
    });
    for (auto& tree : byId)
        tree.second->Finished = true;
    root.PrintReady();
    Y_ENSURE(eof, "truncated log: could not decode event head");
    return 0;
}
