/*
 * 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.service;

import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import com.google.protobuf.Timestamp;
import feast.proto.core.*;
import feast.proto.serving.ServingAPIProto;
import feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesRequest;
import feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesResponse;
import feast.proto.serving.TransformationServiceAPIProto.ValueType;
import feast.proto.serving.TransformationServiceGrpc;
import feast.proto.types.ValueProto;
import feast.serving.registry.RegistryRepository;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Status;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.util.*;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.*;
import org.apache.arrow.vector.ipc.ArrowFileReader;
import org.apache.arrow.vector.ipc.ArrowFileWriter;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;

public class OnlineTransformationService implements TransformationService {

  private static final Logger log =
      org.slf4j.LoggerFactory.getLogger(OnlineTransformationService.class);
  private final TransformationServiceGrpc.TransformationServiceBlockingStub stub;
  private final RegistryRepository registryRepository;
  static final int INT64_BITWIDTH = 64;
  static final int INT32_BITWIDTH = 32;

  public OnlineTransformationService(
      String transformationServiceEndpoint, RegistryRepository registryRepository) {
    if (transformationServiceEndpoint != null) {
      final ManagedChannel channel =
          ManagedChannelBuilder.forTarget(transformationServiceEndpoint).usePlaintext().build();
      this.stub = TransformationServiceGrpc.newBlockingStub(channel);
    } else {
      this.stub = null;
    }
    this.registryRepository = registryRepository;
  }

  /** {@inheritDoc} */
  @Override
  public TransformFeaturesResponse transformFeatures(
      TransformFeaturesRequest transformFeaturesRequest) {
    if (this.stub == null) {
      throw new RuntimeException(
          "Transformation service endpoint must be configured to enable this functionality.");
    }
    return this.stub.transformFeatures(transformFeaturesRequest);
  }

  /** {@inheritDoc} */
  @Override
  public List<ServingAPIProto.FeatureReferenceV2> extractOnDemandFeaturesDependencies(
      List<ServingAPIProto.FeatureReferenceV2> onDemandFeatureReferences) {
    List<ServingAPIProto.FeatureReferenceV2> onDemandFeatureSources = new ArrayList<>();
    for (ServingAPIProto.FeatureReferenceV2 featureReference : onDemandFeatureReferences) {
      OnDemandFeatureViewProto.OnDemandFeatureViewSpec onDemandFeatureViewSpec =
          this.registryRepository.getOnDemandFeatureViewSpec(featureReference);
      Map<String, OnDemandFeatureViewProto.OnDemandSource> sources =
          onDemandFeatureViewSpec.getSourcesMap();

      for (OnDemandFeatureViewProto.OnDemandSource source : sources.values()) {
        OnDemandFeatureViewProto.OnDemandSource.SourceCase sourceCase = source.getSourceCase();
        switch (sourceCase) {
          case REQUEST_DATA_SOURCE:
            // Do nothing. The value should be provided as dedicated request parameter
            break;
          case FEATURE_VIEW_PROJECTION:
            FeatureReferenceProto.FeatureViewProjection projection =
                source.getFeatureViewProjection();
            for (FeatureProto.FeatureSpecV2 featureSpec : projection.getFeatureColumnsList()) {
              String featureName = featureSpec.getName();
              ServingAPIProto.FeatureReferenceV2 onDemandFeatureSource =
                  ServingAPIProto.FeatureReferenceV2.newBuilder()
                      .setFeatureViewName(projection.getFeatureViewName())
                      .setFeatureName(featureName)
                      .build();
              onDemandFeatureSources.add(onDemandFeatureSource);
            }
            break;
          case FEATURE_VIEW:
            FeatureViewProto.FeatureView featureView = source.getFeatureView();
            FeatureViewProto.FeatureViewSpec featureViewSpec = featureView.getSpec();
            String featureViewName = featureViewSpec.getName();
            for (FeatureProto.FeatureSpecV2 featureSpec : featureViewSpec.getFeaturesList()) {
              String featureName = featureSpec.getName();
              ServingAPIProto.FeatureReferenceV2 onDemandFeatureSource =
                  ServingAPIProto.FeatureReferenceV2.newBuilder()
                      .setFeatureViewName(featureViewName)
                      .setFeatureName(featureName)
                      .build();
              onDemandFeatureSources.add(onDemandFeatureSource);
            }
            break;
          default:
            throw Status.INTERNAL
                .withDescription(
                    "OnDemandSource proto source field has an unexpected type: " + sourceCase)
                .asRuntimeException();
        }
      }
    }
    return onDemandFeatureSources;
  }

  /** {@inheritDoc} */
  public void processTransformFeaturesResponse(
      feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesResponse
          transformFeaturesResponse,
      String onDemandFeatureViewName,
      Set<String> onDemandFeatureStringReferences,
      ServingAPIProto.GetOnlineFeaturesResponse.Builder responseBuilder) {
    try {
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      ArrowFileReader reader =
          new ArrowFileReader(
              new ByteArrayReadableSeekableByteChannel(
                  transformFeaturesResponse
                      .getTransformationOutput()
                      .getArrowValue()
                      .toByteArray()),
              allocator);
      reader.loadNextBatch();
      VectorSchemaRoot readBatch = reader.getVectorSchemaRoot();
      Schema responseSchema = readBatch.getSchema();
      List<Field> responseFields = responseSchema.getFields();
      Timestamp now = Timestamp.newBuilder().setSeconds(System.currentTimeMillis() / 1000).build();

      for (Field field : responseFields) {
        String columnName = field.getName();
        String fullFeatureName = columnName.replace("__", ":");
        ArrowType columnType = field.getType();

        // The response will contain all features for the specified ODFV, so we
        // skip the features that were not requested.
        if (!onDemandFeatureStringReferences.contains(fullFeatureName)) {
          continue;
        }

        FieldVector fieldVector = readBatch.getVector(field);
        int valueCount = fieldVector.getValueCount();
        ServingAPIProto.GetOnlineFeaturesResponse.FeatureVector.Builder vectorBuilder =
            responseBuilder.addResultsBuilder();
        List<ValueProto.Value> valueList = Lists.newArrayListWithExpectedSize(valueCount);

        // TODO: support all Feast types
        // TODO: clean up the switch statement
        if (columnType instanceof ArrowType.Int) {
          int bitWidth = ((ArrowType.Int) columnType).getBitWidth();
          switch (bitWidth) {
            case INT64_BITWIDTH:
              for (int i = 0; i < valueCount; i++) {
                long int64Value = ((BigIntVector) fieldVector).get(i);
                valueList.add(ValueProto.Value.newBuilder().setInt64Val(int64Value).build());
              }
              break;
            case INT32_BITWIDTH:
              for (int i = 0; i < valueCount; i++) {
                int int32Value = ((IntVector) fieldVector).get(i);
                valueList.add(ValueProto.Value.newBuilder().setInt32Val(int32Value).build());
              }
              break;
            default:
              throw Status.INTERNAL
                  .withDescription(
                      "Column "
                          + columnName
                          + " is of type ArrowType.Int but has bitWidth "
                          + bitWidth
                          + " which cannot be handled.")
                  .asRuntimeException();
          }
        } else if (columnType instanceof ArrowType.FloatingPoint) {
          FloatingPointPrecision precision = ((ArrowType.FloatingPoint) columnType).getPrecision();
          switch (precision) {
            case DOUBLE:
              for (int i = 0; i < valueCount; i++) {
                double doubleValue = ((Float8Vector) fieldVector).get(i);
                valueList.add(ValueProto.Value.newBuilder().setDoubleVal(doubleValue).build());
              }
              break;
            case SINGLE:
              for (int i = 0; i < valueCount; i++) {
                float floatValue = ((Float4Vector) fieldVector).get(i);
                valueList.add(ValueProto.Value.newBuilder().setFloatVal(floatValue).build());
              }
              break;
            default:
              throw Status.INTERNAL
                  .withDescription(
                      "Column "
                          + columnName
                          + " is of type ArrowType.FloatingPoint but has precision "
                          + precision
                          + " which cannot be handled.")
                  .asRuntimeException();
          }
        }

        for (ValueProto.Value v : valueList) {
          vectorBuilder.addValues(v);
          vectorBuilder.addStatuses(ServingAPIProto.FieldStatus.PRESENT);
          vectorBuilder.addEventTimestamps(now);
        }

        responseBuilder.getMetadataBuilder().getFeatureNamesBuilder().addVal(fullFeatureName);
      }
    } catch (IOException e) {
      log.info(e.toString());
      throw Status.INTERNAL
          .withDescription("Unable to correctly process transform features response: " + e)
          .asRuntimeException();
    }
  }

  /** {@inheritDoc} */
  public ValueType serializeValuesIntoArrowIPC(List<Pair<String, List<ValueProto.Value>>> values) {
    // In order to be serialized correctly, the data must be packaged in a VectorSchemaRoot.
    // We first construct all the columns.
    BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);

    List<Field> columnFields = new ArrayList<>();
    List<FieldVector> columns = new ArrayList<>();

    for (Pair<String, List<ValueProto.Value>> columnEntry : values) {
      // The Python FTS does not expect full feature names, so we extract the feature name.
      String columnName = columnEntry.getKey();

      List<ValueProto.Value> columnValues = columnEntry.getValue();
      FieldVector column;
      ValueProto.Value.ValCase valCase = columnValues.get(0).getValCase();
      // TODO: support all Feast types
      switch (valCase) {
        case INT32_VAL:
          column = new IntVector(columnName, allocator);
          column.setValueCount(columnValues.size());
          for (int idx = 0; idx < columnValues.size(); idx++) {
            ((IntVector) column).set(idx, columnValues.get(idx).getInt32Val());
          }
          break;
        case INT64_VAL:
          column = new BigIntVector(columnName, allocator);
          column.setValueCount(columnValues.size());
          for (int idx = 0; idx < columnValues.size(); idx++) {
            ((BigIntVector) column).set(idx, columnValues.get(idx).getInt64Val());
          }

          break;
        case DOUBLE_VAL:
          column = new Float8Vector(columnName, allocator);
          column.setValueCount(columnValues.size());
          for (int idx = 0; idx < columnValues.size(); idx++) {
            ((Float8Vector) column).set(idx, columnValues.get(idx).getDoubleVal());
          }
          break;
        case FLOAT_VAL:
          column = new Float4Vector(columnName, allocator);
          column.setValueCount(columnValues.size());
          for (int idx = 0; idx < columnValues.size(); idx++) {
            ((Float4Vector) column).set(idx, columnValues.get(idx).getFloatVal());
          }
          break;
        default:
          throw Status.INTERNAL
              .withDescription(
                  "Column " + columnName + " has a type that is currently not handled: " + valCase)
              .asRuntimeException();
      }

      columns.add(column);
      columnFields.add(column.getField());
    }

    VectorSchemaRoot schemaRoot = new VectorSchemaRoot(columnFields, columns);

    // Serialize the VectorSchemaRoot into Arrow IPC format.
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    ArrowFileWriter writer = new ArrowFileWriter(schemaRoot, null, Channels.newChannel(out));
    try {
      writer.start();
      writer.writeBatch();
      writer.end();
    } catch (IOException e) {
      log.info(e.toString());
      throw Status.INTERNAL
          .withDescription("ArrowFileWriter could not write properly; failed with error: " + e)
          .asRuntimeException();
    }
    byte[] byteData = out.toByteArray();
    ByteString sourceData = ByteString.copyFrom(byteData);
    ValueType transformationInput = ValueType.newBuilder().setArrowValue(sourceData).build();
    return transformationInput;
  }
}
