#include "json.h"

#include <saas/util/json/json.h>

#include <library/cpp/testing/unittest/registar.h>
#include <library/cpp/json/json_writer.h>

#include <util/stream/str.h>

using namespace NSearchMapParser;

//#define JSON_FORMAT_UT_DEBUG

namespace {
    const TString CompleteSearchMap = "{\"tests\" : { \"dispatch_target\" : \"distributor\", \"require_auth\" : true,"
        "\"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default0\" : "
        " [{ \"shard_min\" : 0, \"shard_max\" : 65533, \"host\" : \"localhost\", \"search_port\" : 15020, \"indexer_port\" : 19020}]}}}";

    const TString CompleteSearchMapWithShardBy = "{\"tests\" : { \"dispatch_target\" : \"distributor\", \"require_auth\" : true,"
        "\"shard_by\" : \"url_hash\", \"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default0\" : "
        " [{ \"shard_min\" : 0, \"shard_max\" : 65533, \"host\" : \"localhost\", \"search_port\" : 15020, \"indexer_port\" : 19020}]}}}";

    const TString TwoServices = "{ \"fotki\" : { \"search_port\" : 15020, \"indexer_port\" : 19010, \"replicas\" : "
        "{ \"ugr\" : [{ \"host\" : \"arachnid07.yandex.ru\"}]}}, "
        "\"help\" : { \"search_port\" : 17270, \"indexer_port\" : 19270, \"replicas\" : "
        "{ \"myt\" : [{ \"host\" : \"rty03f-prestable.saas.yandex.net\"}]}}}";

    const TString TwoServicesWithSkipPing = "{ \"fotki\" : { \"search_port\" : 15020, \"indexer_port\" : 19010, \"skip_ping\" : true,\"replicas\" : "
        "{ \"ugr\" : [{ \"host\" : \"arachnid07.yandex.ru\"}]}}, "
        "\"help\" : { \"search_port\" : 17270, \"indexer_port\" : 19270, \"replicas\" : "
        "{ \"myt\" : [{ \"host\" : \"rty03f-prestable.saas.yandex.net\"}]}}}";

    const TString OverridenPorts = "{\"tests\" : { \"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default0\" :"
        "[{ \"host\" : \"localhost\", \"search_port\" : 15021, \"indexer_port\" : 12222}]}}}";

    const TString DefaultPortShift = "{\"tests\" : { \"replicas\" : { \"default0\" :"
        "[{ \"host\" : \"localhost\", \"search_port\" : 15021}]}}}";

    const TString InverseOrderPorts = "{\"tests\" : { \"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default0\" :"
        "[{\"host\": \"ws20-199\", \"shard_min\": 16384 }, {\"host\": \"ws39-001\", \"shard_max\": 16383 }]}}}";

    const TString NonUniformSearchMap = "{\"tests\": {\"search_port\": 7300, \"indexer_port\": 7302, \"replicas\": "
        " {\"default0\": [{\"host\": \"ws20-199\", \"shard_min\": 0, \"shard_max\": 16383 }, {\"host\": \"ws39-001\", \"shard_min\": 16384, \"shard_max\": 32767 },"
                         "{\"host\": \"ws39-003\", \"shard_min\": 32768, \"shard_max\": 49151 }, {\"host\": \"ws39-004\", \"shard_min\": 49152, \"shard_max\": 65533 } ], "
        "  \"default1\": [{\"host\": \"ws25-329\", \"shard_min\": 0, \"shard_max\": 77 }, {\"host\": \"ws25-330\", \"shard_min\": 78, \"shard_max\": 65533 } ] } } }";

    const TString Jibberish = "{\"WHAT?\" : \"dunno\"}";

    const TString DisabledSearch = "{\"tests\" : { \"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default0\" :{"
        "\"hosts\" : [{ \"host\" : \"localhost\", \"search_port\" : 15020, \"indexer_port\" : 19020 }],"
        "\"disable_search\" : true, \"disable_indexing\" : true }, \"default1\" :{"
        "\"hosts\" : [{ \"host\" : \"localhost\", \"search_port\" : 15021, \"indexer_port\" : 19021 }]"
        "}   }}}";

    const TString ServiceDiscovery = "{\"tests\" : { \"search_port\" : 15020, \"indexer_port\" : 19020, \"replicas\" : { \"default\" :{"
        "\"hosts\" : [{ \"host\" : \"sas@localhost-yp-1\", \"search_port\" : 15020, \"indexer_port\" : 19020, \"is_sd\" : true},"
        "{ \"host\" : \"localhost\", \"search_port\" : 15021, \"indexer_port\" : 19021 }]"
        "}   }}}";

