#include "library/cpp/json/writer/json_value.h"
#include "proto_dump_impl.h"

#include <kernel/dssm_applier/embeddings_transfer/embeddings_transfer.h>
#include <kernel/dssm_applier/optimized_model/optimized_model_applier.h>

#include <kernel/doom/kosher_offroad_portions/kosher_portion_io_factory.h>
#include <kernel/doom/kosher_offroad_portions/portion_counts_io.h>
#include <kernel/doom/kosher_offroad_portions/portion_attributes_io.h>
#include <kernel/struct_codegen/print/struct_print.h>
#include <kernel/doom/offroad_erf_wad/erf_io.h>

#include <library/cpp/json/json_reader.h>

#include <robot/jupiter/tools/shards_prepare/lib/panther_utils/index_part.h>
#include <robot/jupiter/library/data_printer/json_printer.h>
#include <robot/jupiter/library/tables/shards_prepare.h>
#include <robot/jupiter/protos/content_attrs.pb.h>

#include <util/stream/format.h>
#include <util/string/cast.h>

namespace NRtDoc {
    // This is adapted from robot/jupiter/tools/difftool/lib/diff_printer.cpp

    template<class TProto, class TIo>
    class TOffroadJsonPrinter : public IProtoJsonTransformer {
    public:
        bool Transform(NJson::TJsonValue& /*json*/, const NProtoBuf::Message& /*msg*/) const override {
            return false;
        };
    };

    template<class TIo>
    class TOffroadJsonPrinter<NJupiter::TOffroadKeyInvEntry, TIo> final: public IProtoJsonTransformer {
    public:
        using TProto = NJupiter::TOffroadKeyInvEntry;
        using THit = typename TIo::THit;
        using TReader = typename TIo::TReader;
        using TOffroadIoFactory = NDoom::TKosherPortionIoFactory<TIo>;
    public:
        bool Transform(NJson::TJsonValue& json, const NProtoBuf::Message& msg) const override {
            const TProto& entry = static_cast<const TProto&>(msg);
            const auto& keysBlob = entry.GetKeysBlob();
            const auto& hitsBlob = entry.GetHitsBlob();
            TReader reader(TOffroadIoFactory::HitReaderTable(), TArrayRef<const char>(hitsBlob.data(), hitsBlob.size()),
                    TOffroadIoFactory::KeyReaderTable(), TArrayRef<const char>(keysBlob.data(), keysBlob.size()));

            using TKeyRef = typename TReader::TKeyRef;
            using THit = typename TReader::THit;

            NJson::TJsonValue dump(NJson::EJsonValueType::JSON_MAP);
            TKeyRef key;
            while (reader.ReadKey(&key)) {
                TStringStream ss;
                THit hit;
                while(reader.ReadHit(&hit)) {
                    ss << hit;
                }
                ss.Finish();
                dump.InsertValue(ToString(key), ss.Str());
            }
            json.EraseValue("HitsBlob");
            json.EraseValue("KeysBlob");
            json.InsertValue("Data", dump);

            NJson::TJsonValue parsedPart;
            {
                ui32 index = 0;
                ui32 part = 0;
                NJupiter::NPrivate::TIndexIdPart::Unpack(entry.GetPart(), &index, &part);
                //TStringStream hexPart;
                //hexPart << Hex(entry.GetPart());
                //parsedPart.InsertValue("_Value", hexPart.Str());
                parsedPart.InsertValue("Index", index);
                parsedPart.InsertValue("Part", part);
            }

            json.EraseValue("Part");
            json.InsertValue("Part", parsedPart);

            return true;
        };
    };

    struct TErfDumper : IProtoJsonTransformer {
        using TIo = NDoom::TErf2Io;
        using THit = typename TIo::TReader::THit;
        using THitData = std::remove_pointer<THit>::type;
        using TReaderTable = TIo::TReader::TTable;
        using TSerializer = NDoom::TErfSerializer<THit>;
        using TStructReader = NDoom::TStructReader<THit, TSerializer, TIo::StructType, TIo::CompressionType>;
        using TModel = TIo::TModel;

        TErfDumper() {
            ReaderTable.Reset(NDoom::TStandardIoModelsStorage::Model<TModel>(TIo::DefaultModel));
        }

