#pragma once

#include <saas/library/rtyt/lib/client.h>
#include <saas/library/rtyt/lib/operation/factory.h>
#include <saas/rtyserver_jupi/library/tables/tables.h>

namespace {
    const static ui32 DeletedDocument = -1;
    static const TString FakeMercuryShardStr = "-";
}

class TRTYTDocIdRemapperBase
    : public NYT::IMapper<NYT::TTableReader<NYT::Message>, NYT::TTableWriter<NYT::Message>>
{
protected:
    TString TableName; // key for proto descriptors
    TVector<ui32> Remap;

public:
    Y_SAVELOAD_JOB(TableName, Remap);

    void Reinitialize(const TString& tableName, const TVector<ui32>& map) {
        TableName = tableName;
        Remap = map;
    }

    TRTYTDocIdRemapperBase() = default;
    TRTYTDocIdRemapperBase(const TString& tableName, const TVector<ui32>& map) {
        Reinitialize(tableName, map);
    }
};

template <typename TProto>
class TRTYTDocIdRemapper
    : public TRTYTDocIdRemapperBase
{
public:
    using TRTYTDocIdRemapperBase::TRTYTDocIdRemapperBase;

    void Do(TReader* input, TWriter* output) override;
};

template <typename TProto>
void TRTYTDocIdRemapper<TProto>::Do(TReader* in, TWriter* out) {
    auto descriptor = NFusion::GetRTYTPrepTableTypes()[TableName];
    auto docIdFieldDescr = descriptor->FindFieldByName("DocId");
    Y_ENSURE(docIdFieldDescr, "Currently can set docids only in DocId field, but table " << TableName << " doesn't conain it");
    auto localDocIdFieldDescr = descriptor->FindFieldByName("LocalDocId");
    auto shardFieldDescr = descriptor->FindFieldByName("Shard");
    auto chunkFieldDescr = descriptor->FindFieldByName("Chunk");
    if (!localDocIdFieldDescr) {
        NOTICE_LOG << "There is no LocalDocId in table " << TableName << ", won't fill it" << Endl;
    }
    if (!shardFieldDescr) {
        NOTICE_LOG << "There is no Shard in table " << TableName << ", won't fill it" << Endl;
    }
    if (!chunkFieldDescr) {
        NOTICE_LOG << "There is no Chunk in table " << TableName << ", won't fill it" << Endl;
    }
    TProto inRow;
    auto reflection = inRow.GetReflection();
    for (; in->IsValid(); in->Next()) {
        in->MoveRow<TProto>(&inRow);
        ui32 docid = inRow.GetDocId();
        if (docid != DeletedDocument) {
            docid = Remap[docid];
        }
        inRow.SetDocId(docid);
        if (localDocIdFieldDescr) {
            reflection->SetUInt32(&inRow, localDocIdFieldDescr, docid);
        }
        if (shardFieldDescr) {
            reflection->SetString(&inRow, shardFieldDescr, FakeMercuryShardStr);
        }
        if (chunkFieldDescr) {
            reflection->SetUInt32(&inRow, chunkFieldDescr, 1);
        }
        if (docid != DeletedDocument) {
            out->AddRow<TProto>(inRow);
        }
    }
}
