/*
 * 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.runners.core.construction.BeamUrns.getUrn;
import static org.junit.Assert.assertEquals;

import com.google.auto.service.AutoService;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.model.expansion.v1.ExpansionApi;
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
import org.apache.beam.model.pipeline.v1.ExternalTransforms.ExpansionMethods;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.ParDoTranslation;
import org.apache.beam.runners.core.construction.PipelineTranslation;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.schemas.JavaFieldSchema;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.Field;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.schemas.SchemaTranslation;
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.annotations.SchemaCreate;
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Impulse;
import org.apache.beam.sdk.transforms.InferableFunction;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.util.ByteStringOutputStream;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.grpc.v1p54p0.com.google.protobuf.InvalidProtocolBufferException;
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.Iterables;
import org.junit.Test;

/** Tests for {@link ExpansionServiceSchemaTransformProvider}. */
@SuppressWarnings({
  "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
})
public class ExpansionServiceSchemaTransformProviderTest {

  private static final String TEST_NAME = "TestName";

  private static final String TEST_NAMESPACE = "namespace";

  private static final Schema TEST_SCHEMATRANSFORM_CONFIG_SCHEMA =
      Schema.of(
          Field.of("int1", FieldType.INT32),
          Field.of("int2", FieldType.INT32),
          Field.of("str1", FieldType.STRING),
          Field.of("str2", FieldType.STRING));

  private ExpansionService expansionService = new ExpansionService();

  @DefaultSchema(JavaFieldSchema.class)
  public static class TestSchemaTransformConfiguration {

    public final String str1;
    public final String str2;
    public final Integer int1;
    public final Integer int2;

    @SchemaCreate
    public TestSchemaTransformConfiguration(String str1, String str2, Integer int1, Integer int2) {
      this.str1 = str1;
      this.str2 = str2;
      this.int1 = int1;
      this.int2 = int2;
    }
  }

  /** Registers a SchemaTransform. */
  @AutoService(SchemaTransformProvider.class)
  public static class TestSchemaTransformProvider
      extends TypedSchemaTransformProvider<TestSchemaTransformConfiguration> {

    @Override
    protected Class<TestSchemaTransformConfiguration> configurationClass() {
      return TestSchemaTransformConfiguration.class;
    }

    @Override
    protected SchemaTransform from(TestSchemaTransformConfiguration configuration) {
      return new TestSchemaTransform(
          configuration.str1, configuration.str2, configuration.int1, configuration.int2);
    }

    @Override
    public String identifier() {
      return "dummy_id";
    }

    @Override
    public List<String> inputCollectionNames() {
      return ImmutableList.of("input1");
    }

    @Override
    public List<String> outputCollectionNames() {
      return ImmutableList.of("output1");
    }
  }

  public static class TestDoFn extends DoFn<String, String> {

    public String str1;
    public String str2;
    public int int1;
    public int int2;

    public TestDoFn(String str1, String str2, Integer int1, Integer int2) {
      this.str1 = str1;
      this.str2 = str2;
      this.int1 = int1;
      this.int2 = int2;
    }

    @ProcessElement
    public void processElement(@Element String element, OutputReceiver<String> receiver) {
      receiver.output(element);
    }
  }

  public static class TestSchemaTransform extends SchemaTransform {

    private String str1;
    private String str2;
    private Integer int1;
    private Integer int2;

    public TestSchemaTransform(String str1, String str2, Integer int1, Integer int2) {
      this.str1 = str1;
      this.str2 = str2;
      this.int1 = int1;
      this.int2 = int2;
    }

    @Override
    public PCollectionRowTuple expand(PCollectionRowTuple input) {
      PCollection<Row> outputPC =
          input
              .getAll()
              .values()
              .iterator()
              .next()
              .apply(
                  MapElements.via(
                      new InferableFunction<Row, String>() {
                        @Override
                        public String apply(Row input) throws Exception {
                          return input.getString("in_str");
                        }
                      }))
              .apply(ParDo.of(new TestDoFn(this.str1, this.str2, this.int1, this.int2)))
              .apply(
                  MapElements.via(
                      new InferableFunction<String, Row>() {
                        @Override
                        public Row apply(String input) throws Exception {
                          return Row.withSchema(Schema.of(Field.of("out_str", FieldType.STRING)))
                              .withFieldValue("out_str", input)
                              .build();
                        }
                      }))
              .setRowSchema(Schema.of(Field.of("out_str", FieldType.STRING)));
      return PCollectionRowTuple.of("output1", outputPC);
    }
  }