        bool Transform(NJson::TJsonValue& json, const NProtoBuf::Message& msg) const final {
            auto& wadLumps = static_cast<const NPlutonium::TDocWadLumps&>(msg);
            TString cuteString;
            TStringStream unknownLumps;
            for (auto& lump : wadLumps.GetLumpsList().GetLumps()) {
                const auto id = FromString<NDoom::TWadLumpId>(lump.GetId());
                if (id == NDoom::TWadLumpId(NDoom::EWadIndexType::ErfIndexType, NDoom::EWadLumpRole::Struct)) {
                    const TArrayRef<const char> lumpData(lump.GetData().data(), lump.GetData().size());
                    THit hit;
                    TStructReader reader;
                    // passing of sizeof(THitData) is not always correct
                    // the better way would be to read doc_lumps with global_lumps together
                    reader.Reset(&ReaderTable, sizeof(THitData), lumpData);
                    bool parsed = reader.Read(&hit);
                    Y_ENSURE(parsed, "cannot parse erf");
                    cuteString = TCuteStrStructPrinter<THitData>::PrintFields(*hit, typename THitData::TFieldMask(true));
                } else {
                    unknownLumps << id << ";";
                }
            }
            unknownLumps.Finish();
            if (!unknownLumps.Str().empty()) {
                json.InsertValue("UnknownLumps", unknownLumps.Str());
            }

            json.EraseValue("LumpsList");
            json.InsertValue("LumpsList", cuteString);
            return true;
        }

    private:
        TReaderTable ReaderTable;
    };

    struct TContentAttrsTransformer : IProtoJsonTransformer {
        bool Transform(NJson::TJsonValue& json, const NProtoBuf::Message& msg) const final {
            NOptimizedModel::NProto::TSerializedApplyExResult proto;

            auto& contentAttrs = static_cast<const NJupiter::TContentAttrs&>(msg);
            if (!contentAttrs.HasSerializedProtoOfOptimizedRthubModelApplyResult() ||
                !proto.ParseFromString(contentAttrs.GetSerializedProtoOfOptimizedRthubModelApplyResult())) {
                return false;
            }
            auto applyResult = NOptimizedModel::TApplyExResult::DeserializeFromProto(proto);
            const TStringBuf dssmBertDistillL2 = "doc_embedding_dssm_bert_distill_l2";
            if (const auto* embedding = applyResult.GetEmbedding(dssmBertDistillL2)) {
                NJson::TJsonArray embedData;
                for (auto n : embedding->Data) {
                    embedData.AppendValue(NJson::TJsonValue(n));
                }
                NJson::TJsonMap embedJson;
                embedJson.InsertValue("Version", embedding->Version);
                embedJson.InsertValue("Data", std::move(embedData));
                json.InsertValue("SerializedProtoOfOptimizedRthubModelApplyResult.doc_embedding_dssm_bert_distill_l2", std::move(embedJson));
            }
            return true;
        }
    };

    void TProtoDumper::RegisterOffroad(TJsonStringPrintersRegistry& registry) {
        auto tables = NJupiter::TShardsPrepareTables::Clientless();
        auto countsPrinter = MakeIntrusive<TOffroadJsonPrinter<NJupiter::TOffroadKeyInvEntry, NDoom::TPortionCountsOffroadIo>>();
        registry.RegisterFormat(tables.GetCountsKeyInvTable(0), countsPrinter);
        registry.RegisterFormat(tables.GetKeyInvCountsTable(0), countsPrinter);
        registry.RegisterFormat(tables.GetAnnCountsTable(0), countsPrinter);
        registry.RegisterFormat(tables.GetLinkAnnCountsTable(0), countsPrinter);
        registry.RegisterFormat(tables.GetFactorAnnCountsTable(0), countsPrinter);

        auto attrsPrinter = MakeIntrusive<TOffroadJsonPrinter<NJupiter::TOffroadKeyInvEntry, NDoom::TPortionAttributesOffroadIo>>();
        registry.RegisterFormat(tables.GetAttrsKeyInvTable(0), attrsPrinter);

        registry.RegisterFormat(tables.GetErfDocLumpsTable(0), MakeIntrusive<TErfDumper>());

        registry.RegisterFormat(tables.GetContentAttrsTable(0), MakeIntrusive<TContentAttrsTransformer>());
    }
}
