#include <saas/rtyserver/common/search_area_modifier.h>

#include <search/meta/context.h>
#include <search/meta/mergedres.h>
#include <search/meta/rearrange/rearrange.h>
#include <search/web/core/rule.h>

#include <tuple>

namespace NRTYServer {
    class TAntiDupRule: public IRearrangeRule {
    private:
        class TContext: public IRearrangeRuleContext {
        public:
            TContext(ISearchAreaModifier* searchAreaModifier)
                : SearchAreaModifier(searchAreaModifier)
            {
                CHECK_WITH_LOG(SearchAreaModifier);
            }

            void DoAdjustClientParams(const TAdjustParams& params) override {
                Disabled = LocalScheme()["Disabled"].IsTrue();

                if (!params.ClientRequestAdjuster->ClientFormFieldHas("pron", "earlyurls")) {
                    params.ClientRequestAdjuster->ClientFormFieldInsert("pron", "earlyurls");
                }
            }

            void DoRearrangeAfterMerge(TRearrangeParams& rearrangeParams) override {
                if (Disabled) {
                    return;
                }

                TMetaGrouping* g = rearrangeParams.Current.Grouping;
                EGroupingMode mode = rearrangeParams.Current.Params.Mode;
                if (!g) {
                    return;
                }

                TMap<TMsString, TGroupingDoc> docs;
                TSet<TGroupingDoc> removed;
                THolder<ISearchLocker::ILock> lock;
                for (size_t groupIdx = 0; groupIdx < g->Size(); ++groupIdx) {
                    if (mode != GM_FLAT) {
                        docs.clear();
                    }

                    const TMetaGroup& group = g->GetMetaGroup(groupIdx);
                    for (size_t docIdx = 0; docIdx < group.MetaDocs.size(); ++docIdx) {
                        const TGroupingDoc d = {
                            groupIdx,
                            group.MetaDocs[docIdx]
                        };
                        const TMsString url = d.Doc.Url();
                        if (url.empty()) {
                            rearrangeParams.InsertRearrangeError("EmptyUrl", ToString(groupIdx) + "/" + ToString(docIdx));
                            continue;
                        }

                        auto insertion = docs.emplace(url, d);
                        if (!insertion.second) {
                            if (!lock) {
                                lock.Reset(SearchAreaModifier->GetSearchLock());
                            }

                            const TGroupingDoc& current = insertion.first->second;
                            const TMergedDoc& currentDoc = current.Doc;
                            const ui32 currentTimestamp = SearchAreaModifier->GetTimestamp(currentDoc.ClientNum(), currentDoc->Handle);

                            const TMergedDoc& candidateDoc = d.Doc;
                            const ui32 candidateTimestamp = SearchAreaModifier->GetTimestamp(candidateDoc.ClientNum(), candidateDoc->Handle);

                            if (std::make_tuple(currentTimestamp, currentDoc.ClientNum()) < std::make_tuple(candidateTimestamp, candidateDoc.ClientNum())) {
                                removed.insert(current);
                                docs[url] = d;
                            } else {
                                removed.insert(d);
                            }
                        }
                    }
                }
                lock.Destroy();

                TFoundStatistics64 totalNumDocs = g->NumDocs;
                TFoundStatistics totalNumGroups = g->RelevStat;
                for (auto i = removed.rbegin(); i != removed.rend(); ++i) {
                    const size_t groupIdx = i->GroupIdx;
                    const TMergedDoc& doc = i->Doc;
                    Y_ENSURE(groupIdx < g->Size(), "incorrect GroupIdx " << groupIdx);

                    TMetaGroup& group = g->GetMetaGroup(groupIdx);
                    {
                        group.RemoveDoc(doc.Doc());
                        for (size_t p = 0; p < NUM_PRIORITIES; p++) {
                            if (group.NumDocs[p] > 0) {
                                group.NumDocs[p] -= 1;
                            }
                            if (totalNumDocs[p] > 0) {
                                totalNumDocs[p] -= 1;
                            }
                        }
                        rearrangeParams.InsertWorkedRule("RemovedDoc", ToString(groupIdx) + "/" + doc.Url());
                    }
                    if (group.MetaDocs.empty()) {
                        g->Erase(groupIdx);
                        for (size_t p = 0; p < NUM_PRIORITIES; p++) {
                            if (totalNumGroups[p] > 0) {
                                totalNumGroups[p] -= 1;
                            }
                        }
                        rearrangeParams.InsertWorkedRule("RemoveGroup", ToString(groupIdx));
                    }
                }
                g->SetRelevStats(totalNumDocs, totalNumGroups);
            }

        private:
            struct TGroupingDoc {
                size_t GroupIdx;
                TMergedDoc Doc;

                bool operator<(const TGroupingDoc& other) const {
                    return GroupIdx < other.GroupIdx;
                }
            };

        private:
            ISearchAreaModifier* SearchAreaModifier = nullptr;
            bool Disabled = false;
        };

    public:
        TAntiDupRule(const TString& /*config*/, const TSearchConfig& /*searchConfig*/)
            : SearchAreaModifier(Singleton<TSearchAreaModifier>()->GetModifier())
        {
            CHECK_WITH_LOG(SearchAreaModifier);
        }

        IRearrangeRuleContext* DoConstructContext() const override {
            return new TContext(SearchAreaModifier);
        }

    private:
        ISearchAreaModifier* SearchAreaModifier = nullptr;
    };

    IRearrangeRule* CreateAntiDupRule(const TString& config, const TSearchConfig& searchConfig) {
        return new TAntiDupRule(config, searchConfig);
    }
}

REGISTER_REARRANGE_RULE(RTYServerAntiDup, NRTYServer::CreateAntiDupRule);
