#pragma once

#include "proto_dump.h"

#include <library/cpp/json/json_writer.h>
#include <library/cpp/logger/global/global.h>

#include <robot/jupiter/library/rtdoc/file/reader.h>
#include <robot/jupiter/library/rtdoc/protos/lumps.pb.h>

#include <util/string/hex.h>
#include <util/generic/map.h>
#include <util/generic/set.h>

namespace NRtDoc {
    enum EDumpLevel {
        Blob,
        Row,
        ProtobufTable,
        Max
    };

    class TRtDocDumper {
    private:
        IProtoDumper::TPtr ProtoDumper;
        TSet<TString> Filter;
        EDumpLevel Level = Row;
        bool NoPretty = false;
        bool UseHex = false;

    private:
        struct TOutput {
            // TOutput is a helper that writes the top-level json array.
            // It is done manually because we would like to print rows like a stream
            IOutputStream* Output = nullptr;
            NJson::TJsonWriterConfig JsonOpts;
            bool HasWritten = false;

            inline void Start(IOutputStream& output, bool noPretty);
            inline void Next();
            inline void Finish();
        };

        TOutput Output;

    public:
        void SetNoPretty(bool vl) {
            NoPretty = vl;
        }
        void SetLevel(EDumpLevel l) {
            Level = l;
        }
        void SetPrepIds(const TSet<TString>& filterPrepIds) {
            Filter = filterPrepIds;
        }
        void SetUseHex(bool useHex) {
            UseHex = useHex;
        }

        void Start(IOutputStream& s) {
            if (Level >= EDumpLevel::ProtobufTable) {
                const bool deepParse = Level > EDumpLevel::ProtobufTable;
                ProtoDumper = IProtoDumper::Create(deepParse, UseHex);
            }

            Output.Start(s, NoPretty);
        }

        NJupiter::TMercuryLumps ParseMercuryRow(const TStringBuf& blob) {
            NJupiter::TMercuryLumps row;
            Y_ENSURE(row.ParseFromArray(blob.data(), blob.size()));
            return row;
        }

        void DumpMercuryRow(NJson::TJsonWriter& writer, const TStringBuf& data) {
            NJupiter::TMercuryLumps row = ParseMercuryRow(data);
            TMap<TString, ui32> keysCounts;
            for (const auto& cell : row.get_arr_lumps()) {
                auto p = keysCounts.insert(std::make_pair(cell.GetName(), 1));
                if (!p.second) {
                    ++p.first->second;
                }
            }
            TVector<const NJupiter::TMapType*> values;
            std::transform(row.get_arr_lumps().begin(), row.get_arr_lumps().end(), std::back_inserter(values), [](const NJupiter::TMapType& vl) { return &vl; });
            StableSortBy(values, [](const NJupiter::TMapType* p) { return p->GetName(); });


            auto writeFunc = [this, &writer](const NJupiter::TMapType& cell) {
                if (ProtoDumper && ProtoDumper->IsKnown(cell.GetName())) {
                    NJson::TJsonValue json = ProtoDumper->ToJson(cell);
                    writer.Write(&json);
                } else if (UseHex) {
                    writer.Write(HexEncode(cell.GetData()));
                } else {
                    writer.Write(cell.GetData());
                }
            };

            for (auto it = values.cbegin(); it != values.cend();) {
                const auto& cell = **it;
                const TString cellName = cell.GetName();
                ui32 nValues = keysCounts[cellName];
                if (Filter && !Filter.count(cellName)) {
                    it += nValues;
                    continue;
                }

                writer.Write(cellName);

                if (nValues > 1) {
                    writer.OpenArray();
                    const auto itRangeEnd = it + nValues;
                    for (; it != itRangeEnd && it != values.cend(); ++it) {
                        writeFunc(**it);
                    }
                    writer.CloseArray();
                } else {
                    writeFunc(cell);
                    ++it;
                }
            }
        }

        void Dump(IPrepIterator& item) {
            IPrepIterator::TDocData data = item.GetData();
            Output.Next();
            NJson::TJsonWriter writer(Output.Output, Output.JsonOpts);
            writer.OpenMap();
            writer.Write("DocId", data.FinalDocId);

            bool done = false;
            if (Level >= EDumpLevel::Row) {
                try {
                    DumpMercuryRow(writer, data.Data);
                    done = true;
                } catch (...) {
                    ERROR_LOG << CurrentExceptionMessage() << Endl;
                }
            }
            if (!done) {
                if (!UseHex) {
                    writer.Write("BLOB", data.Data);
                } else {
                    writer.Write("BLOB", HexEncode(data.Data));
                }
            }
            writer.CloseMap();
            writer.Flush();
        }

        void Finish() {
            Output.Finish();
        }
    };

    void TRtDocDumper::TOutput::Start(IOutputStream& output, bool noPretty) {
        JsonOpts = NJson::TJsonWriterConfig();
        JsonOpts.FormatOutput = !noPretty;
        JsonOpts.ValidateUtf8 = false;
        Output = &output;
        *Output << "[";
        if (JsonOpts.FormatOutput) {
            *Output << Endl;
        }
        HasWritten = false;
    }

    void TRtDocDumper::TOutput::Next() {
        Y_ASSERT(Output);
        if (HasWritten) {
            *Output << ",";
            if (JsonOpts.FormatOutput) {
                *Output << Endl;
            }
        }
        HasWritten = true;
    }

    void TRtDocDumper::TOutput::Finish() {
        if (JsonOpts.FormatOutput) {
            *Output << Endl;
        }
        *Output << "]";
        if (JsonOpts.FormatOutput) {
            *Output << Endl;
        }
        Output->Flush();
        Output = nullptr;
    }
}
