#include "shards_prepare.h"

#include <robot/jupiter/library/rtdoc/rt/reader.h>
#include <robot/jupiter/library/rtdoc/rt/writer.h>
#include <robot/jupiter/library/tables/shards_prepare.h>
#include <robot/jupiter/protos/shards.pb.h>

#include <mapreduce/yt/interface/io.h>

#include <google/protobuf/message.h>

namespace NFusion {

    class TShardsPrepareLumpsFiller::TImpl {
    public:
        TImpl(const TVector<NJupiter::TMercuryCmdPtr>& commands, bool heavyDataLumpsEnabled) {
            Init(Input);
            AddProcessors(commands);
            HeavyDataLumpsEnabled = heavyDataLumpsEnabled;
        }

    private:
        using TReaderImpl = NJupiter::TMercuryYtReader;
        using TReader = NYT::TMessageReader;
        using TWriter = NYT::TMessageWriter;

        struct TBoundProcessor {
            TSimpleSharedPtr<NMercury::IProcAdapter> Processor;
            THolder<TReader> Reader;
            THolder<TWriter> Writer;
            TIntrusivePtr<TReaderImpl> ReaderImpl;
        };

        using TCollection = TVector<TBoundProcessor>;

        TCollection Processors;
        NJupiter::TMercuryDataSet Input;
        NJupiter::TMercuryWriterOutput Output;
        bool HeavyDataLumpsEnabled = false;

    private:
        void AddProcessor(const NJupiter::IMercuryCmd& cmd) {
            TIntrusivePtr<NJupiter::TMercuryYtReader> protoreader = MakeIntrusive<NJupiter::TMercuryYtReader>(Input, cmd.GetInputs());
            THolder<TReader> reader = MakeHolder<TReader>(protoreader);
            THolder<TWriter> writer(NJupiter::TMercuryWriterImpl::Create(&Output, cmd.GetOutputs(), cmd.GetDestinations()));
            Processors.push_back(TBoundProcessor{cmd.GetAdapter(), std::move(reader), std::move(writer), std::move(protoreader)});
            Y_ENSURE(Processors.back().Processor);
        }

        void AddProcessors(const TVector<NJupiter::TMercuryCmdPtr>& commands) {
            for(const auto& command : commands) {
                AddProcessor(*command);
            }
        }

        static void Init(NJupiter::TMercuryDataSet& ds) {
            using namespace NJupiter;

            static const TShardsPrepareTables tables = TShardsPrepareTables::Clientless();
            static const TVector<std::pair<const TString, const google::protobuf::Descriptor*>> fields = {
                {tables.GetCalculatedAttrsTable(0), TCalculatedFactors::descriptor()},
                {tables.GetHostDomainDocMappingsTable(0), TGroupAttrHostDomainDocMappings::descriptor()},
                {tables.GetHostMappingTable(0), TGroupAttrStrValueMapping::descriptor()},
                {tables.GetExternalRelevAttrsTable(0), TExternalRelevAttrs::descriptor()},
                {tables.GetContentAttrsTable(0), TContentAttrs::descriptor()},
                {tables.GetUrlHeavyDataTable(0), TUrlHeavyData::descriptor()},
                {tables.GetWalrusTable(0), TWalrus::descriptor()},
                {tables.GetHostAttrsTable(0), THostAttrs::descriptor()},
                {tables.GetMobileUrlTable(0), TMobileUrl::descriptor()},
                {tables.GetStatsTableLink(), TUrlToShardStats::descriptor()},
            };
            for (const auto& [name, descriptor] : fields) {
                ds.AddField(name, descriptor);
            }
            ds.Initialize();
        }

    public:
        static auto CreateRowsDefaultInstances() {
            using namespace NJupiter;
            TVector<THolder<google::protobuf::Message>> ret;
            ret.emplace_back(new TCalculatedFactors());
            ret.emplace_back(new TGroupAttrHostDomainDocMappings());
            ret.emplace_back(new TGroupAttrStrValueMapping());
            ret.emplace_back(new TExternalRelevAttrs());
            ret.emplace_back(new TContentAttrs());
            ret.emplace_back(new TUrlHeavyData());
            ret.emplace_back(new TWalrus());
            ret.emplace_back(new THostAttrs());
            ret.emplace_back(new TMobileUrl());
            ret.emplace_back(new TUrlToShardStats());
            return ret;
        }

        void Fill(const NJupiter::TMercuryLumps& inputData, NJupiter::TMercuryLumps& outputData) {
            using namespace NJupiter;

            static const auto defaultInstances = CreateRowsDefaultInstances();
            std::vector<google::protobuf::Message*> rows(defaultInstances.size(), nullptr);
            TVector<THolder<google::protobuf::Message>> rowsHolders;

            for (const auto& lump : inputData.GetLumps()) {
                if (const auto* field = Input.GetFieldExtPtr(lump.GetName())) {
                    rowsHolders.push_back(THolder(defaultInstances[field->Idx]->New()));
                    Y_PROTOBUF_SUPPRESS_NODISCARD rowsHolders.back()->ParseFromString(lump.GetData());
                    rows[field->Idx] = rowsHolders.back().Get();
                }
            }

            if (auto* ptr = dynamic_cast<TCalculatedFactors*>(rows[0])) {
                ptr->SetDocId(0);
                ptr->SetChunk(0);
                ptr->SetLocalDocId(0);
            }

#define SET_DOC_ID(idx, row_type) \
            if (auto* ptr = dynamic_cast<row_type*>(rows[idx])) { \
                ptr->SetDocId(0); \
            }

            SET_DOC_ID(3, TExternalRelevAttrs);
            SET_DOC_ID(4, TContentAttrs);
            if (HeavyDataLumpsEnabled) {
                SET_DOC_ID(5, TUrlHeavyData);
                SET_DOC_ID(6, TWalrus);
                SET_DOC_ID(7, THostAttrs);
                SET_DOC_ID(8, TMobileUrl);
            } else {
                SET_DOC_ID(5, TWalrus);
                SET_DOC_ID(6, THostAttrs);
                SET_DOC_ID(7, TMobileUrl);
            }

            Input.Set(rows);

            Output.Reset(&outputData);

            for (auto& i : Processors) {
                i.ReaderImpl->NextKey();
                if (!i.ReaderImpl->IsValid()) {
                    continue; // no rows for this processor, go to the next one
                }

                NMercury::IProcAdapter& proc = *i.Processor;
                proc.Start(i.Writer.Get());
                proc.Do(i.Reader.Get(), i.Writer.Get());
                proc.Finish(i.Writer.Get());

                i.Reader->Next(); // to reset reader cache
            }
        }
    };

    TShardsPrepareLumpsFiller::TShardsPrepareLumpsFiller(const TVector<NJupiter::TMercuryCmdPtr>& commands, bool heavyDataLumpsEnabled)
        : Impl(new TImpl(commands, heavyDataLumpsEnabled))
    {
    }

    TShardsPrepareLumpsFiller::~TShardsPrepareLumpsFiller() = default;

    void TShardsPrepareLumpsFiller::Fill(const NJupiter::TMercuryLumps& inputData, NJupiter::TMercuryLumps& outputData) {
        Impl->Fill(inputData, outputData);
    }

}
