#include <robot/pluto/library/common/common.h>
#include <robot/pluto/library/libkeyinv/key_inv.h>
#include <robot/pluto/library/libpanther/panther.h>

#include <kernel/dssm_applier/decompression/decompression.h>

#include <quality/nirvana_tools/conveyor_operations/train_dssm/apply/model_applier.h>

#include <search/web/core/rule.h>

#include <library/cpp/string_utils/base64/base64.h>


namespace {

    inline TVector<TString> ModelFields() {
        // Required fields only
        return {
            "url",
            "title",
            "unigrams",
            "bigrams",
            "region",
            "query"
        };
    }

    NDssmPoolUtil::EModelType GetModelType(const NRearrConf::TRearrangeRuleParams& params, const TStringBuf extension) noexcept {
        if (auto it = params.find("DssmModelType"); it != params.end()) {
            return it->second == "nn_applier" ? NDssmPoolUtil::EModelType::NNApplier : NDssmPoolUtil::EModelType::Dssm3;
        }

        if (extension == "nn_applier") {
            return NDssmPoolUtil::EModelType::NNApplier;
        }

        if (extension == "dssm3") {
            return NDssmPoolUtil::EModelType::Dssm3;
        }

        // default value for a model without extension
        return NDssmPoolUtil::EModelType::Dssm3;
    }

} // namespace

namespace NSaas {

    struct TFulltextPatentSearchData {
        TFulltextPatentSearchData(const TString& options, const TSearchConfig& config) {
            TFsPath workDir = config.WorkDir;
            const NRearrConf::TRearrangeRuleParams params(options);
            workDir /= params.Value("FulltextPatentSearchDir", TStringBuf("data"));

            auto getPath = [&params, &workDir](const auto& key, const TStringBuf defaultValue) {
                return workDir / params.Value(key, defaultValue);
            };

            const auto directTextPurePath = getPath("DirectTextPurePath", "pure/pure.trie");
            const auto directTextConfigPath = getPath("DirectTextConfigPath", "direct_text_config");
            const auto directTextSeoLnkPath = getPath("DirectTextSeoLnkPath", "seolnk");
            const auto metaDescrDictPath = getPath("MetaDescrDictPath", "dict.dict");
            const auto metaDescrStopwordPath = getPath("MetaDescrStopwordPath", "stopword.lst");
            const auto nameExtractorFioPath = getPath("NameExtractorFioPath", "fio/name_extractor.cfg");
            const auto newshostPath = getPath("NewshostPath", "news_hosts.txt");
            const auto numberPath = getPath("NumberPath", "numbers");
            const auto numeratorConfigsPath = getPath("NumeratorConfigsPath", "numerator_config");
            const auto phoneNumberPath = getPath("PhoneNumberPath", "phone_markers.gzt.bin");
            const auto segmentatorPath = getPath("SegmentatorPath", "2ld.list");
            const auto tasixhostPath = getPath("TasixhostPath", "tasix_hosts.txt");

            const auto dssmModelPath = getPath("DssmModelPath", "panther_dwelltime_dssm_model.nn_applier");
            const auto dssmModelType = GetModelType(params, dssmModelPath.GetExtension());

            KeyInv = MakeAtomicShared<NPluto::TKeyInvGenerator>(
                directTextPurePath,
                directTextConfigPath,
                directTextSeoLnkPath,
                metaDescrDictPath,
                metaDescrStopwordPath,
                nameExtractorFioPath,
                newshostPath,
                numberPath,
                numeratorConfigsPath,
                phoneNumberPath,
                segmentatorPath,
                tasixhostPath,
                false); // mapped files prohibited in searchproxy

            DssmApplyer = NDssmPoolUtil::GetModel(
                dssmModelType,
                dssmModelPath,
                ModelFields(),
                NDssmPoolUtil::ELoadType::Preload);

            Component = params.Value("Component", TStringBuf("HNSW"));
            NormalKvReport = params.Value("NormalKvReport", TStringBuf("1"));
            SkipWizard = params.Value("SkipWizard", TStringBuf("1"));
            DiscreteDssm = params.Value("DiscreteDssm", false);
        }

        TAtomicSharedPtr<NPluto::TKeyInvGenerator> KeyInv;
        NDssmPoolUtil::IModelApplierPtr DssmApplyer;