  /** Registers a SchemaTransform. */
  @AutoService(SchemaTransformProvider.class)
  public static class TestSchemaTransformProviderMultiInputMultiOutput
      extends TypedSchemaTransformProvider<TestSchemaTransformConfiguration> {

    @Override
    protected Class<TestSchemaTransformConfiguration> configurationClass() {
      return TestSchemaTransformConfiguration.class;
    }

    @Override
    protected SchemaTransform from(TestSchemaTransformConfiguration configuration) {
      return new TestSchemaTransformMultiInputOutput(
          configuration.str1, configuration.str2, configuration.int1, configuration.int2);
    }

    @Override
    public String identifier() {
      return "dummy_id_multi_input_multi_output";
    }

    @Override
    public List<String> inputCollectionNames() {
      return ImmutableList.of("input1", "input2");
    }

    @Override
    public List<String> outputCollectionNames() {
      return ImmutableList.of("output1", "output2");
    }
  }

  public static class TestSchemaTransformMultiInputOutput extends SchemaTransform {

    private String str1;
    private String str2;
    private Integer int1;
    private Integer int2;

    public TestSchemaTransformMultiInputOutput(
        String str1, String str2, Integer int1, Integer int2) {
      this.str1 = str1;
      this.str2 = str2;
      this.int1 = int1;
      this.int2 = int2;
    }

    @Override
    public PCollectionRowTuple expand(PCollectionRowTuple input) {
      PCollection<Row> outputPC1 =
          input
              .get("input1")
              .apply(
                  MapElements.via(
                      new InferableFunction<Row, String>() {
                        @Override
                        public String apply(Row input) throws Exception {
                          return input.getString("in_str");
                        }
                      }))
              .apply(ParDo.of(new TestDoFn(this.str1, this.str2, this.int1, this.int2)))
              .apply(
                  MapElements.via(
                      new InferableFunction<String, Row>() {
                        @Override
                        public Row apply(String input) throws Exception {
                          return Row.withSchema(Schema.of(Field.of("out_str", FieldType.STRING)))
                              .withFieldValue("out_str", input)
                              .build();
                        }
                      }))
              .setRowSchema(Schema.of(Field.of("out_str", FieldType.STRING)));
      PCollection<Row> outputPC2 =
          input
              .get("input2")
              .apply(
                  MapElements.via(
                      new InferableFunction<Row, String>() {
                        @Override
                        public String apply(Row input) throws Exception {
                          return input.getString("in_str");
                        }
                      }))
              .apply(ParDo.of(new TestDoFn(this.str1, this.str2, this.int1, this.int2)))
              .apply(
                  MapElements.via(
                      new InferableFunction<String, Row>() {
                        @Override
                        public Row apply(String input) throws Exception {
                          return Row.withSchema(Schema.of(Field.of("out_str", FieldType.STRING)))
                              .withFieldValue("out_str", input)
                              .build();
                        }
                      }))
              .setRowSchema(Schema.of(Field.of("out_str", FieldType.STRING)));
      return PCollectionRowTuple.of("output1", outputPC1, "output2", outputPC2);
    }
  }

  @Test
  public void testSchemaTransformDiscovery() {
    ExpansionApi.DiscoverSchemaTransformRequest discoverRequest =
        ExpansionApi.DiscoverSchemaTransformRequest.newBuilder().build();
    ExpansionApi.DiscoverSchemaTransformResponse response =
        expansionService.discover(discoverRequest);
    assertEquals(2, response.getSchemaTransformConfigsCount());
  }

  private void verifyLeafTransforms(ExpansionApi.ExpansionResponse response, int count) {

    int leafTransformCount = 0;
    for (RunnerApi.PTransform transform : response.getComponents().getTransformsMap().values()) {
      if (transform.getSpec().getUrn().equals(PTransformTranslation.PAR_DO_TRANSFORM_URN)) {
        RunnerApi.ParDoPayload parDoPayload;
        try {
          parDoPayload = RunnerApi.ParDoPayload.parseFrom(transform.getSpec().getPayload());
          DoFn doFn = ParDoTranslation.getDoFn(parDoPayload);
          if (!(doFn instanceof TestDoFn)) {
            continue;
          }
          TestDoFn testDoFn = (TestDoFn) doFn;
          assertEquals("aaa", testDoFn.str1);
          assertEquals("bbb", testDoFn.str2);
          assertEquals(111, testDoFn.int1);
          assertEquals(222, testDoFn.int2);
          leafTransformCount++;
        } catch (InvalidProtocolBufferException exc) {
          throw new RuntimeException(exc);
        }
      }
    }
    assertEquals(count, leafTransformCount);
  }

