#pragma once

#include "yt_command.h"
#include "config_files.h"

#include <saas/tools/standalone_indexer/lib/standalone_indexer.h>

#include <saas/rtyserver/components/fullarchive/config.h>
#include <saas/rtyserver/config/common_indexers_config.h>
#include <saas/rtyserver/config/config.h>
#include <saas/library/yt/common/yt_blob.h>
#include <saas/library/yt/common/yt_document.h>
#include <saas/protos/shards.pb.h>
#include <saas/tools/standalone_indexer/protos/sample.pb.h>

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

#include <util/generic/size_literals.h>

class TSaasYTIndexingReducer
    : public NYT::IReducer<NYT::TTableReader<TYTDocument>, NYT::TTableWriter<google::protobuf::Message>>
{
public:
    static const ui32 DOC_IN_INDEX = 0;
    static const ui32 BLOB_OUT_INDEX = 0;
    static const ui32 SAMPLE_OUT_INDEX = 1;
    static const ui32 SHARD_RESOURCES_INDEX = 2;
    static const ui32 INCORRECT_DOCS_INDEX = 3;

    TSaasYTIndexingReducer() = default;
    TSaasYTIndexingReducer(const TRtyConfigBundle& configBundle, bool sampleUrls, ui32 sampleUrlsPerShard, bool failOnIndexError,
                            EYTBlobFormat format, ui32 dumperChunkSize, bool runModules, bool filterIncorrectDocs)
        : ConfigBundle(configBundle)
        , SampleUrls(sampleUrls)
        , SampleUrlsPerShard(sampleUrlsPerShard)
        , FailOnIndexError(failOnIndexError)
        , Format(format)
        , DumperChunkSize(dumperChunkSize)
        , RunModules(runModules)
        , FilterIncorrectDocs(filterIncorrectDocs)
    {
    }
    void Do(TReader* input, TWriter* output) override;

    Y_SAVELOAD_JOB(ConfigBundle, SampleUrls, SampleUrlsPerShard, FailOnIndexError, Format, DumperChunkSize, RunModules, FilterIncorrectDocs);

private:
    TRtyConfigBundle ConfigBundle;
    bool SampleUrls = false;
    ui32 SampleUrlsPerShard = 0;
    bool FailOnIndexError = false;
    EYTBlobFormat Format = EYTBlobFormat::Old;
    ui32 DumperChunkSize = 0;
    bool RunModules = false;
    bool FilterIncorrectDocs = false;
};

class TIndexSegmentsCommand : public TYTCommand {
public:
    TIndexSegmentsCommand(TInputs inputs, TOutputs outputs, bool verbose, NSaas::TYTLaunchReport& report, const TRtyConfigBundle& configBundle,
        const TYTConfigFiles& configFiles, ui32 ramGb, float ramReserveFactor, ui32 tmpfsGb, ui32 cpu, bool sampleUrls,
        ui32 sampleUrlsPerShard, bool failOnIndexError, EYTBlobFormat format, bool skyShareEnabled, ui32 dumperChunkSize,
        bool runModules, bool filterIncorrectDocs, const NYT::TNode& acl)
        : TYTCommand(std::move(inputs), std::move(outputs), "TIndexSegmentsCommand", verbose, report, acl)
        , ConfigBundle(configBundle)
        , ConfigFiles(configFiles)
        , RamGb(ramGb)
        , RamReserveFactor(ramReserveFactor)
        , TmpfsGb(tmpfsGb)
        , Cpu(cpu)
        , SampleUrls(sampleUrls)
        , SampleUrlsPerShard(sampleUrlsPerShard)
        , FailOnIndexError(failOnIndexError)
        , BlobFormat(format)
        , SkyShareEnabled(skyShareEnabled)
        , DumperChunkSize(dumperChunkSize)
        , RunModules(runModules)
        , FilterIncorrectDocs(filterIncorrectDocs)
    {
    }
    virtual ~TIndexSegmentsCommand() = default;

protected:
    virtual void DoPrepareSpec() override {
        auto keys = ShardIdKeyColumns;
        keys.Add("segment_id");

        auto config = ConfigBundle.Parse();
        if(config->GetCommonIndexers().HtmlParserConfigFile)
            JobSpec.AddLocalFile(config->GetCommonIndexers().HtmlParserConfigFile);
        if(config->GetCommonIndexers().XmlParserConfigFile)
            JobSpec.AddLocalFile(config->GetCommonIndexers().XmlParserConfigFile);

        JobSpec.CpuLimit(Cpu);
        JobSpec.MemoryLimit(1_GB * RamGb).ExtraTmpfsSize(1_GB * TmpfsGb);
        JobSpec.MemoryReserveFactor(RamReserveFactor);
        ConfigFiles.FillSpecFiles(JobSpec);
        Spec.AddInput<TYTDocument>(Inputs.at(TSaasYTIndexingReducer::DOC_IN_INDEX));
        if (BlobFormat == EYTBlobFormat::Old) {
            Spec.AddOutput<NSaas::TYTBlobBase>(Outputs.at(TSaasYTIndexingReducer::BLOB_OUT_INDEX));
        } else {
            Spec.AddOutput<NSaas::TYTNewBlobBase>(Outputs.at(TSaasYTIndexingReducer::BLOB_OUT_INDEX));
        }
        Spec.AddOutput<NSaas::TUrlSample>(Outputs.at(TSaasYTIndexingReducer::SAMPLE_OUT_INDEX))
            .AddOutput<NRTYServer::TShardResource>(Outputs.at(TSaasYTIndexingReducer::SHARD_RESOURCES_INDEX));
        if (FilterIncorrectDocs) {
            Spec.AddOutput<NSaas::TYTIndexError>(Outputs.at(TSaasYTIndexingReducer::INCORRECT_DOCS_INDEX));
        }
        Spec.ReducerSpec(JobSpec)
            .ReduceBy(keys);
        Opts.MountSandboxInTmpfs(true);
    }

