#include "yt_document_mapper.h"

#include <saas/library/persqueue/common/common.h>
#include <saas/library/searchmap/parsers/parser.h>
#include <saas/library/searchmap/searchmap.h>

#include <saas/protos/positions.pb.h>
#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/searcher_config.h>
#include <saas/rtyserver/factors/factors_config.h>
#include <saas/rtyserver/pruning_config/pruning_config.h>
#include <saas/rtyserver/docfetcher/library/config.h>
#include <saas/rtyserver/docfetcher/library/stream_basics.h>

namespace {
    bool AssociatedWithStream(const NFusion::TPersQueueStreamConfig& /*stream*/, NRTYServer::TMessage::TDocument& doc) {
        if (doc.HasPosition()) {
            const auto& positionKey = doc.GetPosition().GetKey();
            auto topic = NSaasLB::GetTopicInfo(positionKey);
            return bool(topic);
        }
        return false;
    }
}

REGISTER_MAPPER(TSaasYTDocumentMapper);

class TShardSelector: public NSearchMapParser::ISearchMapProcessor {
public:
    TShardSelector(TShardIndex shardIndex)
        : ShardIndex(shardIndex)
    {
    }

    void Do(const NSearchMapParser::TServiceSpecificOptions& options, const NSearchMapParser::TReplicaSpecificOptions& /*rso*/, const TInterval<TShardIndex>& /*interval*/, const NSearchMapParser::TSearchInformation& si) override {
        if (options.ShardsDispatcher->CheckInterval(ShardIndex, si.Shards)) {
            // the same slot can be processed several times if the search map contains intersecting intervals
            // so we need to save intervals into the hash set (ordered set cannot be used because TShardId doesn't have a proper comparison operator)
            ShardIds.emplace(si.Shards);
        }
    }

    const THashSet<TShardId>& GetShardIds() const {
        return ShardIds;
    }

private:
    const TShardIndex ShardIndex;
    THashSet<TShardId> ShardIds;
};

TSaasYTDocumentMapper::TSaasYTDocumentMapper()
    : Service(nullptr)
{
}

TSaasYTDocumentMapper::TSaasYTDocumentMapper(TContext context, ui64 totalDocumentsCount, const TRtyConfigBundle& configBundle, bool addExtraTimestamp)
    : Service(nullptr)
    , Context(context)
    , TotalDocumentsCount(totalDocumentsCount)
    , ConfigBundle(configBundle)
    , AddExtraTimestamp(addExtraTimestamp)
{
}

TSaasYTDocumentMapper::~TSaasYTDocumentMapper() = default;

void TSaasYTDocumentMapper::Start(TWriter* /*output*/) {
    THolder<NSearchMapParser::ISearchMapParser> searchMapParser = NSearchMapParser::OpenSearchMap(Context.SearchMapText, {Context.ServiceName});
    SearchMap.Reset(new NSearchMapParser::TSearchMap(searchMapParser->GetSearchMap()));
    auto serviceIt = SearchMap->GetServiceMap().find(Context.ServiceName);
    Y_VERIFY(serviceIt != SearchMap->GetServiceMap().end());
    Service = &serviceIt->second;

    auto config = ConfigBundle.Parse();
    auto dfConfigPtr = config->GetModuleConfig<NFusion::TDocFetcherConfig>(NFusion::DocfetcherModuleName);
    DocFetcherConfig = MakeHolder<NFusion::TDocFetcherConfig>(*dfConfigPtr);

    if (config->Pruning->GetType() == TPruningConfig::GROUP_ATTR) {
        GrAttrToFillOnDelete = config->Pruning->ToString();
    }

    if (!!config->GetSearcherConfig().Factors && !!config->GetSearcherConfig().Factors->StaticFactors()) {
        for (const auto& factor: config->GetSearcherConfig().Factors->StaticFactors()) {
            if (!factor.DefaultValue.IsSet()) {
                FactorsToFillOnDelete.insert(factor.Name);
            }
        }
    }

    if (AddExtraTimestamp) {
        auto mapReduceStreams = DocFetcherConfig->MapReduceStreams;
        TVector<NFusion::TMapReduceStreamConfig> enabledMRStreams;
        for (const auto& stream : mapReduceStreams) {
            if (stream.Enabled) {
                enabledMRStreams.push_back(stream);
            }
        }

        auto snapshotStream = DocFetcherConfig->SnapshotStream;
        bool IsDeltaStream = !!snapshotStream && snapshotStream.Get()->UseGroups;
        Y_ENSURE(IsDeltaStream || (enabledMRStreams.size() >= 1), "option --extra-timestamp can't be used when no stream uses extra timestamps");
        Y_ENSURE(!IsDeltaStream || (enabledMRStreams.size() < 1), "Use of option --extra-timestamp is unsupported with more than one stream, which use extra timestamps");
        if (IsDeltaStream) {
            DefaultStreamId = snapshotStream.Get()->StreamId;
        } else {
            DefaultStreamId = enabledMRStreams.front().StreamId;
        }

    }

    RowProcessor.Reset(NSaas::IRowProcessor::TFactory::Construct(Context.Processor, Context.ProcessorOptions));
}

