/*
 * SPDX-License-Identifier: Apache-2.0
 * Copyright 2018-2021 The Feast Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package feast.serving.it;

import static org.awaitility.Awaitility.await;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.jupiter.api.Assertions.*;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import feast.proto.core.FeatureProto;
import feast.proto.core.FeatureViewProto;
import feast.proto.core.RegistryProto;
import feast.proto.serving.ServingAPIProto;
import feast.proto.serving.ServingAPIProto.FieldStatus;
import feast.proto.types.ValueProto;
import feast.serving.registry.LocalRegistryFile;
import feast.serving.util.DataGenerator;
import io.grpc.StatusRuntimeException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.junit.jupiter.api.*;

abstract class ServingBaseTests extends ServingEnvironment {

  protected ServingAPIProto.GetOnlineFeaturesRequest buildOnlineRequest(int driverId) {
    // getOnlineFeatures Information
    String entityName = "driver_id";

    // Instantiate EntityRows
    Map<String, ValueProto.RepeatedValue> entityRows =
        ImmutableMap.of(
            entityName,
            ValueProto.RepeatedValue.newBuilder()
                .addVal(DataGenerator.createInt64Value(driverId))
                .build());

    ImmutableList<String> featureReferences =
        ImmutableList.of("driver_hourly_stats:conv_rate", "driver_hourly_stats:avg_daily_trips");

    // Build GetOnlineFeaturesRequestV2
    return TestUtils.createOnlineFeatureRequest(featureReferences, entityRows);
  }

  static RegistryProto.Registry registryProto = readLocalRegistry();

  private static RegistryProto.Registry readLocalRegistry() {
    return new LocalRegistryFile("src/test/resources/docker-compose/feast10/registry.db")
        .getContent();
  }

  @Test
  public void shouldGetOnlineFeatures() {
    ServingAPIProto.GetOnlineFeaturesResponse featureResponse =
        servingStub.getOnlineFeatures(buildOnlineRequest(1005));

    assertEquals(2, featureResponse.getResultsCount());
    assertEquals(1, featureResponse.getResults(0).getValuesCount());

    assertEquals(
        ImmutableList.of("driver_hourly_stats:conv_rate", "driver_hourly_stats:avg_daily_trips"),
        featureResponse.getMetadata().getFeatureNames().getValList());

    for (int featureIdx : List.of(0, 1)) {
      assertEquals(
          List.of(ServingAPIProto.FieldStatus.PRESENT),
          featureResponse.getResults(featureIdx).getStatusesList());
    }

    assertEquals(0.5, featureResponse.getResults(0).getValues(0).getDoubleVal(), 0.0001);
    assertEquals(500, featureResponse.getResults(1).getValues(0).getInt64Val());
  }

  @Test
  public void shouldGetOnlineFeaturesWithOutsideMaxAgeStatus() {
    ServingAPIProto.GetOnlineFeaturesResponse featureResponse =
        servingStub.getOnlineFeatures(buildOnlineRequest(1001));

    assertEquals(2, featureResponse.getResultsCount());
    assertEquals(1, featureResponse.getResults(0).getValuesCount());

    for (int featureIdx : List.of(0, 1)) {
      assertEquals(
          FieldStatus.OUTSIDE_MAX_AGE, featureResponse.getResults(featureIdx).getStatuses(0));
    }

    assertEquals(0.1, featureResponse.getResults(0).getValues(0).getDoubleVal(), 0.0001);
    assertEquals(100, featureResponse.getResults(1).getValues(0).getInt64Val());
  }

  @Test
  public void shouldGetOnlineFeaturesWithNotFoundStatus() {
    ServingAPIProto.GetOnlineFeaturesResponse featureResponse =
        servingStub.getOnlineFeatures(buildOnlineRequest(-1));

    assertEquals(2, featureResponse.getResultsCount());
    assertEquals(1, featureResponse.getResults(0).getValuesCount());

    for (final int featureIdx : List.of(0, 1)) {
      assertEquals(FieldStatus.NOT_FOUND, featureResponse.getResults(featureIdx).getStatuses(0));
    }
  }

  @Test
  public void shouldRefreshRegistryAndServeNewFeatures() throws InterruptedException {
    updateRegistryFile(
        registryProto
            .toBuilder()
            .addFeatureViews(
                FeatureViewProto.FeatureView.newBuilder()
                    .setSpec(
                        FeatureViewProto.FeatureViewSpec.newBuilder()
                            .setName("new_view")
                            .addEntities("driver_id")
                            .addFeatures(
                                FeatureProto.FeatureSpecV2.newBuilder()
                                    .setName("new_feature")
                                    .setValueType(ValueProto.ValueType.Enum.BOOL))))
            .build());

    ServingAPIProto.GetOnlineFeaturesRequest request = buildOnlineRequest(1005);

    ServingAPIProto.GetOnlineFeaturesRequest requestWithNewFeature =
        request
            .toBuilder()
            .setFeatures(request.getFeatures().toBuilder().addVal("new_view:new_feature"))
            .build();

    await()
        .ignoreException(StatusRuntimeException.class)
        .atMost(5, TimeUnit.SECONDS)
        .until(
            () -> servingStub.getOnlineFeatures(requestWithNewFeature).getResultsCount(),
            equalTo(3));
  }

  /** https://github.com/feast-dev/feast/issues/2253 */
  @Test
  public void shouldGetOnlineFeaturesWithStringEntity() {
    Map<String, ValueProto.RepeatedValue> entityRows =
        ImmutableMap.of(
            "entity",
            ValueProto.RepeatedValue.newBuilder()
                .addVal(DataGenerator.createStrValue("key-1"))
                .build());

    ImmutableList<String> featureReferences =
        ImmutableList.of("feature_view_0:feature_0", "feature_view_0:feature_1");

    ServingAPIProto.GetOnlineFeaturesRequest req =
        TestUtils.createOnlineFeatureRequest(featureReferences, entityRows);

    ServingAPIProto.GetOnlineFeaturesResponse resp = servingStub.getOnlineFeatures(req);

    for (final int featureIdx : List.of(0, 1)) {
      assertEquals(FieldStatus.PRESENT, resp.getResults(featureIdx).getStatuses(0));
    }
  }

  @Test
  public void shouldGetOnlineFeaturesFromAllFeatureViews() {
    Map<String, ValueProto.RepeatedValue> entityRows =
        ImmutableMap.of(
            "entity",
                ValueProto.RepeatedValue.newBuilder()
                    .addVal(DataGenerator.createStrValue("key-1"))
                    .build(),
            "driver_id",
                ValueProto.RepeatedValue.newBuilder()
                    .addVal(DataGenerator.createInt64Value(1005))
                    .build());

    ImmutableList<String> featureReferences =
        ImmutableList.of(
            "feature_view_0:feature_0",
            "feature_view_0:feature_1",
            "driver_hourly_stats:conv_rate",
            "driver_hourly_stats:avg_daily_trips");

    ServingAPIProto.GetOnlineFeaturesRequest req =
        TestUtils.createOnlineFeatureRequest(featureReferences, entityRows);

    ServingAPIProto.GetOnlineFeaturesResponse resp = servingStub.getOnlineFeatures(req);

    for (final int featureIdx : List.of(0, 1, 2, 3)) {
      assertEquals(FieldStatus.PRESENT, resp.getResults(featureIdx).getStatuses(0));
    }
  }

  abstract void updateRegistryFile(RegistryProto.Registry registry);
}
