#pragma once

#include <robot/library/oxygen/indexer/library/pruning_builder/pruning_builder.h>
#include <robot/library/oxygen/indexer/idspace/docid.h>

#include <library/cpp/langs/langs.h>

#include <util/generic/function.h>
#include <util/generic/maybe.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>
#include <util/generic/yexception.h>



using TRankValues = TVector<float>;

struct TPruningZoneParams {
    TVector<ui32> Proportions;
};

class TPruningDataAccessor;


//The internal data structure for TPruningDataAccessor and TFreshPruningBuilder implementations
struct TPruningDataArray {
    struct TEntry {
        ui32 TmpDocId = 0;
        float RegionalSelectionRank; //Value obtained from regional ISelectionRankSource
        TString InvertedUrl;
    };

    TVector<TEntry> Data;
    TVector<TRankValues> Ranks;

    inline float GetRank(ui32 idx, ui32 nSource) const {
        return Ranks[nSource][idx];
    }

    inline bool IsBusy(ui32 idx) const {
        return !IsUndefined(Data[idx].RegionalSelectionRank);
    }

    inline bool IsUndefined(ui32 idx, ui32 nSource) const {
        return IsUndefined(GetRank(idx, nSource));
    }

    inline void SetSource(ui32 idx, ui32 nSource) {
        Data[idx].RegionalSelectionRank = Ranks[nSource][idx];
    }

    inline void ResetSelection() {
        for (auto& entry : Data)
            entry.RegionalSelectionRank = UndefinedRankValue;
    };

    static constexpr float UndefinedRankValue = std::numeric_limits<float>::min();

    inline static bool IsUndefined(float f) {
        return f == UndefinedRankValue;
    }
};

class TPruningDataAccessor : public NOxygen::IPruningDataAccessor, private TPruningDataArray {
public:
    TPruningDataAccessor(TVector<TString> &&rankNames);

    // Returns false if url is bad, like CheckUrl
    bool AddEntry(ui32 tmpDocId, TStringBuf url, const TRankValues& rankValues);
    void Reserve(ui32 nEntries);

    virtual TVector<ui32> GetTmpDocIds() override;
    virtual TString GetInvertedUrl(ui32 docId) override;
    virtual float GetMainRank(ui32 docId) override;
    virtual NOxygen::TPruningRanks GetAllRanks(ui32 docId) override;


    //Additional access methods
    ui32 GetRanksNum() const;

    float GetSelectionRank(ui32 docId, int nSource) const {
        return GetRank(GetIndex(docId), nSource);
    }

    TPruningDataArray& GetInternalArray() {
        return (TPruningDataArray&)(*this);
    }

    inline ui32 GetInternalIndex(ui32 tmpDocId) const {
        return GetIndex(tmpDocId);
    }

    inline const TString& GetRankName(int nSource) const {
        return RankNames[nSource];
    }

private:
    TVector<ui32> DocId2Index;
    TVector<TString> RankNames;

private:
    inline ui32 GetIndex(ui32 tmpDocId) const {
        return DocId2Index[tmpDocId];
    }

    inline TPruningDataArray::TEntry& GetByDocId(ui32 tmpDocId) {
        return Data[GetIndex(tmpDocId)];
    }
};


class TFreshPruningBuilder : public NOxygen::TPruningBuilder {
public:
    TFreshPruningBuilder(const NOxygen::TPruningConfig& pruningConfig, const TPruningZoneParams& zoneParams, TPruningDataAccessor& data);

    void DumpSln(IOutputStream* out) const override;

protected:
    void CalculatePruning(NOxygen::TPruningEntryList* pruningEntryList, const NOxygen::TPruningConfig& config) override;

private:
    TVector<ui32> Proportions;
    TPruningDataAccessor& Data;
};
