#pragma once

#include <saas/rtyserver/common/common_rty.h>
#include <search/session/searcherprops.h>
#include <kernel/externalrelev/relev.h>
#include <util/generic/cast.h>

class IRTYDocProcessor: public IDocProcessor {
public:
    // interface IDocProcessor
    void Process(ui32 docId, const NGroupingAttrs::TDocsAttrs& da, const TAllDocPositions& doc, const TTRIteratorHitInfos* hitInfos) override = 0;

    bool DocumentBanned(ui32 /*docId*/) const override {
        return false;
    }

    bool IsBatchProcessor() const override {
        return false;
    }

public:
    // new methods
    virtual void FillProperties(TSearcherProps* props) const = 0;

    virtual bool IsFilterProcessor() const {
        return !IsBatchProcessor();
    }
};

// Base class for non-filtering doc facets (that should be calculated after filter border)
class TRTYFacetBase: public IRTYDocProcessor {
public:
    bool DocumentBanned(ui32 /*docId*/) const override final {
        return false;
    }

    bool IsBatchProcessor() const override final {
        return true;
    }

    void Process(ui32 docId, const NGroupingAttrs::TDocsAttrs& da, const TAllDocPositions& /*doc*/, const TTRIteratorHitInfos* /*hitInfos*/) override {
        // Most RTY facets do not need hits. Those who need should override Process() insted of ProcessDoc().
        return ProcessDoc(docId, da);
    }

    void ProcessBatch(TArrayRef<const ui32> docIds, const NGroupingAttrs::TDocsAttrs& da) override {
        // default implementation
        for (ui32 docId : docIds) {
            ProcessDoc(docId, da);
        }
    }

    virtual void ProcessDoc(ui32 docId, const NGroupingAttrs::TDocsAttrs& da) = 0;
};

class TDocProcessors: public IRTYDocProcessor {
private:
    TVector<IDocProcessor*> FilterProcessors;
    TVector<IDocProcessor*> BatchProcessors;
    TVector<IRTYDocProcessor*> AllRtyProcessors;
    TVector<TAtomicSharedPtr<IDocProcessor>> AllProcessors;

public:
    virtual void Process(ui32 docId, const NGroupingAttrs::TDocsAttrs& da, const TAllDocPositions& doc, const TTRIteratorHitInfos* hitInfos) override {
        for (auto&& proc : FilterProcessors) {
            proc->Process(docId, da, doc, hitInfos);
        }
    }

    virtual bool DocumentBanned(ui32 docId) const override {
        for (auto&& proc : FilterProcessors) {
            if (proc->DocumentBanned(docId)) {
                return true;
            }
        }
        return false;
    }

    bool IsFilterProcessor() const override final {
        return !FilterProcessors.empty();
    }

    bool IsBatchProcessor() const override final {
        return !BatchProcessors.empty();
    }

    virtual void ProcessBatch(TArrayRef<const ui32> docIds, const NGroupingAttrs::TDocsAttrs& da) override {
        for (auto&& proc : BatchProcessors) {
            proc->ProcessBatch(docIds, da);
        }
    }

    virtual void SerializeFirstStageAttributes(ui32 docId, TArrayRef<const char * const> attrNames, IAttributeWriter& write) override {
        for (auto&& proc : AllProcessors) {
            proc->SerializeFirstStageAttributes(docId, attrNames, write);
        }
    }

    virtual void FillProperties(TSearcherProps* props) const override {
        for (auto&& proc : AllRtyProcessors) {
            proc->FillProperties(props);
        }
    }

protected:
    void AddToStages(IDocProcessor* proc, bool isFilterProcessor, bool isBatchProcessor) {
        if (isFilterProcessor) {
            FilterProcessors.push_back(proc);
        }
        if (isBatchProcessor) {
            BatchProcessors.push_back(proc);
        }
    }

    void Add(TVector<TAtomicSharedPtr<IRTYDocProcessor>>& processors, bool forceNoBatch) {
        for (auto& item : processors) {
            IRTYDocProcessor* rtyProc = item.Get();
            AllProcessors.emplace_back(std::move(item));
            AllRtyProcessors.push_back(rtyProc);
            if (!forceNoBatch) {
                AddToStages(rtyProc, rtyProc->IsFilterProcessor(), rtyProc->IsBatchProcessor());
            } else {
                AddToStages(rtyProc, /*isFilterProcessor=*/true, /*isBatchProcessor=*/false);
            }
        }
    }

public:
    // A method for adding generic IDocProcessor implementations (including nested objects of TDocProcessors)
    void Add(TAtomicSharedPtr<IDocProcessor> proc) {
        CHECK_WITH_LOG(proc);
        AllProcessors.push_back(proc);

        auto* rtyProc = dynamic_cast<IRTYDocProcessor*>(proc.Get());
        if (Y_LIKELY(rtyProc)) {
            AllRtyProcessors.push_back(rtyProc); // for FillProperties
            AddToStages(rtyProc, rtyProc->IsFilterProcessor(), rtyProc->IsBatchProcessor());
        } else {
            Y_ASSERT(0 /* this branch is not used now - remove the assert if your code breaks */);
            AddToStages(proc.Get(), /*isFilterProcessor=*/true, proc->IsBatchProcessor());
        }
    }

    bool Empty() const {
        return AllProcessors.empty();
    }
};