        bool DiscreteDssm;
        TString Component;
        TString NormalKvReport;
        TString SkipWizard;

        TMutex Mutex;
    };


    class TFulltextPatentSearchRuleContext final : public IRearrangeRuleContext {
    public:
        TFulltextPatentSearchRuleContext(TAtomicSharedPtr<TFulltextPatentSearchData> data)
            : Data(std::move(data))
        {}

        void DoAdjustClientParams(const TAdjustParams& params) override {
            const auto& cgi = params.RequestFields.CgiParam;

            const auto patentText = cgi.find("text");
            if (patentText == cgi.end()) {
                return;
            }

            const auto searchSize = cgi.find("search_size");
            if (searchSize == cgi.end()) {
                return;
            }

            const auto topSize = cgi.find("top_size");
            if (topSize == cgi.end()) {
                return;
            }

            const auto lang = cgi.find("lang");
            const auto language = lang != cgi.end() ? LanguageByName(lang->second) : ELanguage::LANG_UNK;

            const auto base64EncodeFields = cgi.find("base64_encode_fields");
            const TString base64EncodeFieldsValue = (base64EncodeFields == cgi.end()) ? "true" : base64EncodeFields->second;

            const auto text = Base64Decode(patentText->second);
            Y_ENSURE(!text.empty());

            const auto d = Data.Get();
            TGuard _(d->Mutex);

            const auto keyInv = Data->KeyInv->Generate(
                text,
                {},
                0,
                0,
                ECharset::CODES_UTF8,
                language,
                language,
                0,
                MimeTypes::MIME_TEXT);
            InsertRearrangeResult("keyinv", Base64EncodeUrl(keyInv));

            const auto [unigrams, bigrams] = NPluto::BuildPantherTerms(keyInv);
            InsertRearrangeResult("unigrams", Base64EncodeUrl(unigrams));
            InsertRearrangeResult("bigrams", Base64EncodeUrl(bigrams));
            const auto dssm = d->DssmApplyer->ApplySample({"", "", unigrams, bigrams, "", ""}, "doc_embedding");
            const auto dssmPacked = NPluto::EncodeDssm(dssm, d->DiscreteDssm);

            auto compSearch = TString::Join(
                "comp:", d->Component, ";",
                "top_size:", topSize->second, ";",
                "search_size:", searchSize->second, ";",
                "base64_encode_fields:", base64EncodeFieldsValue, ";"
                "embed_packed:", dssmPacked, ";",
                "discr_mode:", d->DiscreteDssm ? TStringBuf("yes") : TStringBuf("no"), ";");

            auto cgiField = [&cgi](const TStringBuf key, const auto& defaultValue) {
                const auto it = cgi.find(key);
                return (it != cgi.end()) ? it->second : defaultValue;
            };

            const TCgiParameters result = {
                {"comp_search", std::move(compSearch)},
                {"component", cgiField("component", d->Component)},
                {"normal_kv_report", cgiField("normal_kv_report", d->NormalKvReport)},
                {"skip-wizard", cgiField("skip-wizard", d->SkipWizard)}
            };

            params.ClientCgiInsert(result);
        }

    private:
        const TAtomicSharedPtr<TFulltextPatentSearchData> Data;
    };


    class TFulltextPatentSearchRule final : public IRearrangeRule {
    public:
        TFulltextPatentSearchRule(const TString& options, const TSearchConfig& config) try
            : Data(MakeAtomicShared<TFulltextPatentSearchData>(options, config))
        {
        } catch (...) {
            Y_FAIL("FulltextPatentSearch initialization failed: %s", CurrentExceptionMessage().c_str());
        }

        IRearrangeRuleContext* DoConstructContext() const override {
            return new TFulltextPatentSearchRuleContext(Data);
        }

    private:
        const TAtomicSharedPtr<TFulltextPatentSearchData> Data;
    };

} // namespace NSaas


IRearrangeRule* CreateFulltextPatentSearchRule(const TString& options, const TSearchConfig& config) {
    return new NSaas::TFulltextPatentSearchRule(options, config);
}

REGISTER_REARRANGE_RULE(FulltextPatentSearch, CreateFulltextPatentSearchRule);
