#include "searchproxyserver.h"

#include <saas/library/searchmap/parsers/parser.h>

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

namespace {

const TString SearchMap_2_2 = R"json({
  "services":
    {
      "test_service":
        {
          "shard_by":"url_hash",
          "replicas":
            {
              "default":
                [
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":65533,
                    "host":"vla1-1074.search.yandex.net"
                  },
                  {
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":32765,
                    "host":"vla1-0876.search.yandex.net"
                  },
                  {
                    "shard_min":32766,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":65533,
                    "host":"man1-1599.search.yandex.net"
                  },
                  {
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":32765,
                    "host":"man1-0436.search.yandex.net"
                  }
                ]
            }
        }
    },
  "properties":
    {
      "version":2
    },
  "meta":
    {
    }
})json";

const TString SearchMap_2_3 = R"json({
  "services":
    {
      "test_service":
        {
          "shard_by":"url_hash",
          "replicas":
            {
              "default":
                [
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":65533,
                    "host":"vla1-1074.search.yandex.net"
                  },
                  {
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":32765,
                    "host":"vla1-0876.search.yandex.net"
                  },
                  {
                    "shard_min":43688,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":65533,
                    "host":"man1-1599.search.yandex.net"
                  },
                  {
                    "shard_min":21844,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":43687,
                    "host":"man1-1234.search.yandex.net"
                  },
                  {
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":21843,
                    "host":"man1-0436.search.yandex.net"
                  }
                ]
            }
        }
    },
  "properties":
    {
      "version":2
    },
  "meta":
    {
    }
})json";

const TString SearchMap_2_4 = R"json({
  "services":
    {
      "test_service":
        {
          "shard_by":"url_hash",
          "replicas":
            {
              "default":
                [
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":65533,
                    "host":"vla1-1074.search.yandex.net"
                  },
                  {
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":32765,
                    "host":"vla1-0876.search.yandex.net"
                  },
                  {
                    "shard_min":49149,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":65533,
                    "host":"man1-1599.search.yandex.net"
                  },
                  {
                    "shard_min":32766,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":49148,
                    "host":"man1-2345.search.yandex.net"
                  },
                  {
                    "shard_min":16383,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":32765,
                    "host":"man1-1234.search.yandex.net"
                  },
                  {
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":16382,
                    "host":"man1-0436.search.yandex.net"
                  }
                ]
            }
        }
    },
  "properties":
    {
      "version":2
    },
  "meta":
    {
    }
})json";

const TString SearchMap_2_5 = R"json({
  "services":
    {
      "test_service":
        {
          "shard_by":"url_hash",
          "replicas":
            {
              "default":
                [
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":65533,
                    "host":"vla1-1074.search.yandex.net"
                  },
                  {
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":32765,
                    "host":"vla1-0876.search.yandex.net"
                  },
                  {
                    "shard_min":52424,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":65533,
                    "host":"man1-1599.search.yandex.net"
                  },
                  {
                    "shard_min":13106,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":26211,
                    "host":"man1-1234.search.yandex.net"
                  },
                  {
                    "shard_min":26212,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":39317,
                    "host":"man1-2345.search.yandex.net"
                  },
                  {
                    "shard_min":39318,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":52423,
                    "host":"man1-3456.search.yandex.net"
                  },
                  {
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":13105,
                    "host":"man1-0436.search.yandex.net"
                  }
                ]
            }
        }
    },
  "properties":
    {
      "version":2
    },
  "meta":
    {
    }
})json";

