#include "test_partial_run.h"

#include <kernel/doom/offroad_erf_wad/erf_io.h>
#include <kernel/doom/wad/wad.h>
#include <kernel/urlid/doc_handle.h>

#include <robot/jupiter/library/rtdoc/printer/text_utils.h>

#include <ysite/yandex/erf/flat_hash_table.h>
#include <ysite/yandex/erf_format/erf_format.h>

#include <library/cpp/offroad/flat/flat_ui64_searcher.h>
#include <library/cpp/testing/unittest/env.h>
#include <library/cpp/testing/unittest/registar.h>

#include <util/stream/str.h>


// checks panther dssm routines with old input
class TFusionPantherDssmSmokeTest: public NFusion::TTestPartialRunTestBase {
public:
    UNIT_TEST_SUITE(TFusionPantherDssmSmokeTest)
    UNIT_TEST(TestPantherWithMercuryOldData);
    UNIT_TEST(TestPantherWithMercuryVer35Data);
    UNIT_TEST_SUITE_END();

public:
    TString ModelsDir;
public:
    void Init(const TStringBuf resProtoBin) {
        BundleFilter.Exclusive = true;
        BundleFilter.DisplayNames = {
            "BuildTextKeyInvAndCounts", "BuildTextKeyInvAndCounts", "BuildLinkKeyInvCounts", "BuildLinkKeyInvCounts",
            "MergeCounts", "RebuildCountsSaas", "BuildPanther", "MergePanther", "TempFixPanther",
            "BuildPantherDssmEmbeddings",
            "MergeL1"
        };

        IncomingLumpsFilter = {
            // for counts
            "calculated_attrs_with_keyinv_doc_chunk_mapping_single_chunk",
            "walrus_with_keyinv_doc_chunk_mapping_single_chunk",
            "external_relev_attrs",
            // for panther
            "panther_dssm_url_title_embeddings",
            "erf_herf_features",
            // for l1
            "l1_doc_lumps", "l1_global_lumps"
        };

        ModelsDir = GetTestResource(TFsPath("rty_models_fake_for_panther") / "models.rty");
        TFsPath messagesDump = GetTestResource(TFsPath(resProtoBin) / "messages.protobin");
        messagesDump.CheckExists();
        MakeAllStagesInputs(messagesDump, TFsPath::Cwd() / TestWorkDir);
    }

    void OnTaskAdded(NRtDoc::TBuilderTask& task) override {
        task.SetModelsDir(ModelsDir);
    }

    void SetUp() override {
        Cleanup();
        TFsPath(TestWorkDir).MkDirs();
    }

    void TestPantherWithMercuryOldData() {
        Init("erf_old_model");
        BuildAll();
    }

    void TestPantherWithMercuryVer35Data() {
        Init("mercury_35_rc");
        BuildAll();
    }
};

UNIT_TEST_SUITE_REGISTRATION(TFusionPantherDssmSmokeTest);