void TSaasYTDocumentMapper::FillStreamSpecificFileds(NRTYServer::TMessage::TDocument& doc, const TReader* input) const {
    NFusion::TBaseStreamConfig* associatedStream = nullptr;
    if (DocFetcherConfig) {
        int enabledPQStreams = 0;
        for (const auto& stream : DocFetcherConfig->PersQueueStreams) {
            if (stream.Enabled) {
                enabledPQStreams++;
            }
        }
        Y_VERIFY(enabledPQStreams <= 1, "cannot associate document with 2 logbroker streams SAAS-5496");
        for (auto& stream : DocFetcherConfig->PersQueueStreams) {
            if (stream.Enabled && AssociatedWithStream(stream, doc)) {
                Y_VERIFY(!associatedStream, "two streams associated for one document");
                associatedStream = &stream;
            }
        }
    }

    ui32 streamId = associatedStream ? associatedStream->StreamId : DefaultStreamId;

    if (associatedStream && !doc.HasStreamId()) {
        doc.SetStreamId(NFusion::GetSubStream(streamId, NFusion::SubStreamDocumentAux));
    }

    if (AddExtraTimestamp) {
        if (!doc.HasStreamId())
            doc.SetStreamId(NFusion::GetSubStream(streamId, NFusion::SubStreamDocumentAux));
        auto timestamp = doc.AddTimestamps();
        timestamp->SetStream(NFusion::GetSubStream(streamId, NFusion::SubStreamPositionAux));
        timestamp->SetValueEx(Context.ProcessorOptions.Timestamp);
        if (Context.IsDelta) {
            timestamp->SetValue(input->GetRowIndex());
        } else {
            timestamp->SetValue(Max<ui64>());
        }
   }
}

void TSaasYTDocumentMapper::Do(TReader* input, TWriter* output) {
    Y_VERIFY(Service);
    Y_VERIFY(!!RowProcessor);
    for (; input->IsValid(); input->Next()) {
        const NYT::TNode& row = input->GetRow();
        if (RowProcessor->ShouldSkipRow(row)) {
            continue;
        }

        NSaas::TAction action = RowProcessor->ProcessRowSingle(row);

        TYTDocument outRow;
        NRTYServer::TMessage::TDocument& doc = *outRow.MutableDocument();
        doc.CopyFrom(action.ToProtobuf().GetDocument());

        FillStreamSpecificFileds(doc, input);

        if (action.GetActionType() == NSaas::TAction::atDelete && !Context.SaveDeletedDocuments) {
            continue;
        }
        if (action.GetActionType() == NSaas::TAction::atDelete && GrAttrToFillOnDelete) {
            NRTYServer::TAttribute* attribute = doc.AddGroupAttributes();
            attribute->set_name(GrAttrToFillOnDelete);
            attribute->set_value("0");
            attribute->set_type(::NRTYServer::TAttribute::INTEGER_ATTRIBUTE);
        }
        if (action.GetActionType() == NSaas::TAction::atDelete && FactorsToFillOnDelete) {
            NRTYServer::TMessage::TErfInfo* factors = doc.MutableFactors();
            for (const auto& it : FactorsToFillOnDelete) {
                factors->AddNames(it);
                factors->MutableValues()->AddValues()->SetValue(0);
            }
        }
        outRow.SetIsDeleted(action.GetActionType() == NSaas::TAction::atDelete);
        const TShardIndex idx = Service->ShardsDispatcher->GetShard(action.ToProtobuf());
        TShardSelector s(idx);
        Service->ProcessSearchMap(s);
        Y_ENSURE(!s.GetShardIds().empty(), "ShardId not found for ShardIndex" << idx);
        for (const auto& shardId : s.GetShardIds()) {
            Y_VERIFY(idx >= shardId.GetMin() && idx <= shardId.GetMax());

            /* Example:
             * MaxDocumentsPerSegment: 30'000
             * TotalDocumentsCount: 2'503'656
             * shardId: 8191-16381
             * shardLength: 8191
             * documentsInThisShard: 312928
             * segmentsInThisShard: 11
             * (idx =  8191) => (segmentId =  0)
             * (idx =  8200) => (segmentId =  0)
             * (idx =  9000) => (segmentId =  1)
             * (idx = 16381) => (segmentId =  10)
            */
            const ui64 shardLength = shardId.GetLength() + 1; // borders are inclusive, length of (1,10) is 10, not 9
            const ui64 fullLength = NSearchMapParser::FullInterval.GetLength() + 1;
            const ui64 documentsInThisShard = (TotalDocumentsCount * shardLength) / fullLength;
            const ui32 segmentsInThisShard = Max<ui32>(
                Context.MinSegmentsCount,
                (documentsInThisShard / Context.MaxDocumentsPerSegment) + (documentsInThisShard % Context.MaxDocumentsPerSegment != 0)
            );
            const ui32 segmentId = (idx - shardId.GetMin()) * segmentsInThisShard / shardLength;

            outRow.SetShardIndex(idx);
            outRow.SetShardId(shardId);
            outRow.SetSegmentId(segmentId);

            output->AddRow(outRow);
        };
    }
}