const TString SearchMap_SD_2_2 = R"json({
  "services":
    {
      "test_service":
        {
          "shard_by":"url_hash",
          "replicas":
            {
              "default":
                [
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":65533,
                    "is_sd": true,
                    "host":"vla@endpointset2-yp"
                  },
                  {
                    "shard_min":32766,
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"SAS-0@1@1@1",
                    "shard_max":65533,
                    "is_sd": true,
                    "host":"sas@endpointset2-yp"
                  },
                  {
                    "search_port":27820,
                    "indexer_port":27822,
                    "group":"VLA-0@1@1@1",
                    "shard_max":32765,
                    "host":"vla1-0876.search.yandex.net"
                  },
                  {
                    "shard_min":32766,
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":65533,
                    "host":"man1-1599.search.yandex.net"
                  },
                  {
                    "search_port":33840,
                    "indexer_port":33842,
                    "group":"MAN-1@0@0@1",
                    "shard_max":32765,
                    "host":"man1-0436.search.yandex.net"
                  }
                ]
            }
        }
    },
  "properties":
    {
      "version":2
    },
  "meta":
    {
    }
})json";

struct TFakeHostResolver : IHostResolver {
    bool CanBeResolved(const TString& /*host*/, ui16 /*port*/) noexcept final {
        return true;
    }
};

void DoTest(const TString& searchMapJson, const TVector<std::pair<TString, TString>>& expectedSources) {
    auto searchMap = NSearchMapParser::OpenSearchMap(searchMapJson)->GetSearchMap(false, true);
    TFakeHostResolver hostResolver;
    TSourcesConfig sourcesConfig = Default<TSearchServerConfig>();
    TSearchSourcesFiller searchSourcesFiller(searchMap, hostResolver);
    searchSourcesFiller.FillSearchSources("test_service", sourcesConfig, {}, TServiceConfig::TConsistencyPolicy::cpNo, true, true);

    ASSERT_EQ(sourcesConfig.SearchSources.size(), expectedSources.size() + 1);
    for (size_t i = 0; i < expectedSources.size(); ++i) {
        ASSERT_EQ(sourcesConfig.SearchSources[i + 1].ProtoConfig_.GetServerGroup(), expectedSources[i].first);
        ASSERT_EQ(sourcesConfig.SearchSources[i + 1].ProtoConfig_.GetSearchScript(), expectedSources[i].second);
    }
}

TEST(TSearchSourcesFillerSuite, Replicas_2_2) {
    DoTest(SearchMap_2_2, {
        {"0-32765", "(tcp2://man1-0436.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"32766-65533", "(tcp2://man1-1599.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "}
    });
}

TEST(TSearchSourcesFillerSuite, Replicas_2_3) {
    DoTest(SearchMap_2_3, {
        {"0-21843", "(tcp2://man1-0436.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"21844-32765", "(tcp2://man1-1234.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"32766-43687", "(tcp2://man1-1234.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "},
        {"43688-65533", "(tcp2://man1-1599.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "}
    });
}

TEST(TSearchSourcesFillerSuite, Replicas_2_4) {
    DoTest(SearchMap_2_4, {
        {"0-16382", "(tcp2://man1-0436.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"16383-32765", "(tcp2://man1-1234.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"32766-49148", "(tcp2://man1-2345.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "},
        {"49149-65533", "(tcp2://man1-1599.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "}
    });
}

TEST(TSearchSourcesFillerSuite, Replicas_2_5) {
    DoTest(SearchMap_2_5, {
        {"0-13105", "(tcp2://man1-0436.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"13106-26211", "(tcp2://man1-1234.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"26212-32765", "(tcp2://man1-2345.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"32766-39317", "(tcp2://man1-2345.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "},
        {"39318-52423", "(tcp2://man1-3456.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "},
        {"52424-65533", "(tcp2://man1-1599.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-1074.search.yandex.net:27821 )@0@1@1@1 "}
    });
}
TEST(TSearchSourcesFillerSuite, Replicas_SD_2_2) {
    DoTest(SearchMap_SD_2_2, {
        {"0-32765", "(tcp2://man1-0436.search.yandex.net:33841 )@1@0@0@1 (tcp2://vla1-0876.search.yandex.net:27821 )@0@1@1@1 "},
        {"32766-65533", "sd://vla@endpointset2-yp@0@1@1@1 sd://sas@endpointset2-yp@0@1@1@1 "}
    });
}
}
