#pragma once

#include <saas/rtyserver/search/processors/processor.h>
#include <library/cpp/logger/global/global.h>
#include <kernel/groupattrs/docsattrs.h>
#include <search/rank/docalc_docrelev.h>
#include <search/idl/meta.pb.h>

class THitsBuilder: public IRTYDocProcessor {
private:
    ui64 Counter;
public:

    void FillProperties(TSearcherProps* props) const override {
        props->Set("rty_hits_count", ToString(Counter));
    }

    THitsBuilder()
        : Counter(0)
    {
    }

    using IRTYDocProcessor::Process;

    void Process(ui32 /*docId*/, const NGroupingAttrs::TDocsAttrs& /*da*/, const TAllDocPositions& /*doc*/, const TTRIteratorHitInfos* /*hitInfos*/) override {
        ++Counter;
    }
};

class THitsBuilderDetail: public IRTYDocProcessor {
private:

    struct TDocHitsStorage {
        TVector<TVector<ui32>> HitCountByLevel;
        ui64 HitMask;
        TDocHitsStorage()
            : HitCountByLevel(64,
                TVector<ui32>(static_cast<ui32>(EFormClass::NUM_FORM_CLASSES))
            ) {
        }
        void ProcessHit(const TFullPositionEx& pos) {
            HitCountByLevel[pos.WordIdx][static_cast<int>(TWordPosition::Form(pos.Pos.Beg))]++;
            HitMask |= 1 << pos.WordIdx;
        }
        ui32& GetHitCount(int wordIdx, int hitLevel) {
            return HitCountByLevel[wordIdx][hitLevel];
        }

        ui32 GetHitCount(int wordIdx, int hitLevel) const {
            return HitCountByLevel[wordIdx][hitLevel];
        }
        ui64 GetHitMask() const {
            return HitMask;
        }
        NJson::TJsonValue SerializeAsJson(const TMaybe<TTRIteratorHitInfos>& HitInfos) {
            NJson::TJsonValue jsonHits(NJson::JSON_MAP);
            for (size_t i = 0; i < HitCountByLevel.size(); i++) {
                ui32 totalHits = Accumulate(HitCountByLevel[i].begin(), HitCountByLevel[i].end(), 0);
                if (totalHits > 0) {
                    NJson::TJsonValue countsArray;
                    for (auto hitCount : HitCountByLevel[i]) {
                        countsArray.AppendValue(hitCount);
                    }
                    NJson::TJsonValue wordHitInfo;
                    wordHitInfo.InsertValue("count_by_level", countsArray);
                    wordHitInfo.InsertValue("count", totalHits);

                    if (HitInfos) {
                        wordHitInfo.InsertValue("text", HitInfos->at(i).Text);
                        wordHitInfo.InsertValue("lemma", HitInfos->at(i).Lemma);
                        wordHitInfo.InsertValue("word_weight", HitInfos->at(i).WordWeight);
                    }
                    jsonHits[ToString(i)] = wordHitInfo;
                }
            }
            return jsonHits;
        }
    };

    void SerializeHitsInfo(ui32 docId, TArrayRef<const char * const> attrNames, IAttributeWriter& write);
    void SerializeQueryHitsInfo(ui32 docId, TArrayRef<const char * const> attrNames, IAttributeWriter& write);

    ui64 Counter;
    ui64 CounterAll;
    ui64 CounterBreaks;
    THashMap<ui32, TVector<TFullPositionEx>> Hits;
    // map below: TextHitsByWord[docId].GetHitCount(wordIdx, hitLevel) = hitCount
    THashMap<ui32, TDocHitsStorage> TextHits;
    THashMap<ui32, TDocHitsStorage> AnnHits;
    TMaybe<TTRIteratorHitInfos> HitInfos = Nothing();
public:

    void SerializeFirstStageAttributes(ui32 docId, TArrayRef<const char * const> attrNames, IAttributeWriter& write) override;

    void FillProperties(TSearcherProps* props) const override {
        props->Set("rty_hits_count", ToString(Counter));
        props->Set("rty_hits_count_full", ToString(CounterAll));
        props->Set("rty_hits_count_sents", ToString(CounterBreaks));

        if (HitInfos.Defined()) {
            NJson::TJsonValue result;
            for (ui32 i = 0; i < HitInfos->size(); i++) {
                NJson::TJsonValue wordHitInfo;
                wordHitInfo.InsertValue("text", HitInfos->at(i).Text);
                wordHitInfo.InsertValue("lemma", HitInfos->at(i).Lemma);
                wordHitInfo.InsertValue("word_weight", HitInfos->at(i).WordWeight);
                result[ToString(i)] = wordHitInfo;
            }
            props->Set("rty_query_stats", result.GetStringRobust());
        }
    }

    THitsBuilderDetail()
        : Counter(0)
        , CounterAll(0)
        , CounterBreaks(0)
    {
    }

    void Process(ui32 docId, const NGroupingAttrs::TDocsAttrs& /*da*/, const TAllDocPositions& doc, const TTRIteratorHitInfos* hitInfos) override {
        if (HitInfos.Empty()) {
            HitInfos = *hitInfos;
        }
        ++Counter;
        CounterAll += doc.TextHits.Count;
        TSet<i32> breaks;
        for (int i = 0; i < doc.TextHits.Count; ++i) {
            const TFullPositionEx& pos = doc.TextHits.Pos[i];
            Hits[docId].push_back(pos);
            breaks.insert(TWordPosition::Break(pos.Pos.Beg));
            TextHits[docId].ProcessHit(pos);
        }
        for (int i = 0; i < doc.QueryHits.Count; ++i) { // QueryHit -- hit from annotations
            const TFullPositionEx& pos = doc.QueryHits.Pos[i];
            AnnHits[docId].ProcessHit(pos);
        }
        CounterBreaks += breaks.size();
    }
};
