#include "erf_component.h"
#include "erf_config.h"
#include "erf_disk.h"
#include "erf_builder.h"
#include "erf_parsed_entity.h"
#include "erf_mem.h"

#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/config/merger_config.h>
#include <saas/rtyserver/config/searcher_config.h>
#include <saas/rtyserver/factors/factors_config.h>
#include <saas/rtyserver/factors/rank_model.h>
#include <saas/rtyserver/merger/library/plain_file_block_merger.h>

#include <kernel/web_factors_info/factor_names.h>

namespace {
    bool IsUsed(const TRTYServerConfig& Config) {
        return !!Config.GetSearcherConfig().Factors && Config.GetSearcherConfig().Factors->StaticFactors().size()
            && Config.ComponentsConfig.Get<TRTYErfIndexComponentConfig>(ERF_COMPONENT_NAME)->IsUsed;
    }
}

TRTYErfIndexComponent::TRTYErfIndexComponent(const TRTYServerConfig& config, bool isUsed, const TString& fileName, const IRTYStaticFactors& staticFactors)
    : IIndexComponent(isUsed)
    , Config(config)
    , StaticFactors(staticFactors)
    , BaseFileName(fileName)
{
    const bool defaultChecked = true;
    IndexFiles.insert(TIndexFile(BaseFileName, defaultChecked, TIndexFile::ppLock));
    IndexFiles.insert(TIndexFile(BaseFileName + ".hdr", defaultChecked, TIndexFile::ppDisable));
}

TRTYErfIndexComponent::TRTYErfIndexComponent(const TRTYServerConfig& config, const TString& fileName, const IRTYStaticFactors& staticFactors)
    : TRTYErfIndexComponent(config, IsUsed(config), fileName, staticFactors)
{
}

NRTYServer::IParsedEntity::TPtr TRTYErfIndexComponent::BuildParsedEntity(NRTYServer::IParsedEntity::TConstructParams& params) const {
    return new TRTYErfParsedEntity(params);
}

NRTYServer::IComponentParser::TPtr TRTYErfIndexComponent::BuildParser() const {
    return new TRTYErfComponentParser(*this);
}

THolder<NRTYServer::IIndexComponentManager> TRTYErfIndexComponent::CreateManager(const NRTYServer::TManagerConstructionContext& context) const {
    TRTYErfDiskManager::TCreationContext cc(context.Dir, BaseFileName, &StaticFactors, Config.IsReadOnly,
                                        (context.UseGlobalFileMapping && context.IndexType == IIndexController::FINAL));
    return MakeHolder<TRTYErfDiskManager>(cc, GetName());
}

THolder<NRTYServer::IIndexComponentBuilder> TRTYErfIndexComponent::CreateBuilder(const NRTYServer::TBuilderConstructionContext& context) const {
    const ui32 blocksCountStart = Min<ui32>(context.Config.MaxDocuments * 2, (1 << DOC_LEVEL_Bits));
    auto manager = new TRTYErfMemoryManager(blocksCountStart, &StaticFactors, GetName());
    if (context.Config.GetType() == "memory") {
        return MakeHolder<TRTYMemoryErfBuilder>(manager, GetName());
    } else {
        TRTYErfDiskManager::TCreationContext diskManagerCtx(TPathName{context.DirConfig.TempDir}, BaseFileName, &StaticFactors, Config.IsReadOnly);
        return MakeHolder<TRTYDiskErfBuilder>(manager, diskManagerCtx, GetName());
    }
}

const TRTYErfIndexComponent::TIndexFiles& TRTYErfIndexComponent::GetIndexFiles() const {
    return IndexFiles;
}

class TErfDestWriter : public NRTYMerger::TPlainFileDestWriter {
public:
    TErfDestWriter(const TFsPath& fileName, const TThrottle::TOptions& writeOption, const IRTYStaticFactors* factors)
        : NRTYMerger::TPlainFileDestWriter(fileName, writeOption)
        , Header(factors, fileName.Basename())
        , HeaderPath(fileName.Parent())
    {
        Header.LoadHeader(HeaderPath);
    }

    bool WriteHeader(ui32 docsCount) override {
        Header.SetDocsCount(docsCount);
        Header.SaveHeader(HeaderPath);
        return true;
    }

private:
    TRTYErfDiskHeader Header;
    TString HeaderPath;
};

bool TRTYErfIndexComponent::DoMerge(const NRTYServer::TMergeContext& rtyContext) const {
    const TRTYMerger::TContext& context = rtyContext.Context;
    NRTYMerger::TBlockMerger::TSourceReaders sources;
    NRTYMerger::TBlockMerger::TDestWriters dests;
    TRTYErfDiskHeader header(&StaticFactors, BaseFileName);
    size_t blockSize = Max<size_t>();
    for (TVector<TString>::const_iterator i = context.Sources.begin(), e = context.Sources.end(); i != e; ++i) {
        header.LoadHeader(*i);
        size_t newBlockSize = header.GetSumSizeBytes();
        VERIFY_WITH_LOG(blockSize == Max<size_t>() || blockSize == newBlockSize, "Inconsist indexes");
        blockSize = newBlockSize;
        sources.push_back(new NRTYMerger::TPlainFileSourceReader(*i + "/" + BaseFileName, blockSize));
    }

    for (size_t i = 0; i < context.Dests.size(); ++i) {
        const TFsPath dest = context.Dests[i];
        dests.push_back(new TErfDestWriter(dest / BaseFileName, Config.GetMergerConfig().WriteOption, &StaticFactors));
    }
    NRTYMerger::TBlockMerger merger(sources, dests, *context.Decoder);
    NRTYMerger::TContext mergerContext;
    mergerContext.StopFlag = rtyContext.RigidStopSignal;
    return merger.Merge(mergerContext);
}