  @Test
  public void testSchemaTransformExpansion() {
    Pipeline p = Pipeline.create();
    p.apply(Impulse.create());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);

    String inputPcollId =
        Iterables.getOnlyElement(
            Iterables.getOnlyElement(pipelineProto.getComponents().getTransformsMap().values())
                .getOutputsMap()
                .values());
    Row configRow =
        Row.withSchema(TEST_SCHEMATRANSFORM_CONFIG_SCHEMA)
            .withFieldValue("int1", 111)
            .withFieldValue("int2", 222)
            .withFieldValue("str1", "aaa")
            .withFieldValue("str2", "bbb")
            .build();

    ByteStringOutputStream outputStream = new ByteStringOutputStream();
    try {
      SchemaCoder.of(configRow.getSchema()).encode(configRow, outputStream);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    ExternalTransforms.SchemaTransformPayload payload =
        ExternalTransforms.SchemaTransformPayload.newBuilder()
            .setIdentifier("dummy_id")
            .setConfigurationRow(outputStream.toByteString())
            .setConfigurationSchema(
                SchemaTranslation.schemaToProto(TEST_SCHEMATRANSFORM_CONFIG_SCHEMA, true))
            .build();

    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(
                        RunnerApi.FunctionSpec.newBuilder()
                            .setUrn(getUrn(ExpansionMethods.Enum.SCHEMA_TRANSFORM))
                            .setPayload(payload.toByteString()))
                    .putInputs("input1", inputPcollId))
            .setNamespace(TEST_NAMESPACE)
            .build();
    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();

    assertEquals(3, expandedTransform.getSubtransformsCount());
    assertEquals(1, expandedTransform.getInputsCount());
    assertEquals(1, expandedTransform.getOutputsCount());
    verifyLeafTransforms(response, 1);
  }

  @Test
  public void testSchemaTransformExpansionMultiInputMultiOutput() {
    Pipeline p = Pipeline.create();
    p.apply(Impulse.create());
    p.apply(Impulse.create());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);

    List<String> inputPcollIds = new ArrayList<>();
    for (RunnerApi.PTransform transform :
        pipelineProto.getComponents().getTransformsMap().values()) {
      inputPcollIds.add(Iterables.getOnlyElement(transform.getOutputsMap().values()));
    }
    assertEquals(2, inputPcollIds.size());

    Row configRow =
        Row.withSchema(TEST_SCHEMATRANSFORM_CONFIG_SCHEMA)
            .withFieldValue("int1", 111)
            .withFieldValue("int2", 222)
            .withFieldValue("str1", "aaa")
            .withFieldValue("str2", "bbb")
            .build();

    ByteStringOutputStream outputStream = new ByteStringOutputStream();
    try {
      SchemaCoder.of(configRow.getSchema()).encode(configRow, outputStream);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    ExternalTransforms.SchemaTransformPayload payload =
        ExternalTransforms.SchemaTransformPayload.newBuilder()
            .setIdentifier("dummy_id_multi_input_multi_output")
            .setConfigurationRow(outputStream.toByteString())
            .setConfigurationSchema(
                SchemaTranslation.schemaToProto(TEST_SCHEMATRANSFORM_CONFIG_SCHEMA, true))
            .build();

    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(
                        RunnerApi.FunctionSpec.newBuilder()
                            .setUrn(getUrn(ExpansionMethods.Enum.SCHEMA_TRANSFORM))
                            .setPayload(payload.toByteString()))
                    .putInputs("input1", inputPcollIds.get(0))
                    .putInputs("input2", inputPcollIds.get(1)))
            .setNamespace(TEST_NAMESPACE)
            .build();
    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();

    assertEquals(6, expandedTransform.getSubtransformsCount());
    assertEquals(2, expandedTransform.getInputsCount());
    assertEquals(2, expandedTransform.getOutputsCount());
    verifyLeafTransforms(response, 2);
  }
}
