/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 *
 *     http://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 org.apache.beam.sdk.expansion.service;

import static org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods.Enum.SCHEMA_TRANSFORM;
import static org.apache.beam.sdk.schemas.transforms.SchemaTransformTranslation.SchemaTransformPayloadTranslator;
import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
import static org.apache.beam.sdk.util.construction.PTransformTranslation.READ_TRANSFORM_URN;

import com.google.auto.service.AutoService;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.ServiceLoader;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.beam.fn.harness.ExternalWorkerService;
import org.apache.beam.model.expansion.v1.ExpansionApi;
import org.apache.beam.model.expansion.v1.ExpansionApi.DiscoverSchemaTransformRequest;
import org.apache.beam.model.expansion.v1.ExpansionApi.DiscoverSchemaTransformResponse;
import org.apache.beam.model.expansion.v1.ExpansionApi.SchemaTransformConfig;
import org.apache.beam.model.expansion.v1.ExpansionServiceGrpc;
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
import org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods;
import org.apache.beam.model.pipeline.v1.ExternalTransforms.ExternalConfigurationPayload;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.SchemaApi;
import org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.RowCoder;
import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
import org.apache.beam.sdk.expansion.service.JavaClassLookupTransformProvider.AllowList;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.schemas.NoSuchSchemaException;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.Field;
import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.schemas.SchemaRegistry;
import org.apache.beam.sdk.schemas.SchemaTranslation;
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.util.construction.BeamUrns;
import org.apache.beam.sdk.util.construction.Environments;
import org.apache.beam.sdk.util.construction.PTransformTranslation;
import org.apache.beam.sdk.util.construction.PTransformTranslation.TransformPayloadTranslator;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
import org.apache.beam.sdk.util.construction.PipelineTranslation;
import org.apache.beam.sdk.util.construction.RehydratedComponents;
import org.apache.beam.sdk.util.construction.SdkComponents;
import org.apache.beam.sdk.util.construction.SplittableParDo;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.alts.AltsServerBuilder;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.CaseFormat;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Converter;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A service that allows pipeline expand transforms from a remote SDK. */
@SuppressWarnings({
  "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
})
public class ExpansionService extends ExpansionServiceGrpc.ExpansionServiceImplBase
    implements AutoCloseable {

  private static final Logger LOG = LoggerFactory.getLogger(ExpansionService.class);

  private static final SchemaRegistry SCHEMA_REGISTRY = SchemaRegistry.createDefault();

  /**
   * A registrar that creates {@link TransformProvider} instances from {@link
   * RunnerApi.FunctionSpec}s.
   *
   * <p>Transform authors have the ability to provide a registrar by creating a {@link
   * ServiceLoader} entry and a concrete implementation of this interface.
   *
   * <p>It is optional but recommended to use one of the many build time tools such as {@link
   * AutoService} to generate the necessary META-INF files automatically.
   */
  public interface ExpansionServiceRegistrar {
    Map<String, TransformProvider> knownTransforms();
  }

  /**
   * Exposes Java transforms via {@link org.apache.beam.sdk.expansion.ExternalTransformRegistrar}.
   */
  @AutoService(ExpansionService.ExpansionServiceRegistrar.class)
  public static class ExternalTransformRegistrarLoader
      implements ExpansionService.ExpansionServiceRegistrar {

    @Override
    public Map<String, TransformProvider> knownTransforms() {
      Map<String, TransformProvider> providers = new HashMap<>();

      // First check and register ExternalTransformBuilder in ServiceLoader style, converting
      // to TransformProvider after validation.
      Map<String, ExternalTransformBuilder> registeredBuilders = loadTransformBuilders();
      for (Map.Entry<String, ExternalTransformBuilder> registeredBuilder :
          registeredBuilders.entrySet()) {
        providers.put(
            registeredBuilder.getKey(),
            new TransformProviderForBuilder(registeredBuilder.getValue()));
      }

      // Now load payload translators and convert them to TransformProviders after validating for
      // duplicates
      List<String> deprecatedTransformURNs = ImmutableList.of(READ_TRANSFORM_URN);
      Map<? extends Class<? extends PTransform>, ? extends TransformPayloadTranslator>
          registeredPayloadTranslators = PTransformTranslation.getKnownPayloadTranslators();

      for (Map.Entry<? extends Class<? extends PTransform>, ? extends TransformPayloadTranslator>
          entry : registeredPayloadTranslators.entrySet()) {

        @Initialized TransformPayloadTranslator translator = entry.getValue();

        if (translator == null) {
          continue;
        }

        String urn;
        try {
          urn = translator.getUrn();
          if (urn == null) {
            LOG.debug(
                "Could not load the TransformPayloadTranslator {} to the Expansion Service since it did not produce a unique URN.",
                translator);
            continue;
          } else if (urn.equals(BeamUrns.getUrn(SCHEMA_TRANSFORM))
              && translator instanceof SchemaTransformPayloadTranslator) {
            urn = ((SchemaTransformPayloadTranslator) translator).provider().identifier();
          }
        } catch (Exception e) {
          LOG.info(
              "Could not load the TransformPayloadTranslator {} to the Expansion Service.",
              translator,
              e);
          continue;
        }

        if (deprecatedTransformURNs.contains(urn)) {
          continue;
        }
        final String finalUrn = urn;
        TransformProvider transformProvider = new TransformProviderForPayloadTranslator(translator);
        providers.put(finalUrn, transformProvider);
      }
      return providers;
    }

    /**
     * Returns all {@link ExternalTransformBuilder} associations from any registrar set up via the
     * service loader interface.
     */
    private static Map<String, ExternalTransformBuilder> loadTransformBuilders() {

      Map<String, ExternalTransformBuilder> registeredBuilders = new HashMap<>();

      ImmutableSet.Builder<String> conflictingRegistrations = new ImmutableSet.Builder<>();

      for (ExternalTransformRegistrar registrar :
          ServiceLoader.load(ExternalTransformRegistrar.class)) {

        for (Map.Entry<String, ExternalTransformBuilder<?, ?, ?>> entry :
            registrar.knownBuilderInstances().entrySet()) {

          String urn = entry.getKey();
          ExternalTransformBuilder newBuilder = entry.getValue();
          @Nullable ExternalTransformBuilder existingBuilder = registeredBuilders.get(urn);

          if (existingBuilder == null) {
            registeredBuilders.put(urn, newBuilder);
          } else {
            LOG.error(
                "Conflicting registrations for {}: {} and {}", urn, existingBuilder, newBuilder);
            conflictingRegistrations.add(urn);
          }
        }
      }

      Set<String> conflictingRegistrationSet = conflictingRegistrations.build();
      if (!conflictingRegistrationSet.isEmpty()) {
        throw new IllegalArgumentException(
            String.format(
                "Conflicting registrations for: %s",
                Joiner.on(", ").join(conflictingRegistrationSet)));
      }
      return registeredBuilders;
    }
  }

  private static class TransformProviderForPayloadTranslator<
          InputT extends PInput, OutputT extends POutput>
      implements TransformProvider<InputT, OutputT> {

    private final TransformPayloadTranslator<PTransform<InputT, OutputT>> payloadTranslator;

    // Returns true if the underlying transform represented by this is a schema-aware transform.
    private boolean isSchemaTransform() {
      return (payloadTranslator instanceof SchemaTransformPayloadTranslator);
    }

    private TransformProviderForPayloadTranslator(
        TransformPayloadTranslator<PTransform<InputT, OutputT>> payloadTranslator) {
      this.payloadTranslator = payloadTranslator;
    }

    @Override
    public PTransform<InputT, OutputT> getTransform(
        RunnerApi.FunctionSpec spec, PipelineOptions options) {
      if (isSchemaTransform()) {
        return ExpansionServiceSchemaTransformProvider.of().getTransform(spec, options);
      } else {
        try {
          ExternalConfigurationPayload payload =
              ExternalConfigurationPayload.parseFrom(spec.getPayload());
          Row configRow =
              RowCoder.of(SchemaTranslation.schemaFromProto(payload.getSchema()))
                  .decode(new ByteArrayInputStream(payload.getPayload().toByteArray()));
          PTransform transformFromRow = payloadTranslator.fromConfigRow(configRow, options);
          if (transformFromRow != null) {
            return transformFromRow;
          } else {
            throw new RuntimeException(
                String.format(
                    "A transform cannot be initiated using the provided config row %s and the"
                        + " TransformPayloadTranslator %s",
                    configRow, payloadTranslator));
          }
        } catch (Exception e) {
          throw new RuntimeException(
              String.format(
                  "Failed to build transform %s from spec %s: %s",
                  spec.getUrn(), spec, e.getMessage()),
              e);
        }
      }
    }

    @Override
    public InputT createInput(Pipeline p, Map<String, PCollection<?>> inputs) {
      if (isSchemaTransform()) {
        return (InputT) ExpansionServiceSchemaTransformProvider.of().createInput(p, inputs);
      } else {
        return TransformProvider.super.createInput(p, inputs);
      }
    }

    @Override
    public Map<String, PCollection<?>> extractOutputs(OutputT output) {
      if (isSchemaTransform()) {
        return ExpansionServiceSchemaTransformProvider.of()
            .extractOutputs((PCollectionRowTuple) output);
      } else {
        return TransformProvider.super.extractOutputs(output);
      }
    }

    @Override
    public boolean equals(@Nullable Object other) {
      return other instanceof TransformProviderForPayloadTranslator
          && Objects.equals(
              payloadTranslator, ((TransformProviderForPayloadTranslator) other).payloadTranslator);
    }

    @Override
    public int hashCode() {
      return payloadTranslator.hashCode();
    }
  }

  private static class TransformProviderForBuilder implements TransformProvider {

    private final ExternalTransformBuilder transformBuilder;

    private TransformProviderForBuilder(ExternalTransformBuilder transformBuilder) {
      this.transformBuilder = transformBuilder;
    }

    @Override
    public PTransform getTransform(RunnerApi.FunctionSpec spec, PipelineOptions options) {
      try {
        Class configClass = getConfigClass(transformBuilder);
        return transformBuilder.buildExternal(
            payloadToConfig(
                ExternalConfigurationPayload.parseFrom(spec.getPayload()), configClass));
      } catch (Exception e) {
        throw new RuntimeException(
            String.format("Failed to build transform from spec %s: %s", spec, e.getMessage()), e);
      }
    }

    @Override
    public List<String> getDependencies(RunnerApi.FunctionSpec spec, PipelineOptions options) {
      try {
        Class configClass = getConfigClass(transformBuilder);
        Optional<List<String>> dependencies =
            transformBuilder.getDependencies(
                payloadToConfig(
                    ExternalConfigurationPayload.parseFrom(spec.getPayload()), configClass),
                options);
        return dependencies.orElseGet(() -> TransformProvider.super.getDependencies(spec, options));
      } catch (Exception e) {
        throw new RuntimeException(
            String.format("Failed to get dependencies for spec %s", spec), e);
      }
    }

    @Override
    public boolean equals(@Nullable Object other) {
      return other instanceof TransformProviderForBuilder
          && Objects.equals(
              transformBuilder, ((TransformProviderForBuilder) other).transformBuilder);
    }

    @Override
    public int hashCode() {
      return transformBuilder.hashCode();
    }
  }

  private static <ConfigT> Class<ConfigT> getConfigClass(
      ExternalTransformBuilder<ConfigT, ?, ?> transformBuilder) {
    Class<ConfigT> configurationClass = null;
    for (Method method : transformBuilder.getClass().getMethods()) {
      if (method.getName().equals("buildExternal")) {
        Preconditions.checkState(
            method.getParameterCount() == 1,
            "Build method for ExternalTransformBuilder %s must have exactly one parameter, but"
                + " had %s parameters.",
            transformBuilder.getClass().getSimpleName(),
            method.getParameterCount());
        configurationClass = (Class<ConfigT>) method.getParameterTypes()[0];
        if (Object.class.equals(configurationClass)) {
          continue;
        }
        break;
      }
    }

    if (configurationClass == null) {
      throw new AssertionError("Failed to find buildExternal method.");
    }

    return configurationClass;
  }

  static Row decodeConfigObjectRow(SchemaApi.Schema schema, ByteString payload) {
    Schema payloadSchema = SchemaTranslation.schemaFromProto(schema);

    if (payloadSchema.getFieldCount() == 0) {
      return Row.withSchema(Schema.of()).build();
    }

    // Coerce field names to camel-case
    Converter<String, String> camelCaseConverter =
        CaseFormat.LOWER_UNDERSCORE.converterTo(CaseFormat.LOWER_CAMEL);
    payloadSchema =
        payloadSchema.getFields().stream()
            .map(
                (field) -> {
                  Preconditions.checkNotNull(field.getName());
                  if (field.getName().contains("_")) {
                    @Nullable String newName = camelCaseConverter.convert(field.getName());
                    assert newName != null
                        : "@AssumeAssertion(nullness): converter type is imprecise; it is"
                            + " nullness-preserving";
                    return field.withName(newName);
                  } else {
                    return field;
                  }
                })
            .collect(Schema.toSchema());

    Row configRow;
    try {
      configRow = RowCoder.of(payloadSchema).decode(payload.newInput());
    } catch (IOException e) {
      throw new RuntimeException("Error decoding payload", e);
    }
    return configRow;
  }

  /**
   * Attempt to create an instance of {@link ConfigT} from an {@link ExternalConfigurationPayload}.
   * If a schema is registered for {@link ConfigT} this method will attempt to ise it. Throws an
   * {@link IllegalArgumentException} if the schema in {@code payload} is not {@link
   * Schema#assignableTo(Schema) assignable to} the registered schema.
   *
   * <p>If no Schema is registered, {@link ConfigT} must have a zero-argument constructor and
   * setters corresponding to each field in the row encoded by {@code payload}. Note {@link ConfigT}
   * may have additional setters not represented in the {@code payload} schema.
   *
   * <p>Exposed for testing only. No backwards compatibility guarantees.
   */
  @VisibleForTesting
  public static <ConfigT> ConfigT payloadToConfig(
      ExternalConfigurationPayload payload, Class<ConfigT> configurationClass) {
    try {
      return payloadToConfigSchema(payload, configurationClass);
    } catch (NoSuchSchemaException schemaException) {
      LOG.warn(
          "Configuration class '{}' has no schema registered. Attempting to construct with setter"
              + " approach.",
          configurationClass.getName());
      try {
        return payloadToConfigSetters(payload, configurationClass);
      } catch (ReflectiveOperationException e) {
        throw new IllegalArgumentException(
            String.format(
                "Failed to construct instance of configuration class '%s'",
                configurationClass.getName()),
            e);
      }
    }
  }

  private static <ConfigT> ConfigT payloadToConfigSchema(
      ExternalConfigurationPayload payload, Class<ConfigT> configurationClass)
      throws NoSuchSchemaException {
    Schema configSchema = SCHEMA_REGISTRY.getSchema(configurationClass);
    SerializableFunction<Row, ConfigT> fromRowFunc =
        SCHEMA_REGISTRY.getFromRowFunction(configurationClass);

    Row payloadRow = decodeConfigObjectRow(payload.getSchema(), payload.getPayload());

    if (!payloadRow.getSchema().assignableTo(configSchema)) {
      throw new IllegalArgumentException(
          String.format(
              "Schema in expansion request payload is not assignable to the schema for the "
                  + "configuration object.%n%nPayload Schema: %s%n%nConfiguration Schema: %s",
              payloadRow.getSchema(), configSchema));
    }

    return fromRowFunc.apply(payloadRow);
  }

  private static <ConfigT> ConfigT payloadToConfigSetters(
      ExternalConfigurationPayload payload, Class<ConfigT> configurationClass)
      throws ReflectiveOperationException {
    Row configRow = decodeConfigObjectRow(payload.getSchema(), payload.getPayload());

    Constructor<ConfigT> constructor = configurationClass.getDeclaredConstructor();
    constructor.setAccessible(true);

    ConfigT config = constructor.newInstance();
    for (Field field : configRow.getSchema().getFields()) {
      String key = field.getName();
      @Nullable Object value = configRow.getValue(field.getName());

      String fieldName = key;

      Coder coder = SchemaCoder.coderForFieldType(field.getType());
      Class type = coder.getEncodedTypeDescriptor().getRawType();

      String setterName =
          "set" + Character.toUpperCase(fieldName.charAt(0)) + fieldName.substring(1);
      Method method;
      try {
        // retrieve the setter for this field
        method = config.getClass().getMethod(setterName, type);
      } catch (NoSuchMethodException e) {
        throw new IllegalArgumentException(
            String.format(
                "The configuration class %s is missing a setter %s for %s with type %s",
                config.getClass(),
                setterName,
                fieldName,
                coder.getEncodedTypeDescriptor().getType().getTypeName()),
            e);
      }
      invokeSetter(config, value, method);
    }
    return config;
  }

  // Checker framework is conservative for Method#invoke, args are NonNull
  // See https://checkerframework.org/manual/#reflection-resolution
  @SuppressWarnings("nullness")
  private static <ConfigT> void invokeSetter(ConfigT config, @Nullable Object value, Method method)
      throws IllegalAccessException, InvocationTargetException {
    method.invoke(config, value);
  }

  private @MonotonicNonNull Map<String, TransformProvider> registeredTransforms;
  private final PipelineOptions commandLineOptions;
  private final @Nullable String loopbackAddress;

  public ExpansionService() {
    this(new String[] {});
  }

  public ExpansionService(String[] args) {
    this(PipelineOptionsFactory.fromArgs(args).create());
  }

  public ExpansionService(PipelineOptions opts) {
    this(opts, null);
  }

  public ExpansionService(PipelineOptions opts, @Nullable String loopbackAddress) {
    this.commandLineOptions = opts;
    this.loopbackAddress = loopbackAddress;
  }

  private Map<String, TransformProvider> getRegisteredTransforms() {
    if (registeredTransforms == null) {
      registeredTransforms = loadRegisteredTransforms();
    }
    return registeredTransforms;
  }

  private Iterable<SchemaTransformProvider> getRegisteredSchemaTransforms() {
    return ExpansionServiceSchemaTransformProvider.of().getAllProviders();
  }

  private Map<String, TransformProvider> loadRegisteredTransforms() {
    ImmutableMap.Builder<String, TransformProvider> registeredTransformsBuilder =
        ImmutableMap.builder();
    for (ExpansionServiceRegistrar registrar :
        ServiceLoader.load(ExpansionServiceRegistrar.class)) {
      registeredTransformsBuilder.putAll(registrar.knownTransforms());
    }
    ImmutableMap<String, TransformProvider> registeredTransforms =
        registeredTransformsBuilder.build();
    LOG.info("Registering external transforms: {}", registeredTransforms.keySet());
    return registeredTransforms;
  }

  @VisibleForTesting
  /*package*/ ExpansionApi.ExpansionResponse expand(ExpansionApi.ExpansionRequest request) {
    final String urn = request.getTransform().getSpec().getUrn();
    LOG.info("Expanding '{}' with URN '{}'", request.getTransform().getUniqueName(), urn);
    LOG.debug("Full transform: {}", request.getTransform());
    Set<String> existingTransformIds = request.getComponents().getTransformsMap().keySet();

    PipelineOptions pipelineOptionsFromRequest =
        PipelineOptionsTranslation.fromProto(request.getPipelineOptions());
    Pipeline pipeline = createPipeline(pipelineOptionsFromRequest);

    boolean isUseDeprecatedRead =
        ExperimentalOptions.hasExperiment(commandLineOptions, "use_deprecated_read")
            || ExperimentalOptions.hasExperiment(
                commandLineOptions, "beam_fn_api_use_deprecated_read");
    if (!isUseDeprecatedRead) {
      ExperimentalOptions.addExperiment(
          pipeline.getOptions().as(ExperimentalOptions.class), "beam_fn_api");
      // TODO(https://github.com/apache/beam/issues/20530): Remove this when we address performance
      // issue.
      ExperimentalOptions.addExperiment(
          pipeline.getOptions().as(ExperimentalOptions.class), "use_sdf_read");
    } else {
      LOG.warn(
          "Using use_depreacted_read in portable runners is runner-dependent. The "
              + "ExpansionService will respect that, but if your runner does not have support for "
              + "native Read transform, your Pipeline will fail during Pipeline submission.");
    }

    RehydratedComponents rehydratedComponents =
        RehydratedComponents.forComponents(request.getComponents()).withPipeline(pipeline);

    Map<String, PCollection<?>> inputs =
        request.getTransform().getInputsMap().entrySet().stream()
            .collect(
                Collectors.toMap(
                    Map.Entry::getKey,
                    input -> {
                      try {
                        return rehydratedComponents.getPCollection(input.getValue());
                      } catch (IOException exn) {
                        throw new RuntimeException(exn);
                      }
                    }));

    TransformProvider transformProvider = getRegisteredTransforms().get(urn);
    if (transformProvider == null) {
      if (getUrn(ExpansionMethods.Enum.JAVA_CLASS_LOOKUP).equals(urn)) {
        AllowList allowList =
            commandLineOptions.as(ExpansionServiceOptions.class).getJavaClassLookupAllowlist();
        assert allowList != null;
        transformProvider = new JavaClassLookupTransformProvider(allowList);
      } else if (getUrn(SCHEMA_TRANSFORM).equals(urn)) {
        try {
          String underlyingIdentifier =
              ExternalTransforms.SchemaTransformPayload.parseFrom(
                      request.getTransform().getSpec().getPayload())
                  .getIdentifier();
          transformProvider = getRegisteredTransforms().get(underlyingIdentifier);
        } catch (InvalidProtocolBufferException e) {
          throw new RuntimeException(e);
        }
        transformProvider =
            transformProvider != null
                ? transformProvider
                : ExpansionServiceSchemaTransformProvider.of();
      } else {
        throw new UnsupportedOperationException("Unknown urn: " + urn);
      }
    }

    // Use expansion config file provided in commandLineOptions if not available
    // in the expansion request options.
    String configFileFromPipelineOptions =
        pipeline.getOptions().as(ExpansionServiceOptions.class).getExpansionServiceConfigFile();
    String configFileFromCommandLineOptions =
        commandLineOptions.as(ExpansionServiceOptions.class).getExpansionServiceConfigFile();

    if (configFileFromPipelineOptions == null && configFileFromCommandLineOptions != null) {
      pipeline
          .getOptions()
          .as(ExpansionServiceOptions.class)
          .setExpansionServiceConfigFile(configFileFromCommandLineOptions);
    }

    List<String> classpathResources =
        transformProvider.getDependencies(request.getTransform().getSpec(), pipeline.getOptions());
    pipeline.getOptions().as(PortablePipelineOptions.class).setFilesToStage(classpathResources);

    Map<String, PCollection<?>> outputs =
        transformProvider.apply(
            pipeline,
            request.getTransform().getUniqueName(),
            request.getTransform().getSpec(),
            inputs);

    // Needed to find which transform was new...
    SdkComponents sdkComponents =
        rehydratedComponents
            .getSdkComponents(request.getRequirementsList())
            .withNewIdPrefix(request.getNamespace());
    RunnerApi.Environment defaultEnvironment =
        Environments.createOrGetDefaultEnvironment(
            pipeline.getOptions().as(PortablePipelineOptions.class));
    if (commandLineOptions.as(ExpansionServiceOptions.class).getAlsoStartLoopbackWorker()) {
      PortablePipelineOptions externalOptions =
          PipelineOptionsFactory.create().as(PortablePipelineOptions.class);
      externalOptions.setDefaultEnvironmentType(Environments.ENVIRONMENT_EXTERNAL);
      externalOptions.setDefaultEnvironmentConfig(loopbackAddress);
      defaultEnvironment =
          Environments.createAnyOfEnvironment(
              defaultEnvironment, Environments.createOrGetDefaultEnvironment(externalOptions));
    }
    sdkComponents.registerEnvironment(defaultEnvironment);
    Map<String, String> outputMap =
        outputs.entrySet().stream()
            .collect(
                Collectors.toMap(
                    Map.Entry::getKey,
                    output -> {
                      try {
                        return sdkComponents.registerPCollection(output.getValue());
                      } catch (IOException exn) {
                        throw new RuntimeException(exn);
                      }
                    }));

    if (isUseDeprecatedRead) {
      SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(pipeline);
    }

    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, sdkComponents);
    String expandedTransformId =
        Iterables.getOnlyElement(
            pipelineProto.getRootTransformIdsList().stream()
                .filter(id -> !existingTransformIds.contains(id))
                .collect(Collectors.toList()));
    RunnerApi.Components components = pipelineProto.getComponents();
    RunnerApi.PTransform expandedTransform =
        components
            .getTransformsOrThrow(expandedTransformId)
            .toBuilder()
            .setUniqueName(expandedTransformId)
            .clearOutputs()
            .putAllOutputs(outputMap)
            .build();
    LOG.debug("Expanded to {}", expandedTransform);

    return ExpansionApi.ExpansionResponse.newBuilder()
        .setComponents(components.toBuilder().removeTransforms(expandedTransformId))
        .setTransform(expandedTransform)
        .addAllRequirements(pipelineProto.getRequirementsList())
        .build();
  }

  protected Pipeline createPipeline(PipelineOptions requestOptions) {
    // We expect the ExpansionRequest to contain a valid set of options to be used for this
    // expansion.
    // Additionally, we override selected options using options values set via command line or
    // ExpansionService wide overrides.

    PortablePipelineOptions requestPortablePipelineOptions =
        requestOptions.as(PortablePipelineOptions.class);
    PortablePipelineOptions commandLinePortablePipelineOptions =
        commandLineOptions.as(PortablePipelineOptions.class);
    Optional.ofNullable(commandLinePortablePipelineOptions.getDefaultEnvironmentType())
        .ifPresent(requestPortablePipelineOptions::setDefaultEnvironmentType);
    Optional.ofNullable(commandLinePortablePipelineOptions.getDefaultEnvironmentConfig())
        .ifPresent(requestPortablePipelineOptions::setDefaultEnvironmentConfig);
    List<String> filesToStage = commandLinePortablePipelineOptions.getFilesToStage();
    if (filesToStage != null) {
      requestPortablePipelineOptions
          .as(PortablePipelineOptions.class)
          .setFilesToStage(filesToStage);
    }
    requestPortablePipelineOptions
        .as(ExperimentalOptions.class)
        .setExperiments(commandLineOptions.as(ExperimentalOptions.class).getExperiments());
    requestPortablePipelineOptions.setRunner(NotRunnableRunner.class);
    requestPortablePipelineOptions
        .as(ExpansionServiceOptions.class)
        .setExpansionServiceConfig(
            commandLineOptions.as(ExpansionServiceOptions.class).getExpansionServiceConfig());
    return Pipeline.create(requestOptions);
  }

  @Override
  public void expand(
      ExpansionApi.ExpansionRequest request,
      StreamObserver<ExpansionApi.ExpansionResponse> responseObserver) {
    try {
      responseObserver.onNext(expand(request));
      responseObserver.onCompleted();
    } catch (RuntimeException exn) {
      responseObserver.onNext(
          ExpansionApi.ExpansionResponse.newBuilder()
              .setError(Throwables.getStackTraceAsString(exn))
              .build());
      responseObserver.onCompleted();
    }
  }

  DiscoverSchemaTransformResponse discover(DiscoverSchemaTransformRequest request) {
    ExpansionServiceSchemaTransformProvider transformProvider =
        ExpansionServiceSchemaTransformProvider.of();
    DiscoverSchemaTransformResponse.Builder responseBuilder =
        DiscoverSchemaTransformResponse.newBuilder();
    for (org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider provider :
        transformProvider.getAllProviders()) {
      SchemaTransformConfig.Builder schemaTransformConfigBuilder =
          SchemaTransformConfig.newBuilder();
      schemaTransformConfigBuilder.setDescription(provider.description());
      schemaTransformConfigBuilder.setConfigSchema(
          SchemaTranslation.schemaToProto(provider.configurationSchema(), true));
      schemaTransformConfigBuilder.addAllInputPcollectionNames(provider.inputCollectionNames());
      schemaTransformConfigBuilder.addAllOutputPcollectionNames(provider.outputCollectionNames());
      responseBuilder.putSchemaTransformConfigs(
          provider.identifier(), schemaTransformConfigBuilder.build());
    }

    return responseBuilder.build();
  }

  @Override
  public void discoverSchemaTransform(
      DiscoverSchemaTransformRequest request,
      StreamObserver<DiscoverSchemaTransformResponse> responseObserver) {
    try {
      responseObserver.onNext(discover(request));
      responseObserver.onCompleted();
    } catch (RuntimeException exn) {
      responseObserver.onNext(
          ExpansionApi.DiscoverSchemaTransformResponse.newBuilder()
              .setError(Throwables.getStackTraceAsString(exn))
              .build());
      responseObserver.onCompleted();
    }
  }

  @Override
  public void close() throws Exception {
    // Nothing to do because the expansion service is stateless.
  }

  public static void main(String[] args) throws Exception {
    if (args.length < 1) {
      printUsage();
      System.exit(1);
    }
    int port = Integer.parseInt(args[0]);
    System.out.println("Starting expansion service at localhost:" + port);

    // Register the options class used by the expansion service.
    PipelineOptionsFactory.register(ExpansionServiceOptions.class);
    @SuppressWarnings({"nullness"})
    PipelineOptions options =
        PipelineOptionsFactory.fromArgs(Arrays.copyOfRange(args, 1, args.length)).create();

    @SuppressWarnings("nullness")
    ExpansionService service = new ExpansionService(options, "localhost:" + port);

    StringBuilder registeredTransformsLog = new StringBuilder();
    boolean registeredTransformsFound = false;
    registeredTransformsLog.append("\n");
    registeredTransformsLog.append("Registered transforms:");

    for (Map.Entry<String, TransformProvider> entry :
        service.getRegisteredTransforms().entrySet()) {
      registeredTransformsFound = true;
      registeredTransformsLog.append("\n\t" + entry.getKey() + ": " + entry.getValue());
    }

    StringBuilder registeredSchemaTransformProvidersLog = new StringBuilder();
    boolean registeredSchemaTransformProvidersFound = false;
    registeredSchemaTransformProvidersLog.append("\n");
    registeredSchemaTransformProvidersLog.append("Registered SchemaTransformProviders:");

    for (SchemaTransformProvider provider : service.getRegisteredSchemaTransforms()) {
      registeredSchemaTransformProvidersFound = true;
      registeredSchemaTransformProvidersLog.append("\n\t" + provider.identifier());
    }

    if (registeredTransformsFound) {
      System.out.println(registeredTransformsLog.toString());
    }
    if (registeredSchemaTransformProvidersFound) {
      System.out.println(registeredSchemaTransformProvidersLog.toString());
    }
    if (!registeredTransformsFound && !registeredSchemaTransformProvidersFound) {
      System.out.println("\nDid not find any registered transforms or SchemaTransforms.\n");
    }

    boolean useAlts = options.as(ExpansionServiceOptions.class).getUseAltsServer();
    ServerBuilder serverBuilder =
        useAlts ? AltsServerBuilder.forPort(port) : ServerBuilder.forPort(port);

    if (useAlts) {
      LOG.info("Running with gRPC ALTS authentication.");
    }

    serverBuilder.addService(service).addService(new ArtifactRetrievalService());
    if (options.as(ExpansionServiceOptions.class).getAlsoStartLoopbackWorker()) {
      serverBuilder.addService(new ExternalWorkerService(options));
    }
    Server server = serverBuilder.build();
    server.start();
    server.awaitTermination();
  }

  private static void printUsage() {
    System.err.println("Usage: java -jar <expansion-service-jar> port [pipeline-options]");
  }

  private static class NotRunnableRunner extends PipelineRunner<PipelineResult> {
    public static NotRunnableRunner fromOptions(@SuppressWarnings("unused") PipelineOptions opts) {
      return new NotRunnableRunner();
    }

    @Override
    public PipelineResult run(Pipeline pipeline) {
      throw new UnsupportedOperationException();
    }
  }
}