bool TRTYErfIndexComponent::DoAllRight(const NRTYServer::TNormalizerContext& context) const {
    TRTYErfDiskHeader edh(&StaticFactors, BaseFileName);
    VERIFY_WITH_LOG(edh.LoadHeader(context.Dir.PathName()), "Broken erf header");
    if (edh.NeedInRemap())
        return false;
    const TString filePath = context.Dir.PathName() + "/" + BaseFileName;
    const TString hdrPath = filePath + ".hdr";
    return NFs::Exists(filePath) && NFs::Exists(hdrPath);
}

void TRTYErfIndexComponent::CheckAndFix(const NRTYServer::TNormalizerContext& context) const {
    TRTYErfDiskHeader edh(&StaticFactors, BaseFileName);
    VERIFY_WITH_LOG(edh.LoadHeader(context.Dir.PathName()), "Broken erf header");
    if (edh.NeedInRemap()) {
        NRTYServer::IIndexOwner::TGuardIndexModification g(context.Index);
        edh.Remap();
    }
    const TString filePath = context.Dir.PathName() + "/" + BaseFileName;
    const TString hdrPath = filePath + ".hdr";
    const bool isExistsFile = NFs::Exists(filePath);
    const bool isExistsHdr = NFs::Exists(hdrPath);
    bool isReadOnly = Config.IsReadOnly;
    if (!isExistsFile || !isExistsHdr) {
        ui32 docCount = context.Managers.GetDocumentsCount();
        NRTYServer::IIndexOwner::TGuardIndexModification g(context.Index);
        if (!isExistsHdr) {
            DEBUG_LOG << hdrPath << "does not exist. Create it" << Endl;
            edh.SetDocsCount(docCount);
            edh.SaveHeader(context.Dir.PathName());
        }
        if (!isExistsFile) {
            DEBUG_LOG << filePath << "does not exist. Create it" << Endl;
            VERIFY_WITH_LOG(!isReadOnly, "Can't fix nonexistent file in read-only mode");

            TFile f(filePath, OpenAlways | RdWr);
            const size_t len = edh.GetFactorsCount() * docCount * sizeof(float);
            f.Resize(len);
            TFileMap map(f, TMemoryMapCommon::oRdWr);
            map.Map(0, len);
            memset(map.Ptr(), 0, len);
        }
    }
}

bool TRTYErfIndexComponent::GetInfoChecker(NRTYServer::TInfoChecker& info) const {
    info.Version = 1;
    info.AdditionalHash = StaticFactors.GetHash();
    return true;
}

const NRTYFactors::TRankModelHolder* TRTYErfIndexComponent::GetPruningFormula() const {
    if (Config.Pruning->GetType() == TPruningConfig::FORMULA) {
        return Config.GetSearcherConfig().Factors->GetRankModel(Config.Pruning->ToString().substr(NRTYServer::FormulaPruningPrefix.size()));
    } else {
        return nullptr;
    }
}

const IRTYStaticFactors& TRTYErfIndexComponent::GetStaticFactors() const {
    return StaticFactors;
}

TRTYStaticFactorsConfigHolder::TRTYStaticFactorsConfigHolder(THolder<IRTYStaticFactors> factorsConfig)
    : FactorsConfig(std::move(factorsConfig))
{}

TRTYStaticFactorsConfigHolder::~TRTYStaticFactorsConfigHolder() = default;

TRTYStaticErfIndexComponent::TRTYStaticErfIndexComponent(const TRTYServerConfig& config)
    : TRTYStaticFactorsConfigHolder(MakeHolder<TRTYStaticFactorsConfig>(config.GetSearcherConfig().Factors.Get()))
    , TRTYErfIndexComponent(config, "indexerf.rty", *FactorsConfig)
{
}

class TFormulaPruningCalcer : public TPruningConfig::ICalcer {
public:
    TFormulaPruningCalcer(const IRTYErfManager* manager, const NRTYFactors::TConfig& factors, const NRTYFactors::TRankModelHolder& formula)
        : Manager(manager)
        , Factors(factors)
        , Formula(formula)
    {
        Y_UNUSED(Factors);
    }

    double PruningRank(ui32 docid) const override {
        CHECK_WITH_LOG(Manager);
        return Calculate([this, docid](TFactorView& factorStorage) {
            Manager->Read(factorStorage, docid);
        });
    }

    double PruningRankByDoc(const TPruningConfig::IDocument* document) const override {
        return Calculate([document](TFactorView& factorStorage) {
            document->FillFactorStorage(factorStorage);
        });
    }

private:
    template <class F>
    double Calculate(F fill) const {
        //Note: factorStorage here uses the common relev.conf layout (it is not an ERF block)
        TBasicFactorStorage factorStorage(N_FACTOR_COUNT);
        TFactorView factorView(factorStorage);
        fill(factorView);

        float result = 0;
        Formula.MultiCalc(&factorStorage.factors, &result, 1);
        return static_cast<double>(result);
    }

private:
    const IRTYErfManager* Manager;
    const NRTYFactors::TConfig& Factors;
    const NRTYFactors::TRankModelHolder& Formula;
};

THolder<TPruningConfig::ICalcer> TRTYStaticErfIndexComponent::CreatePruningCalcer(const NRTYServer::IIndexManagersStorage* managers) const {
    if (Config.Pruning->GetType() == TPruningConfig::FORMULA) {
        const auto formula = GetPruningFormula();
        return MakeHolder<TFormulaPruningCalcer>(managers ? managers->GetManager<IRTYErfManager>(ERF_COMPONENT_NAME) : nullptr, *Config.GetSearcherConfig().Factors, *formula);
    } else {
        return nullptr;
    }
}