    const TString BrokenSearchMap = "{ \"fotki\" : { \"search_port\" : 15020, \"indexer_port\" : 19010, \"replicas\" : "
        "{ \"ugr\" : [{ \"host\" : \"arachnid07.yandex.ru\"}]}}, "
        "\"help\" : { \"search_port\" : 19010, \"indexer_port\" : 19270, \"replicas\" : "
        "{ \"myt\" : [{ \"host\" : \"arachnid07.yandex.ru\"}]}}}";

    const TString MetaService = "{ \"services\": {\"ppb-fresh\":{\"replicas\": {\"default0\": [{ \"shard_min\": 0, \"shard_max\": 65533, \"host\": \"localhost\", \"search_port\": 15020, \"indexer_port\": 19020}]}}, \"ppb-archive\": {\"replicas\":{\"fuckhead\":[{\"search_port\":7300,\"indexer_port\":7302,\"host\":\"sas1-4261\"}]}}}, \"meta\": {   \"ppb\": { \"components\": [{\"service\": \"ppb-fresh\", \"priority\": 0},{\"service\": \"ppb-archive\", \"priority\": 1}]}}}";

#define SEARCHMAP_FROM_STRING(X) \
    TStringInput ss(X); \
    TJsonSearchMapParser parser; \
    parser.ParseSearchMap((IInputStream&)ss);\
    const TSearchMap& searchMap = parser.GetSearchMap();\
    const TSearchMap::TServiceMap& serviceMap = searchMap.GetServiceMap();\
    Y_UNUSED(serviceMap);
}

