/*
 * 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.sdk.util.construction.PipelineOptionsTranslation.PIPELINE_OPTIONS_URN_PREFIX;
import static org.apache.beam.sdk.util.construction.PipelineOptionsTranslation.PIPELINE_OPTIONS_URN_SUFFIX;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.hasProperty;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.notNullValue;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import com.google.auto.service.AutoService;
import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.beam.model.expansion.v1.ExpansionApi;
import org.apache.beam.model.pipeline.v1.ExternalTransforms;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.schemas.AutoValueSchema;
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.transforms.Count;
import org.apache.beam.sdk.transforms.Impulse;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.util.ByteStringOutputStream;
import org.apache.beam.sdk.util.construction.PipelineTranslation;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Value;
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.Iterables;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Resources;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.hamcrest.Matchers;
import org.junit.Assert;
import org.junit.Test;

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

  private static final String TEST_URN = "test:beam:transforms:count";
  private static final String TEST_OPTIONS_URN = "test:beam:transforms:test_options";

  private static final String TEST_NAME = "TestName";

  private static final String TEST_NAMESPACE = "namespace";

  private ExpansionService expansionService = new ExpansionService();
  public static final List<byte[]> BYTE_LIST =
      ImmutableList.of("testing", "compound", "coders").stream()
          .map(str -> str.getBytes(StandardCharsets.UTF_8))
          .collect(Collectors.toList());
  public static final Map<String, Long> BYTE_KV_LIST =
      ImmutableList.of("testing", "compound", "coders").stream()
          .collect(Collectors.toMap((str) -> str, (str) -> (long) str.length()));
  public static final Map<String, List<Long>> BYTE_KV_LIST_WITH_LIST_VALUE =
      ImmutableList.of("testing", "compound", "coders").stream()
          .collect(
              Collectors.toMap(str -> str, str -> Collections.singletonList((long) str.length())));

  /** Registers a single test transformation. */
  @AutoService(ExpansionService.ExpansionServiceRegistrar.class)
  public static class TestTransformRegistrar implements ExpansionService.ExpansionServiceRegistrar {

    static final String EXPECTED_STRING_VALUE = "abcde";
    static final Boolean EXPECTED_BOOLEAN_VALUE = true;
    static final Integer EXPECTED_INTEGER_VALUE = 12345;

    @Override
    public Map<String, TransformProvider> knownTransforms() {
      return ImmutableMap.of(
          TEST_URN, (spec, options) -> Count.perElement(),
          TEST_OPTIONS_URN,
              (spec, options) ->
                  new TestOptionsTransform(
                      EXPECTED_STRING_VALUE, EXPECTED_BOOLEAN_VALUE, EXPECTED_INTEGER_VALUE));
    }
  }

  public interface TestOptions extends PipelineOptions {
    String getStringOption();

    void setStringOption(String value);

    Boolean getBooleanOption();

    void setBooleanOption(Boolean value);

    Integer getIntegerOption();

    void setIntegerOption(Integer value);
  }

  public static class TestOptionsTransform
      extends PTransform<PCollection<String>, PCollection<String>> {
    String expectedStringValue;

    Boolean expectedBooleanValue;

    Integer expectedIntegerValue;

    public TestOptionsTransform(
        String expectedStringValue, Boolean expectedBooleanValue, Integer expectedIntegerValue) {
      this.expectedStringValue = expectedStringValue;
      this.expectedBooleanValue = expectedBooleanValue;
      this.expectedIntegerValue = expectedIntegerValue;
    }

    @Override
    public PCollection<String> expand(PCollection<String> input) {
      TestOptions testOption = input.getPipeline().getOptions().as(TestOptions.class);

      Assert.assertEquals(expectedStringValue, testOption.getStringOption());
      Assert.assertEquals(expectedBooleanValue, testOption.getBooleanOption());
      Assert.assertEquals(expectedIntegerValue, testOption.getIntegerOption());

      return input;
    }
  }

  @Test
  public void testConstruct() {
    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());
    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(TEST_URN))
                    .putInputs("input", inputPcollId))
            .setNamespace(TEST_NAMESPACE)
            .build();
    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();
    assertEquals(TEST_NAMESPACE + TEST_NAME, expandedTransform.getUniqueName());

    // Verify it has the right input.
    assertThat(expandedTransform.getInputsMap().values(), contains(inputPcollId));

    // Verify it has the right output.
    assertThat(expandedTransform.getOutputsMap().keySet(), contains("output"));

    // Loose check that it's composite, and its children are represented.
    assertThat(expandedTransform.getSubtransformsCount(), greaterThan(0));
    for (String subtransform : expandedTransform.getSubtransformsList()) {
      assertTrue(response.getComponents().containsTransforms(subtransform));
    }
    // Check that any newly generated components are properly namespaced.
    Set<String> originalIds = allIds(request.getComponents());
    for (String id : allIds(response.getComponents())) {
      assertTrue(id, id.startsWith(TEST_NAMESPACE) || originalIds.contains(id));
    }
  }

  @Test
  public void testConstructWithPipelineOptions() {
    PipelineOptionsFactory.register(TestOptions.class);
    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());

    Struct optionsStruct =
        Struct.newBuilder()
            .putFields(
                PIPELINE_OPTIONS_URN_PREFIX + "string_option" + PIPELINE_OPTIONS_URN_SUFFIX,
                Value.newBuilder()
                    .setStringValue(TestTransformRegistrar.EXPECTED_STRING_VALUE)
                    .build())
            .putFields(
                PIPELINE_OPTIONS_URN_PREFIX + "boolean_option" + PIPELINE_OPTIONS_URN_SUFFIX,
                Value.newBuilder()
                    .setBoolValue(TestTransformRegistrar.EXPECTED_BOOLEAN_VALUE)
                    .build())
            .putFields(
                PIPELINE_OPTIONS_URN_PREFIX + "integer_option" + PIPELINE_OPTIONS_URN_SUFFIX,
                Value.newBuilder()
                    .setNumberValue(TestTransformRegistrar.EXPECTED_INTEGER_VALUE)
                    .build())
            .build();
    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setPipelineOptions(optionsStruct)
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(TEST_OPTIONS_URN))
                    .putInputs("input", inputPcollId))
            .setNamespace(TEST_NAMESPACE)
            .build();
    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();
    assertEquals(TEST_NAMESPACE + TEST_NAME, expandedTransform.getUniqueName());

    // Verify it has the right input.
    assertThat(expandedTransform.getInputsMap().values(), contains(inputPcollId));

    // Verify it has the right output.
    assertThat(expandedTransform.getOutputsMap().keySet(), contains("output"));
  }

  @Test
  public void testConstructGenerateSequenceWithRegistration() {
    ExternalTransforms.ExternalConfigurationPayload payload =
        encodeRowIntoExternalConfigurationPayload(
            Row.withSchema(
                    Schema.of(
                        Field.of("start", FieldType.INT64),
                        Field.nullable("stop", FieldType.INT64)))
                .withFieldValue("start", 0L)
                .withFieldValue("stop", 1L)
                .build());

    Pipeline p = Pipeline.create();
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(
                        RunnerApi.FunctionSpec.newBuilder()
                            .setUrn(GenerateSequence.External.URN)
                            .setPayload(payload.toByteString())))
            .setNamespace(TEST_NAMESPACE)
            .build();
    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();
    assertEquals(TEST_NAMESPACE + TEST_NAME, expandedTransform.getUniqueName());
    assertThat(expandedTransform.getInputsCount(), Matchers.is(0));
    assertThat(expandedTransform.getOutputsCount(), Matchers.is(1));
    assertThat(expandedTransform.getSubtransformsCount(), greaterThan(0));
  }

  @Test
  public void testCompoundCodersForExternalConfiguration_setters() throws Exception {
    ExternalTransforms.ExternalConfigurationPayload externalConfig =
        encodeRowIntoExternalConfigurationPayload(
            Row.withSchema(
                    Schema.of(
                        Field.nullable("config_key1", FieldType.INT64),
                        Field.nullable("config_key2", FieldType.iterable(FieldType.BYTES)),
                        Field.of("config_key3", FieldType.map(FieldType.STRING, FieldType.INT64)),
                        Field.of(
                            "config_key4",
                            FieldType.map(FieldType.STRING, FieldType.array(FieldType.INT64)))))
                .withFieldValue("config_key1", 1L)
                .withFieldValue("config_key2", BYTE_LIST)
                .withFieldValue("config_key3", BYTE_KV_LIST)
                .withFieldValue("config_key4", BYTE_KV_LIST_WITH_LIST_VALUE)
                .build());

    TestConfigSetters config =
        ExpansionService.payloadToConfig(externalConfig, TestConfigSetters.class);

    assertThat(config.configKey1, Matchers.is(1L));
    assertThat(config.configKey2, contains(BYTE_LIST.toArray()));
    assertThat(config.configKey3, is(notNullValue()));
    // no-op for checker framework
    if (config.configKey3 != null) {
      assertThat(
          config.configKey3.entrySet(),
          containsInAnyOrder(
              BYTE_KV_LIST.entrySet().stream()
                  .map(
                      (entry) ->
                          allOf(
                              hasProperty("key", equalTo(entry.getKey())),
                              hasProperty("value", equalTo(entry.getValue()))))
                  .collect(Collectors.toList())));
    }
    assertThat(config.configKey4, is(notNullValue()));
    // no-op for checker framework
    if (config.configKey4 != null) {
      assertThat(
          config.configKey4.entrySet(),
          containsInAnyOrder(
              BYTE_KV_LIST_WITH_LIST_VALUE.entrySet().stream()
                  .map(
                      (entry) ->
                          allOf(
                              hasProperty("key", equalTo(entry.getKey())),
                              hasProperty("value", equalTo(entry.getValue()))))
                  .collect(Collectors.toList())));
    }
  }

  private static class TestConfigSetters {
    private @Nullable Long configKey1 = null;
    private @Nullable Iterable<byte[]> configKey2 = null;
    private @Nullable Map<byte[], Long> configKey3 = null;
    private @Nullable Map<byte[], List<Long>> configKey4 = null;

    public void setConfigKey1(@Nullable Long configKey1) {
      this.configKey1 = configKey1;
    }

    public void setConfigKey2(@Nullable Iterable<byte[]> configKey2) {
      this.configKey2 = configKey2;
    }

    public void setConfigKey3(@Nullable Map<byte[], Long> configKey3) {
      this.configKey3 = configKey3;
    }

    public void setConfigKey4(@Nullable Map<byte[], List<Long>> configKey4) {
      this.configKey4 = configKey4;
    }
  }

  @Test
  public void testCompoundCodersForExternalConfiguration_schemas() throws Exception {
    ExternalTransforms.ExternalConfigurationPayload externalConfig =
        encodeRowIntoExternalConfigurationPayload(
            Row.withSchema(
                    Schema.of(
                        Field.nullable("configKey1", FieldType.INT64),
                        Field.nullable("configKey2", FieldType.iterable(FieldType.BYTES)),
                        Field.of("configKey3", FieldType.map(FieldType.STRING, FieldType.INT64)),
                        Field.of(
                            "configKey4",
                            FieldType.map(FieldType.STRING, FieldType.array(FieldType.INT64)))))
                .withFieldValue("configKey1", 1L)
                .withFieldValue("configKey2", BYTE_LIST)
                .withFieldValue("configKey3", BYTE_KV_LIST)
                .withFieldValue("configKey4", BYTE_KV_LIST_WITH_LIST_VALUE)
                .build());
    TestConfigSchema config =
        ExpansionService.payloadToConfig(externalConfig, TestConfigSchema.class);

    assertThat(config.getConfigKey1(), Matchers.is(1L));
    assertThat(config.getConfigKey2(), contains(BYTE_LIST.toArray()));
    Map<String, Long> configKey3 = config.getConfigKey3();
    assertThat(configKey3, is(notNullValue()));
    // no-op for checker framework
    if (configKey3 != null) {
      assertThat(
          configKey3.entrySet(),
          containsInAnyOrder(
              BYTE_KV_LIST.entrySet().stream()
                  .map(
                      (entry) ->
                          allOf(
                              hasProperty("key", equalTo(entry.getKey())),
                              hasProperty("value", equalTo(entry.getValue()))))
                  .collect(Collectors.toList())));
    }
    Map<String, List<Long>> configKey4 = config.getConfigKey4();
    assertThat(configKey4, is(notNullValue()));
    // no-op for checker framework
    if (configKey4 != null) {
      assertThat(
          configKey4.entrySet(),
          containsInAnyOrder(
              BYTE_KV_LIST_WITH_LIST_VALUE.entrySet().stream()
                  .map(
                      (entry) ->
                          allOf(
                              hasProperty("key", equalTo(entry.getKey())),
                              hasProperty("value", equalTo(entry.getValue()))))
                  .collect(Collectors.toList())));
    }
  }

  @DefaultSchema(AutoValueSchema.class)
  @AutoValue
  abstract static class TestConfigSchema {
    abstract @Nullable Long getConfigKey1();

    abstract @Nullable Iterable<byte[]> getConfigKey2();

    abstract @Nullable Map<String, Long> getConfigKey3();

    abstract @Nullable Map<String, List<Long>> getConfigKey4();
  }

  @Test
  public void testExternalConfiguration_simpleSchema() throws Exception {
    ExternalTransforms.ExternalConfigurationPayload externalConfig =
        encodeRowIntoExternalConfigurationPayload(
            Row.withSchema(
                    Schema.of(
                        Field.of("bar", FieldType.STRING),
                        Field.of("foo", FieldType.INT64),
                        Field.of("list", FieldType.array(FieldType.STRING))))
                .withFieldValue("foo", 1L)
                .withFieldValue("bar", "test string")
                .withFieldValue("list", ImmutableList.of("abc", "123"))
                .build());

    TestConfigSimpleSchema config =
        ExpansionService.payloadToConfig(externalConfig, TestConfigSimpleSchema.class);

    assertThat(config.getFoo(), Matchers.is(1L));
    assertThat(config.getBar(), Matchers.is("test string"));
    assertThat(config.getList(), Matchers.is(ImmutableList.of("abc", "123")));
  }

  @Test
  public void testExpansionServiceConfig() throws Exception {
    URL expansionServiceConfigFile = Resources.getResource("./test_expansion_service_config.yaml");
    ExpansionServiceConfig config =
        ExpansionServiceConfig.parseFromYamlStream(
            Files.newInputStream(Paths.get(expansionServiceConfigFile.getPath())));
    assertEquals(3, config.getAllowlist().size());
    assertTrue(config.getAllowlist().contains("beam:transform:my_dummy_transform_1"));
    assertTrue(config.getAllowlist().contains("beam:transform:my_dummy_transform_2"));
    assertTrue(config.getAllowlist().contains("beam:transform:my_dummy_transform_3"));

    assertEquals(3, config.getDependencies().size());
    assertTrue(config.getDependencies().containsKey("beam:transform:my_dummy_transform_2"));
    assertTrue(config.getDependencies().containsKey("beam:transform:my_dummy_transform_3"));

    assertEquals(1, config.getDependencies().get("beam:transform:my_dummy_transform_2").size());
    assertEquals(
        "jars/my_dummy_transform_2_dep1.jar",
        config.getDependencies().get("beam:transform:my_dummy_transform_2").get(0).getPath());
    assertEquals(2, config.getDependencies().get("beam:transform:my_dummy_transform_3").size());

    ArrayList<String> expectedDepsOfTransform3 =
        new ArrayList<>(
            Arrays.asList(
                "jars/my_dummy_transform_3_dep1.jar", "jars/my_dummy_transform_3_dep2.jar"));

    assertTrue(
        expectedDepsOfTransform3.contains(
            config.getDependencies().get("beam:transform:my_dummy_transform_3").get(0).getPath()));
    assertTrue(
        expectedDepsOfTransform3.contains(
            config.getDependencies().get("beam:transform:my_dummy_transform_3").get(1).getPath()));
  }

  @DefaultSchema(AutoValueSchema.class)
  @AutoValue
  abstract static class TestConfigSimpleSchema {
    abstract long getFoo();

    abstract String getBar();

    abstract List<String> getList();
  }

  private static ExternalTransforms.ExternalConfigurationPayload
      encodeRowIntoExternalConfigurationPayload(Row row) {
    ByteStringOutputStream outputStream = new ByteStringOutputStream();
    try {
      SchemaCoder.of(row.getSchema()).encode(row, outputStream);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    return ExternalTransforms.ExternalConfigurationPayload.newBuilder()
        .setSchema(SchemaTranslation.schemaToProto(row.getSchema(), true))
        .setPayload(outputStream.toByteString())
        .build();
  }

  public Set<String> allIds(RunnerApi.Components components) {
    Set<String> all = new HashSet<>();
    all.addAll(components.getTransformsMap().keySet());
    all.addAll(components.getPcollectionsMap().keySet());
    all.addAll(components.getCodersMap().keySet());
    all.addAll(components.getWindowingStrategiesMap().keySet());
    all.addAll(components.getEnvironmentsMap().keySet());
    return all;
  }
}
