#include "shards_processor.h"
#include "synchronizer.h"
#include "config.h"
#include <saas/rtyserver/config/realm_config.h>
#include <saas/rtyserver/indexer_core/merger_interfaces.h>
#include <saas/rtyserver/merger/doc_extractor.h>

namespace NRTYServer {
    TShardsProcessor::TShardsProcessor(TSynchronizer* synchronizer)
        : ISynchronizerProcessor(synchronizer)
    {}

    static void ProcessTask(TAsyncTaskExecutor::TTask& task, const ui32 shardMin, const ui32 shardMax, const NSaas::TShardsDispatcher::TContext& sharding, TShardsProcessor::TAction action, const TString& path) {
        const IIndexStorage& is = TMergerEngine::GetMerger()->GetIndexStorage();
        TVector<TString> dirs = is.GetFinalIndexes(path);
        ui32 docsCount = 0;
        TDocumentsExtractor::TContext ctx(sharding);
        TDocumentsExtractor docsExtractor(ctx);
        for (ui32 i = 0; i < dirs.size(); ++i) {
            TVector<ui32> docIdsShards;
            TDocPlaceMap shards;
            TIndexControllerPtr index = is.GetIndexController(dirs[i]);
            docsExtractor.GetDocPlaceMap(*index, shards);
            for (ui32 doc = 0; doc < shards.size(); ++doc) {
                if (shards[doc].Shard >= shardMin && shards[doc].Shard <= shardMax)
                    docIdsShards.push_back(doc);
            }
            docsCount += docIdsShards.size();
            switch (action) {
                case TShardsProcessor::spaCheck: {
                ui32 docsInIndex = index->GetDocumentsCount(false);
                if (docIdsShards.size() != docsInIndex) {
                    ERROR_LOG << "Shards check fail: docIdsShards.size() != docsInIndex " << docIdsShards.size() << "!=" << docsInIndex << Endl;
                    task.GetReply()->InsertValue("shards_result", "FAIL");
                    task.GetReply()->InsertValue("failed_dir", dirs[i]);
                    return;
                }
                break;
            }
            case TShardsProcessor::spaRemove:
                index->RemoveDocIds(docIdsShards);
                break;
            default:
                FAIL_LOG("Incorrect action %i", (int)action);
            }
        }
        task.GetReply()->InsertValue("shards_result", "OK");
        task.GetReply()->InsertValue("doc_ids_shards", docsCount);
    }

    bool TShardsProcessor::ProcessRequest(const TCgiParameters& cgi, TStringBuf postBuffer, TAsyncTaskExecutor::TTask& task) {
        Y_UNUSED(postBuffer);
        DEBUG_LOG << "Start shards request processing " << cgi.Print() << Endl;
        TString minShardS = cgi.Get("min_shard");
        TString maxShardS = cgi.Get("max_shard");
        NSaas::TShardsDispatcher::TContext sharding = NSaas::TShardsDispatcher::TContext::FromString(cgi.Get("sharding_type"));
        TString actionS = cgi.Get("act_shards");
        ui32 minShard, maxShard;

        try {
            minShard = FromString<ui32>(minShardS);
            maxShard = FromString<ui32>(maxShardS);
        } catch (...) {
            ythrow yexception() << "incorrect &min_shard or &max_shard";
        }

        TAction action = spaCheck;

        if (actionS == "check")
            action = spaCheck;
        else if(actionS == "remove")
            action = spaRemove;
        else
            ythrow yexception() << "incorrect &act_shards";

        if (sharding.Type == NSaas::KeyPrefix && !Synchronizer->GetRTYConfig().IsPrefixedIndex)
                ythrow yexception() << "incorrect &sharding_type for non prefixed index";

        ProcessTask(task, minShard, maxShard, sharding, action, Synchronizer->GetRTYConfig().GetRealmListConfig().GetMainRealmConfig().RealmDirectory);
        DEBUG_LOG << "Finish shards request processing " << cgi.Print() << Endl;
        return true;
    }
}