    NYT::TTableSchema CreateBlobSchema() const {
        NYT::TSortColumns keyColumns = ShardIdKeyColumns;
        keyColumns.Add("segment_id");

        NYT::TTableSchema schema;
        if (BlobFormat == EYTBlobFormat::New) {
            keyColumns.Add("filename");
            keyColumns.Add("part_index");
            schema = NYT::CreateTableSchema<NSaas::TYTNewBlobBase>(keyColumns);
            schema.AddColumn(NYT::TColumnSchema().Name("sha1").Type(NYT::VT_STRING));
            schema.AddColumn(NYT::TColumnSchema().Name("md5").Type(NYT::VT_STRING));
            for (auto &column : schema.MutableColumns()) {
                column.Group(column.Name() == "data" ? "data" : "meta");
            }
        } else {
            keyColumns.Add("name");
            keyColumns.Add("part_index");
            schema = NYT::CreateTableSchema<NSaas::TYTBlobBase>(keyColumns);
        }
        return schema;
    }

    void CreateIndexBlobTable(NYT::IClientBase* client) {
        NYT::TYPath deltaAttribute = Inputs.at(TSaasYTIndexingReducer::DOC_IN_INDEX).Path_ + "/@is_delta";
        bool isDelta = (client->Exists(deltaAttribute) && client->Get(deltaAttribute).AsBool());
        NYT::TNode options;
        if (SkyShareEnabled) {
            options = NYT::TNode()("enable_skynet_sharing", true)
                                ("optimize_for", "scan");
        }
        options("schema", CreateBlobSchema().ToNode())("is_delta", isDelta)("blob_format", ToString(BlobFormat));

        client->Create(Outputs.at(TSaasYTIndexingReducer::BLOB_OUT_INDEX).Path_, NYT::ENodeType::NT_TABLE,
                        NYT::TCreateOptions().Recursive(true).Force(true).Attributes(options));
    }

    void CreateFilteredDocsTable(NYT::IClientBase* client) {
        client->CreateTableWriter<NSaas::TYTIndexError>(Outputs.at(TSaasYTIndexingReducer::INCORRECT_DOCS_INDEX),
                                                            NYT::TTableWriterOptions().InferSchema(true))->Finish();
    }
    virtual void DoRun(NYT::IClientBase* client) override {
        CreateIndexBlobTable(client);
        if (FilterIncorrectDocs) {
            CreateFilteredDocsTable(client);
        }
        Reduce(client, Spec, new TSaasYTIndexingReducer(ConfigBundle, SampleUrls, SampleUrlsPerShard, FailOnIndexError,
                       BlobFormat, DumperChunkSize, RunModules, FilterIncorrectDocs), Opts);
        if (FilterIncorrectDocs) {
            i64 droppedDocsCount = client->Get(Outputs.at(TSaasYTIndexingReducer::INCORRECT_DOCS_INDEX).Path_ + "/@row_count").AsInt64();
            Report.AddSensorValue("dropped-docs-count", droppedDocsCount);
        }
        {
            NYT::TTableReaderPtr<NRTYServer::TShardResource> reader = client->CreateTableReader<NRTYServer::TShardResource>(Outputs.at(TSaasYTIndexingReducer::SHARD_RESOURCES_INDEX));
            ui64 maxAgeTs = 0;
            for (; reader->IsValid(); reader->Next()) {
                const auto& row = reader->GetRow();
                if (row.HasStat()) {
                    TStringInput si(row.GetStat());
                    NJson::TJsonValue stat;
                    try {
                        NJson::ReadJsonTree(&si, &stat);
                        maxAgeTs = Max<ui64>(maxAgeTs, FromString<ui64>(stat["ts_max"].GetStringRobust()));
                    } catch (const yexception& e) {
                        ERROR_LOG << e.what() << Endl;
                    }
                }
            }
            if (maxAgeTs > 0)
                Report.AddSensorValue("index-age", Now().Seconds() - maxAgeTs);
        }
    }

private:
    const TRtyConfigBundle& ConfigBundle;
    const TYTConfigFiles& ConfigFiles;
    const ui32 RamGb;
    const float RamReserveFactor;
    const ui32 TmpfsGb;
    const ui32 Cpu;
    NYT::TReduceOperationSpec Spec;
    NYT::TUserJobSpec JobSpec;
    NYT::TOperationOptions Opts;

    const bool SampleUrls;
    const ui32 SampleUrlsPerShard;
    const bool FailOnIndexError;
    const EYTBlobFormat BlobFormat;
    const bool SkyShareEnabled;
    const ui32 DumperChunkSize;
    const bool RunModules;
    const bool FilterIncorrectDocs;
};
