#include "prepare_ann.h"

#include <google/protobuf/descriptor.h>

#include <kernel/indexann/protos/data.pb.h>
#include <kernel/indexann_data/data.h>

#include <library/cpp/protobuf/util/simple_reflection.h>
#include <library/cpp/protobuf/json/json2proto.h>
#include <library/cpp/json/json_reader.h>
#include <library/cpp/testing/unittest/registar.h>

#include <util/folder/path.h>
#include <util/string/join.h>


namespace {
    void FillAllFields(const google::protobuf::Descriptor* desc, google::protobuf::Message* data) {
        const auto* emptyDummyDescriptor = NIndexAnn::TDummyData::descriptor();
        for (int fieldNo = 0; fieldNo < desc->field_count(); ++fieldNo) {
            const google::protobuf::FieldDescriptor* fld = desc->field(fieldNo);
            const google::protobuf::Reflection* refl = data->GetReflection();
            if (fld->type() == google::protobuf::FieldDescriptor::TYPE_MESSAGE) {
                if (fld->message_type() == emptyDummyDescriptor) {
                    const google::protobuf::Message* child = refl->MutableMessage(data, fld);
                    Y_UNUSED(child); // create an instance of the empty type
                } else {
                    google::protobuf::Message* child = refl->MutableMessage(data, fld);
                    FillAllFields(fld->message_type(), child);
                }
            } else {
                NProtoBuf::TMutableField field(*data, fld);
                field.Set(1);
            }
        }
    }

    TVector<NIndexAnn::EDataType> GetSaasAnnFeaturesFromReflection(const NIndexAnn::IFeaturesProfile& featuresFilter) {
        using TData = NIndexAnn::TRegionData;
        NIndexAnn::TAnnotationRec annRec;
        TData& protoAnn = *annRec.AddData();
        const google::protobuf::Descriptor* msg = TData::descriptor();
        FillAllFields(msg, &protoAnn);

        TVector<NIndexAnn::EDataType> streams;
        NIndexAnn::EnumerateFeatures(annRec, streams);

        TVector<NIndexAnn::EDataType> saasStreams;
        saasStreams.reserve(streams.size());
        for (NIndexAnn::EDataType id : streams) {
            if (featuresFilter.IsEssential(id)) {
                saasStreams.push_back(id);
            }
        }
        return saasStreams;
    }
}

class TSaasRobotAnn: public NUnitTest::TTestBase {
    UNIT_TEST_SUITE(TSaasRobotAnn)
    UNIT_TEST(TestFeaturesListIsActual);
    UNIT_TEST(TestUnpackToSaasData);
    UNIT_TEST(TestSkipNonSaas);
    UNIT_TEST_SUITE_END();

public:
    void TestFeaturesListIsActual() {
        TVector<TString> expectedFeatures, usedFeatures;
        const auto& saasProfile = NIndexAnn::GetSaasProfile();
        TVector<NIndexAnn::EDataType> reflection = GetSaasAnnFeaturesFromReflection(saasProfile);
        for(const NIndexAnn::EDataType id : reflection) {
            expectedFeatures.push_back(NIndexAnn::TDataTypeConcern::Instance().GetName(id));
        }

        for(const NIndexAnn::EDataType id : NIndexAnn::GetSaasFeatures()) {
            usedFeatures.push_back(NIndexAnn::TDataTypeConcern::Instance().GetName(id));
        }
        Sort(expectedFeatures);
        Sort(usedFeatures);
        UNIT_ASSERT_VALUES_EQUAL(JoinSeq(" ", expectedFeatures), JoinSeq(" ", usedFeatures));
    }

    //
    // Obtained as:
    // yt read --proxy arnold //tmp/yrum0612_ann[:#1] --format json | jq .column4 -r |
    //                  jq '.Recs[0]' | egrep -v '"[^"]*": null,?$' | tr -d ' \n' |
    //                  sed -r 's:,([}]):\1:g' | jq . -c | python octal_encode.py -x
    //
    // (where .column4 is NIndexAnn.TIndexAnnSiteData in Json form)
    //
    const TString SomeJupiterAnnotationsRec =
        "{\"Data\":[{\"Region\":0,\"ClickMachine\":{\"VpcgCorrectedClicksSLP\":0.9927266240119934}},{\"Region\":225,\"ClickMachine"
        "\":{\"VpcgCorrectedClicksSLP\":0.9917650818824768}}],\"Text\":\"\303\220\302\262\303\220\302\267\303\220\302\262\303\220\302"
        "\265\303\220\302\271\303\221\302\201\303\221\302\217\303\220\302\272\303\220\302\276\303\221\302\201\303\221\302\202\303\221"
        "\302\200\303\220\302\260\303\220\302\274\303\220\302\270\303\221\302\201\303\220\302\270\303\220\302\275\303\220\302\270\303"
        "\220\302\265\303\220\302\275\303\220\302\276\303\221\302\207\303\220\302\270\",\"TextLanguage\":1}";