Y_UNIT_TEST_SUITE(TestSearchMapJsonFormat) {
    void CheckCompleteParameterSet(const TString& stringSearchMap, NSaas::ShardingType sharding = NSaas::KeyPrefix) {
        SEARCHMAP_FROM_STRING(stringSearchMap);
        UNIT_ASSERT_C(serviceMap.size() == 1, "incorrect services count");
        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
#define CHECK_FIELD(name, value) UNIT_ASSERT_C(serviceIterator->second.name == value, "incorrect " #name)
        CHECK_FIELD(SearchProtocol, "proto");
        CHECK_FIELD(ShardsDispatcher->GetContext().Type, sharding);
        CHECK_FIELD(IndexingTarget, Distributor);
        CHECK_FIELD(RequireAuth, true);
        CHECK_FIELD(SkipPing, false);
#undef  CHECK_FIELD
        const TSearchCluster& cluster = serviceIterator->second;
        TSearchReplica::const_iterator i = cluster.Replicas[0].begin();
        UNIT_ASSERT(i != cluster.Replicas[0].end());
        const TVector<TSearchInformation>& searchInfo = i->second;
        UNIT_ASSERT(searchInfo.size() == 1);
        UNIT_ASSERT_C(searchInfo[0].IndexerPort == 19020, "wrong indexer port");
        UNIT_ASSERT_C(searchInfo[0].SearchPort == 15020, "wrong search port");
    }

    Y_UNIT_TEST(TestCompleteParameterSet) {
        CheckCompleteParameterSet(CompleteSearchMap);
        CheckCompleteParameterSet(CompleteSearchMapWithShardBy, NSaas::UrlHash);
    }

    Y_UNIT_TEST(TestTwoServices) {
        {
            SEARCHMAP_FROM_STRING(TwoServices);
            UNIT_ASSERT_C(serviceMap.size() == 2, "incorrect services count");
            TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("fotki");
            UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service fotki not found");
            const TSearchCluster& cluster = serviceIterator->second;
            UNIT_ASSERT_C(!cluster.SkipPing, "Incorrect skip ping");
        }
        {
            SEARCHMAP_FROM_STRING(TwoServicesWithSkipPing);
            UNIT_ASSERT_C(serviceMap.size() == 2, "incorrect services count");
            TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("fotki");
            UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service fotki not found");
            const TSearchCluster& cluster = serviceIterator->second;
            UNIT_ASSERT_C(cluster.SkipPing, "Incorrect skip ping");
        }

    }

    Y_UNIT_TEST(TestFilter) {
        TParsingSettings settings;
        for (const TString singleService: {"help", "fotki"}) {
            settings.FilterServices = {singleService};
            TJsonSearchMapParser parser(settings);
            parser.ParseSearchMap(TwoServices);
            TSearchMap searchMap = parser.GetSearchMap();
            const auto& serviceMap = searchMap.GetServiceMap();

            UNIT_ASSERT(serviceMap.size() == 1);
            UNIT_ASSERT(serviceMap.find(singleService) != serviceMap.end());
        }
    }

    Y_UNIT_TEST(TestInverseOrder) {
        SEARCHMAP_FROM_STRING(TwoServices);
    }

    Y_UNIT_TEST(TestDefaultPortShift) {
        SEARCHMAP_FROM_STRING(DefaultPortShift);
        UNIT_ASSERT_C(serviceMap.size() == 1, "incorrect services count");
        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = serviceIterator->second;
        TSearchReplica::const_iterator i = cluster.Replicas[0].begin();
        UNIT_ASSERT(i != cluster.Replicas[0].end());
        const TVector<TSearchInformation>& searchInfo = i->second;
        UNIT_ASSERT(searchInfo.size() == 1);
        UNIT_ASSERT_C(searchInfo[0].SearchPort == 15021, "wrong search port");
        UNIT_ASSERT_C(searchInfo[0].IndexerPort == 15023, "wrong indexer port");
    }

    Y_UNIT_TEST(TestOverridenPorts) {
        SEARCHMAP_FROM_STRING(OverridenPorts);
        UNIT_ASSERT_C(serviceMap.size() == 1, "incorrect services count");
        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = serviceIterator->second;
        TSearchReplica::const_iterator i = cluster.Replicas[0].begin();
        UNIT_ASSERT(i != cluster.Replicas[0].end());
        const TVector<TSearchInformation>& searchInfo = i->second;
        UNIT_ASSERT(searchInfo.size() == 1);
        UNIT_ASSERT_C(searchInfo[0].IndexerPort == 12222, "wrong indexer port");
        UNIT_ASSERT_C(searchInfo[0].SearchPort == 15021, "wrong search port");
    }

    Y_UNIT_TEST(TestSerialization) {
        TString serializedSearchMap;
        {
            SEARCHMAP_FROM_STRING(CompleteSearchMap);
            NJson::TJsonValue serialized = parser.SerializeToJson();
            TStringOutput stringOutput(serializedSearchMap);
            NJson::WriteJson(&stringOutput, &serialized);
        }
        CheckCompleteParameterSet(serializedSearchMap);

        TString serializedSearchMap1;
        {
            SEARCHMAP_FROM_STRING(CompleteSearchMapWithShardBy);
            NJson::TJsonValue serialized = parser.SerializeToJson();
            TStringOutput stringOutput(serializedSearchMap1);
            NJson::WriteJson(&stringOutput, &serialized);
        }
        CheckCompleteParameterSet(serializedSearchMap1, NSaas::UrlHash);

        TString serializedSearchMap2;
        {
            SEARCHMAP_FROM_STRING(MetaService);
            NJson::TJsonValue serialized = parser.SerializeToJson();
            TStringOutput stringOutput(serializedSearchMap2);
            NJson::WriteJson(&stringOutput, &serialized);
        }
        {
            SEARCHMAP_FROM_STRING(serializedSearchMap2);
            auto& meta = searchMap.GetMetaServices();
            UNIT_ASSERT(meta.size() == 1);
            UNIT_ASSERT(meta[0].Name == "ppb");
            UNIT_ASSERT(meta[0].Components[1].ServiceName == "ppb-archive");
        }
    }

    Y_UNIT_TEST(TestNonUniformSearchMap) {
        SEARCHMAP_FROM_STRING(NonUniformSearchMap);

        TSearchMap::TServiceMap::const_iterator service = serviceMap.find("tests");
        UNIT_ASSERT_C(service != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = service->second;
        Y_UNUSED(cluster);
        /* disabled after compiled searchmap refactoring
        const TSearchReplica& replica = cluster.Replicas[0];
        UNIT_ASSERT(replica.GetSearchInterval("ws25-329:7300").size() == 1);
        UNIT_ASSERT(replica.GetSearchInterval("ws25-330:7300").size() == 4);
        UNIT_ASSERT(replica.GetSearchInterval("ws20-199:7300").size() == 2);
        UNIT_ASSERT(replica.DoesCoverInterval(TInterval<ui16>(0, SearchMapShards - 1)));

        for (TSearchReplica::const_iterator interval = replica.begin(); interval != replica.end(); ++interval)
            UNIT_ASSERT(interval->second.size() == 2);
        */
    }

    Y_UNIT_TEST(TestIgnoreErrors) {
        TStringInput ss(Jibberish);
        TParsingSettings settings;
        settings.IgnoreParsingErrors = true;
        TJsonSearchMapParser parser(settings);
        parser.ParseSearchMap((IInputStream&)ss);
        UNIT_ASSERT(parser.GetSearchMap(false, false).GetInternalSearchMap().ysize() == 1);
    }

    Y_UNIT_TEST(TestDisabledSearch) {
        NJson::TJsonValue original = NUtil::JsonFromString(DisabledSearch);
        SEARCHMAP_FROM_STRING(DisabledSearch);
        NJson::TJsonValue serialized = parser.SerializeToJson();
#if defined(JSON_FORMAT_UT_DEBUG)
        Cerr << "original: " << NUtil::JsonToString(original) << Endl;
        Cerr << "serialized: " << NUtil::JsonToString(serialized) << Endl;
#endif
        //UNIT_ASSERT(serialized == original);

        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = serviceIterator->second;
        const TSearchReplica& replica = cluster.Replicas[0];
        TSearchReplica::const_iterator i = replica.begin();
        UNIT_ASSERT(i != replica.end());
        //UNIT_ASSERT(replica.CustomConfig == "custom");
        //UNIT_ASSERT(replica.DisableSearch == true);
        auto hosts = replica.begin()->second;
        UNIT_ASSERT(hosts.size() == 2);
        UNIT_ASSERT(hosts[1].DisableSearch == true);
    }

    Y_UNIT_TEST(TestServiceDiscovery) {
        NJson::TJsonValue original = NUtil::JsonFromString(ServiceDiscovery);
        SEARCHMAP_FROM_STRING(ServiceDiscovery);
        NJson::TJsonValue serialized = parser.SerializeToJson();
#if defined(JSON_FORMAT_UT_DEBUG)
        Cerr << "original: " << NUtil::JsonToString(original) << Endl;
        Cerr << "serialized: " << NUtil::JsonToString(serialized) << Endl;
#endif
        //UNIT_ASSERT(serialized == original);

        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = serviceIterator->second;
        const TSearchReplica& replica = cluster.Replicas[0];
        TSearchReplica::const_iterator i = replica.begin();
        UNIT_ASSERT(i != replica.end());
        //UNIT_ASSERT(replica.CustomConfig == "custom");
        //UNIT_ASSERT(replica.DisableSearch == true);
        auto hosts = replica.begin()->second;
        UNIT_ASSERT(hosts.size() == 2);
        UNIT_ASSERT(hosts[0].IsSd == true);
        UNIT_ASSERT(hosts[1].IsSd == false);
    }

    Y_UNIT_TEST(TestAllHosts) {
        SEARCHMAP_FROM_STRING(NonUniformSearchMap);

        TSearchMap::TServiceMap::const_iterator serviceIterator = serviceMap.find("tests");
        UNIT_ASSERT_C(serviceIterator != serviceMap.end(), "service tests not found");
        const TSearchCluster& cluster = serviceIterator->second;
        UNIT_ASSERT(cluster.GetAllHosts().size() == 6);
    }

    Y_UNIT_TEST(TestSearcMapDuplicatedHosts) {
        TString error;
        try {
            SEARCHMAP_FROM_STRING(TwoServices);
        } catch (const yexception&) {
            UNIT_ASSERT_C(false, "Unexpected exception " + CurrentExceptionMessage());
        }

        bool failed = false;
        try {
            SEARCHMAP_FROM_STRING(BrokenSearchMap);
        } catch (const yexception&) {
            failed = true;
        }
        UNIT_ASSERT_C(failed, "Should fail");
    }

    Y_UNIT_TEST(TestRescaleSearchMap) {
        SEARCHMAP_FROM_STRING(NonUniformSearchMap);
        parser.RescaleSearchMap(SearchMapShards);
        TString error;
        UNIT_ASSERT_C(parser.VerifySearchMap(error), error.data());
    }

    Y_UNIT_TEST(TestMergeReplicas) {
        SEARCHMAP_FROM_STRING(NonUniformSearchMap);
        TSearchMap sm = parser.GetSearchMap(false, false);
        UNIT_ASSERT(sm.GetInternalSearchMap().size() == 1);
        UNIT_ASSERT(sm.GetInternalSearchMap()[0].Replicas.size() == 1);
        UNIT_ASSERT(sm.GetInternalSearchMap()[0].Replicas[0].Hosts.size() == 6);
    }

    Y_UNIT_TEST(TestSlotsPool) {
        TSlotsPool pool;
        const NJson::TJsonValue& source =
            NUtil::JsonFromString("[{\"search_port\":10139,\"indexer_port\":10141,\"shard_max\":32765,\"host\":\"localhost\"},"
                                       "{\"shard_min\":32766,\"search_port\":10143,\"indexer_port\":10145,\"host\":\"localhost\"}]");
        UNIT_ASSERT(pool.DeserializeFromJson(source));
    }

    Y_UNIT_TEST(TestMetaService) {
        SEARCHMAP_FROM_STRING(MetaService);
        auto& meta = searchMap.GetMetaServices();
        UNIT_ASSERT(meta.size() == 1);
        UNIT_ASSERT(meta[0].Name == "ppb");
        UNIT_ASSERT(meta[0].Components[1].ServiceName == "ppb-archive");
    }
};
