/*
 * 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.runners.dataflow.worker.streaming.config;

import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import com.google.api.services.dataflow.model.StreamingComputationConfig;
import com.google.api.services.dataflow.model.StreamingConfigTask;
import com.google.api.services.dataflow.model.WorkItem;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import org.apache.beam.runners.dataflow.worker.OperationalLimits;
import org.apache.beam.runners.dataflow.worker.WorkUnitClient;
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.joda.time.Duration;
import org.junit.After;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.internal.stubbing.answers.Returns;

@RunWith(JUnit4.class)
public class StreamingEngineComputationConfigFetcherTest {

  private final WorkUnitClient mockDataflowServiceClient =
      mock(WorkUnitClient.class, new Returns(Optional.empty()));
  private StreamingEngineComputationConfigFetcher streamingEngineConfigFetcher;

  private StreamingEngineComputationConfigFetcher createConfigFetcher(
      boolean waitForInitialConfig,
      long globalConfigRefreshPeriod,
      StreamingGlobalConfigHandleImpl globalConfigHandle) {
    return StreamingEngineComputationConfigFetcher.forTesting(
        !waitForInitialConfig,
        globalConfigRefreshPeriod,
        mockDataflowServiceClient,
        globalConfigHandle,
        ignored -> Executors.newSingleThreadScheduledExecutor());
  }

  @After
  public void cleanUp() {
    streamingEngineConfigFetcher.stop();
  }

  @Test
  public void testStart_requiresInitialConfig() throws IOException, InterruptedException {
    WorkItem initialConfig =
        new WorkItem()
            .setJobId("job")
            .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(10L));
    CountDownLatch waitForInitialConfig = new CountDownLatch(1);
    Set<StreamingGlobalConfig> receivedPipelineConfig = new HashSet<>();
    when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
        .thenReturn(Optional.of(initialConfig));
    StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
    globalConfigHandle.registerConfigObserver(
        config -> {
          try {
            receivedPipelineConfig.add(config);
            waitForInitialConfig.await();
          } catch (InterruptedException e) {
            throw new RuntimeException(e);
          }
        });
    streamingEngineConfigFetcher =
        createConfigFetcher(/* waitForInitialConfig= */ true, 0, globalConfigHandle);
    Thread asyncStartConfigLoader = new Thread(streamingEngineConfigFetcher::start);
    asyncStartConfigLoader.start();
    waitForInitialConfig.countDown();
    asyncStartConfigLoader.join();
    StreamingGlobalConfig.Builder configBuilder =
        StreamingGlobalConfig.builder()
            .setOperationalLimits(
                OperationalLimits.builder()
                    .setMaxWorkItemCommitBytes(
                        initialConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
                    .build());
    assertThat(receivedPipelineConfig).containsExactly(configBuilder.build());
  }

  @Test
  public void testStart_startsPeriodicConfigRequests() throws IOException, InterruptedException {
    WorkItem firstConfig =
        new WorkItem()
            .setJobId("job")
            .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(10L));
    WorkItem secondConfig =
        new WorkItem()
            .setJobId("job")
            .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(15L));
    WorkItem thirdConfig =
        new WorkItem()
            .setJobId("job")
            .setStreamingConfigTask(new StreamingConfigTask().setMaxWorkItemCommitBytes(100L));
    CountDownLatch numExpectedRefreshes = new CountDownLatch(3);
    Set<StreamingGlobalConfig> receivedPipelineConfig = new HashSet<>();
    when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
        .thenReturn(Optional.of(firstConfig))
        .thenReturn(Optional.of(secondConfig))
        // ConfigFetcher should still fetch subsequent configs on error.
        .thenThrow(new IOException("something bad happened."))
        .thenReturn(Optional.of(thirdConfig))
        // ConfigFetcher should not do anything with a config that doesn't contain a
        // StreamingConfigTask.
        .thenReturn(Optional.of(new WorkItem().setJobId("jobId")));
    StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
    globalConfigHandle.registerConfigObserver(
        config -> {
          receivedPipelineConfig.add(config);
          numExpectedRefreshes.countDown();
        });
    streamingEngineConfigFetcher =
        createConfigFetcher(
            /* waitForInitialConfig= */ true, Duration.millis(100).getMillis(), globalConfigHandle);

    Thread asyncStartConfigLoader = new Thread(streamingEngineConfigFetcher::start);
    asyncStartConfigLoader.start();
    numExpectedRefreshes.await();
    asyncStartConfigLoader.join();
    assertThat(receivedPipelineConfig)
        .containsExactly(
            StreamingGlobalConfig.builder()
                .setOperationalLimits(
                    OperationalLimits.builder()
                        .setMaxWorkItemCommitBytes(
                            firstConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
                        .build())
                .build(),
            StreamingGlobalConfig.builder()
                .setOperationalLimits(
                    OperationalLimits.builder()
                        .setMaxWorkItemCommitBytes(
                            secondConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
                        .build())
                .build(),
            StreamingGlobalConfig.builder()
                .setOperationalLimits(
                    OperationalLimits.builder()
                        .setMaxWorkItemCommitBytes(
                            thirdConfig.getStreamingConfigTask().getMaxWorkItemCommitBytes())
                        .build())
                .build());
  }

  @Test
  public void testGetComputationConfig() throws IOException {
    StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
    streamingEngineConfigFetcher =
        createConfigFetcher(/* waitForInitialConfig= */ false, 0, globalConfigHandle);
    String computationId = "computationId";
    String stageName = "stageName";
    String systemName = "systemName";
    StreamingComputationConfig pipelineConfig =
        new StreamingComputationConfig()
            .setComputationId(computationId)
            .setStageName(stageName)
            .setSystemName(systemName)
            .setInstructions(ImmutableList.of());

    WorkItem workItem =
        new WorkItem()
            .setStreamingConfigTask(
                new StreamingConfigTask()
                    .setStreamingComputationConfigs(ImmutableList.of(pipelineConfig)));

    when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString()))
        .thenReturn(Optional.of(workItem));
    Optional<ComputationConfig> actualPipelineConfig =
        streamingEngineConfigFetcher.fetchConfig(computationId);

    assertTrue(actualPipelineConfig.isPresent());
    assertThat(actualPipelineConfig.get())
        .isEqualTo(
            ComputationConfig.create(
                StreamingEngineComputationConfigFetcher.createMapTask(pipelineConfig),
                ImmutableMap.of(),
                ImmutableMap.of()));
  }

  @Test
  public void testGetComputationConfig_noComputationPresent() throws IOException {
    Set<StreamingGlobalConfig> receivedPipelineConfig = new HashSet<>();
    StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
    globalConfigHandle.registerConfigObserver(receivedPipelineConfig::add);
    streamingEngineConfigFetcher =
        createConfigFetcher(/* waitForInitialConfig= */ false, 0, globalConfigHandle);
    when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString()))
        .thenReturn(Optional.empty());
    Optional<ComputationConfig> pipelineConfig =
        streamingEngineConfigFetcher.fetchConfig("someComputationId");
    assertFalse(pipelineConfig.isPresent());
    assertThat(receivedPipelineConfig).isEmpty();
  }

  @Test
  public void testGetComputationConfig_fetchConfigFromDataflowError() throws IOException {
    StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
    streamingEngineConfigFetcher =
        createConfigFetcher(/* waitForInitialConfig= */ false, 0, globalConfigHandle);
    RuntimeException e = new RuntimeException("something bad happened.");
    when(mockDataflowServiceClient.getStreamingConfigWorkItem(anyString())).thenThrow(e);
    Throwable fetchConfigError =
        assertThrows(
            RuntimeException.class,
            () -> streamingEngineConfigFetcher.fetchConfig("someComputationId"));
    assertThat(fetchConfigError).isSameInstanceAs(e);
  }

  @Test
  public void test_streamingEngineStateTagEncodingVersion()
      throws IOException, InterruptedException {
    for (Optional<Integer> tagEncoding :
        Arrays.<Optional<Integer>>asList(Optional.empty(), Optional.of(1), Optional.of(2))) {
      StreamingConfigTask streamingConfigTask =
          new StreamingConfigTask().setMaxWorkItemCommitBytes(100L);
      tagEncoding.ifPresent(
          version -> streamingConfigTask.setStreamingEngineStateTagEncodingVersion(version));
      WorkItem initialConfig =
          new WorkItem().setJobId("job").setStreamingConfigTask(streamingConfigTask);
      CountDownLatch waitForInitialConfig = new CountDownLatch(1);
      Set<StreamingGlobalConfig> receivedPipelineConfig = new HashSet<>();
      when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
          .thenReturn(Optional.of(initialConfig));
      StreamingGlobalConfigHandleImpl globalConfigHandle = new StreamingGlobalConfigHandleImpl();
      globalConfigHandle.registerConfigObserver(
          config -> {
            receivedPipelineConfig.add(config);
            waitForInitialConfig.countDown();
          });
      streamingEngineConfigFetcher =
          createConfigFetcher(/* waitForInitialConfig= */ true, 0, globalConfigHandle);
      Thread asyncStartConfigLoader = new Thread(streamingEngineConfigFetcher::start);
      asyncStartConfigLoader.start();
      waitForInitialConfig.await();
      asyncStartConfigLoader.join();
      assertEquals(1, receivedPipelineConfig.size());
      assertEquals(
          Objects.equals(2, streamingConfigTask.getStreamingEngineStateTagEncodingVersion()),
          receivedPipelineConfig.iterator().next().enableStateTagEncodingV2());
    }
  }
}