    // canonized
    const TString ExpectedRtyAnnData =
        "Sentences { Text: \"\134303\134220\134302\134262\134303\134220\134302\134267\134303\134220\134302\134262\134303\134220\134"
        "302\134265\134303\134220\134302\134271\134303\134221\134302\134201\134303\134221\134302\134217\134303\134220\134302\13427"
        "2\134303\134220\134302\134276\134303\134221\134302\134201\134303\134221\134302\134202\134303\134221\134302\134200\134303\134"
        "220\134302\134260\134303\134220\134302\134274\134303\134220\134302\134270\134303\134221\134302\134201\134303\134220\13430"
        "2\134270\134303\134220\134302\134275\134303\134220\134302\134270\134303\134220\134302\134265\134303\134220\134302\134275\134"
        "303\134220\134302\134276\134303\134221\134302\134207\134303\134220\134302\134270\" TextLanguage: 1 StreamsByRegion { Regi"
        "on: 0 Streams { Name: \"DT_VPCG_CORRECTED_CLICKS_SLP\" Value { Value: \"/Q==\" } } } StreamsByRegion { Region: 225 Stream"
        "s { Name: \"DT_VPCG_CORRECTED_CLICKS_SLP\" Value { Value: \"/A==\" } } } }";


    void TestUnpackToSaasData() {
        NIndexAnn::TIndexAnnSiteData annoSite;
        *annoSite.AddRecs() = NProtobufJson::Json2Proto<NIndexAnn::TAnnotationRec>(SomeJupiterAnnotationsRec);
        NRTYServer::TMessage::TAnnData saasAnnoSite;

        NIndexAnn::UnpackToSaasData(saasAnnoSite, annoSite, NIndexAnn::GetSaasProfile());

        UNIT_ASSERT_VALUES_EQUAL(ExpectedRtyAnnData, saasAnnoSite.ShortDebugString());
    }

    void TestSkipNonSaas() {
        NRTYServer::TMessage::TAnnData saasAnnoSite;
        NIndexAnn::TIndexAnnSiteData annoSite;
        NIndexAnn::TAnnotationRec* rec;
        NIndexAnn::TRegionData* reg;
        // rec1 is dropped entirely - no supported data
        rec = annoSite.AddRecs();
        rec->SetText("rec1");
        rec->SetTextLanguage(2);
        reg = rec->AddData();
        reg->SetRegion(0);
        reg->MutableBannerMinusWords(); //not supported
        reg->MutableClickMachine()->SetLongClick(0); // supported, but ommitted if zero

        UNIT_ASSERT_VALUES_EQUAL(true, NIndexAnn::UnpackToSaasData(saasAnnoSite, annoSite, NIndexAnn::TAnyFeaturesProfile()));
        UNIT_ASSERT_VALUES_EQUAL(false, NIndexAnn::UnpackToSaasData(saasAnnoSite, annoSite, NIndexAnn::GetSaasProfile()));

        // rec2 has only one region and one stream left
        rec = annoSite.AddRecs();
        rec->SetText("rec2");
        rec->SetTextLanguage(2);
        reg = rec->AddData();
        reg->SetRegion(200);
        reg->MutableBannerMinusWords(); //not supported
        reg = rec->AddData();
        reg->SetRegion(0);
        reg->MutableClickMachine()->SetLongClick(1); // supported

        bool ok = NIndexAnn::UnpackToSaasData(saasAnnoSite, annoSite, NIndexAnn::GetSaasProfile());
        if (!ok) {
            saasAnnoSite.Clear();
        }
        Cout << saasAnnoSite.ShortDebugString() << Endl;

        const TStringBuf expectedRtyAnnData =
                "Sentences { Text: \"rec2\" TextLanguage: 2 "
                "StreamsByRegion { Region: 0 Streams { Name: \"DT_LONG_CLICK\" Value { Value: \"/w==\" } } } }";

        UNIT_ASSERT_VALUES_EQUAL(expectedRtyAnnData, saasAnnoSite.ShortDebugString());
    }
};

UNIT_TEST_SUITE_REGISTRATION(TSaasRobotAnn);

