#include "searcher.h"

#include "config.h"
#include "manager.h"
#include "trie_search_request.h"

#include <saas/rtyserver/components/fullarchive/manager.h>

#include <library/cpp/threading/future/future.h>

namespace NRTYServer {
    template<typename TResult>
    struct TParallelExecutor {
        using TFullResult = TVector<TResult>;

        TParallelExecutor(size_t nThreads) {
            Queue = MakeHolder<TThreadPool>(TThreadPool::TParams().SetBlocking(true).SetCatching(false).SetThreadName("TrieSearcher"));
            Queue->Start(nThreads, nThreads * 100);
        }

        template<typename TJob>
        NThreading::TFuture<TResult> AddJob(TJob job) {
            auto promise = NThreading::NewPromise<TResult>();
            bool added = Queue->AddFunc([promise, job = std::move(job)]() mutable {
                try {
                    promise.SetValue(job());
                } catch (...) {
                    promise.SetException(std::current_exception());
                }
            });
            if (!added) {
                promise.SetException("Can not add a search request to the queue because it is going to stop");
            }
            return promise.GetFuture();
        }

    private:
        THolder<IThreadPool> Queue;
    };

    struct TTrieSingleThreadSearcher : ITrieMultiSegmentSearcher {
        explicit TTrieSingleThreadSearcher(TString indexGenerator)
            : IndexGenerator(std::move(indexGenerator))
        {
        }

        void DoSearch(TTrieSearchRequest& request,
                      const TCgiParameters& cgiParams,
                      const TVector<IIndexController::TPtr>& controllers) final {
            for (auto& controller : controllers) {
                auto trie = dynamic_cast<const TTrieIndexManager*>(controller->GetManager(TrieComponentName));
                auto fullArc = dynamic_cast<const IFAManager*>(controller->GetManager(IndexGenerator));
                Y_VERIFY(trie);
                Y_VERIFY(fullArc);
                ui32 nFound = trie->DoSearch(request, *fullArc, cgiParams, request.GetReport());
                if (request.DecreaseMaxDocs(nFound) == 0) {
                    break;
                }
            }
        }

    private:
        TString IndexGenerator;
    };

    using TReportPiece = TList<NMetaProtocol::TDocument>;

    struct TReportPieceBuilder : ISimpleReportBuilder {
        void AddDocument(NMetaProtocol::TDocument& doc) final {
            Documents.emplace_back();
            Documents.back().Swap(&doc);
        }

        void AddDocuments(TReportPiece documents) final {
            Documents.splice(Documents.end(), std::move(documents));
        }

        TReportPiece Take() {
            return std::move(Documents);
        }

    private:
        TReportPiece Documents;
    };

    struct TTrieMultiThreadSearcher : ITrieMultiSegmentSearcher {
        TTrieMultiThreadSearcher(TString indexGenerator, ui32 nThreads)
            : Executor(nThreads)
            , IndexGenerator(std::move(indexGenerator))
        {
        }

        void DoSearch(TTrieSearchRequest& mutableRequest,
                      const TCgiParameters& cgiParams,
                      const TVector<IIndexController::TPtr>& controllers) final {
            const auto& request = mutableRequest;
            size_t nControllers = controllers.size();
            TVector<NThreading::TFuture<TReportPiece>> futures(nControllers);
            for (size_t i = 0; i < nControllers; ++i) {
                auto trie = dynamic_cast<const TTrieIndexManager*>(controllers[i]->GetManager(TrieComponentName));
                auto fullArc = dynamic_cast<const IFAManager*>(controllers[i]->GetManager(IndexGenerator));
                Y_VERIFY(trie);
                Y_VERIFY(fullArc);
                futures[i] = Executor.AddJob([&request, trie, fullArc, &cgiParams]() -> TReportPiece {
                    TReportPieceBuilder piece;
                    trie->DoSearch(request, *fullArc, cgiParams, piece);
                    return piece.Take();
                });
            }
            auto& report = mutableRequest.GetReport();
            for (size_t i = 0; i < nControllers; ++i) {
                auto reportPiece = futures[i].ExtractValue(TDuration::Max());
                if (reportPiece.size() <= request.GetMaxDocs()) {
                    mutableRequest.DecreaseMaxDocs(reportPiece.size());
                    report.AddDocuments(std::move(reportPiece));
                } else {
                    for (auto& doc : reportPiece) {
                        if (request.GetMaxDocs() == 0) {
                            return;
                        }
                        report.AddDocument(doc);
                        mutableRequest.DecreaseMaxDocs(1);
                    }
                }
            }
        }

    private:
        TParallelExecutor<TReportPiece> Executor;
        TString IndexGenerator;
    };

    THolder<ITrieMultiSegmentSearcher> CreateTrieSearcher(TString indexGenerator, ui32 nThreads) {
        if (nThreads == 0) {
            return MakeHolder<TTrieSingleThreadSearcher>(std::move(indexGenerator));
        }
        return MakeHolder<TTrieMultiThreadSearcher>(std::move(indexGenerator), nThreads);
    }
}
