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

import java.io.File;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.Permission;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
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.util.construction.resources.PipelineResources;
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.flink.client.cli.CliFrontend;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.RestOptions;
import org.apache.flink.runtime.client.JobStatusMessage;
import org.apache.flink.runtime.minicluster.MiniClusterConfiguration;
import org.apache.flink.runtime.minicluster.RpcServiceSharing;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.rules.Timeout;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** End-to-end submission test of Beam jobs on a Flink cluster. */
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
})
public class FlinkSubmissionTest {

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

  @ClassRule public static final TemporaryFolder TEMP_FOLDER = new TemporaryFolder();
  private static final Map<String, String> ENV = System.getenv();
  private static final SecurityManager SECURITY_MANAGER = System.getSecurityManager();

  /** Flink cluster that runs over the lifespan of the tests. */
  private static transient RemoteMiniCluster flinkCluster;

  /** Each test has a timeout of 60 seconds (for safety). */
  @Rule public Timeout timeout = new Timeout(60, TimeUnit.SECONDS);

  /** Counter which keeps track of the number of jobs submitted. */
  private static int expectedNumberOfJobs;

  @BeforeClass
  public static void beforeClass() throws Exception {
    Configuration config = new Configuration();
    // Avoid port collision in parallel tests on the same machine
    config.set(RestOptions.PORT, 0);

    MiniClusterConfiguration clusterConfig =
        new MiniClusterConfiguration.Builder()
            .setConfiguration(config)
            .setNumTaskManagers(1)
            .setNumSlotsPerTaskManager(1)
            // Create a shared actor system for all cluster services
            .setRpcServiceSharing(RpcServiceSharing.SHARED)
            .build();

    flinkCluster = new RemoteMiniClusterImpl(clusterConfig);
    flinkCluster.start();
    prepareEnvironment();
  }

  @AfterClass
  public static void afterClass() throws Exception {
    restoreEnvironment();
    flinkCluster.close();
    flinkCluster = null;
  }

  @Test
  public void testSubmissionBatch() throws Exception {
    runSubmission(false, false);
  }

  @Test
  public void testSubmissionStreaming() throws Exception {
    runSubmission(false, true);
  }

  @Test
  public void testDetachedSubmissionBatch() throws Exception {
    runSubmission(true, false);
  }

  @Test
  public void testDetachedSubmissionStreaming() throws Exception {
    runSubmission(true, true);
  }

  private void runSubmission(boolean isDetached, boolean isStreaming) throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(FlinkPipelineOptions.class).setStreaming(isStreaming);
    options.setTempLocation(TEMP_FOLDER.getRoot().getPath());
    String jarPath =
        Iterables.getFirst(
            PipelineResources.detectClassPathResourcesToStage(getClass().getClassLoader(), options),
            null);

    try {
      throwExceptionOnSystemExit();
      ImmutableList.Builder<String> argsBuilder = ImmutableList.builder();
      argsBuilder.add("run").add("-c").add(getClass().getName());
      if (isDetached) {
        argsBuilder.add("-d");
      }
      argsBuilder.add(jarPath);
      argsBuilder.add("--runner=flink");

      if (isStreaming) {
        argsBuilder.add("--streaming");
      }

      FlinkSubmissionTest.expectedNumberOfJobs++;
      ImmutableList<String> args = argsBuilder.build();
      // Run end-to-end test
      CliFrontend.main(args.toArray(new String[0]));
    } catch (SystemExitException e) {
      // The CliFrontend exited and we can move on to check if the job has finished
    } finally {
      restoreDefaultSystemExitBehavior();
    }

    waitUntilJobIsCompleted();
  }

  private void waitUntilJobIsCompleted() throws Exception {
    while (true) {
      Collection<JobStatusMessage> allJobsStates = flinkCluster.listJobs().get();
      if (allJobsStates.size() == expectedNumberOfJobs
          && allJobsStates.stream()
              .allMatch(jobStatus -> jobStatus.getJobState().isTerminalState())) {
        LOG.info(
            "All job finished with statuses: {}",
            allJobsStates.stream().map(j -> j.getJobState().name()).collect(Collectors.toList()));
        return;
      }
      Thread.sleep(50);
    }
  }

  /** The Flink program which is executed by the CliFrontend. */
  public static void main(String[] args) {
    FlinkPipelineOptions options =
        PipelineOptionsFactory.fromArgs(args).withValidation().as(FlinkPipelineOptions.class);
    options.setRunner(FlinkRunner.class);
    options.setParallelism(1);
    Pipeline p = Pipeline.create(options);
    p.apply(GenerateSequence.from(0).to(1));
    p.run();
  }

  private static void prepareEnvironment() throws Exception {
    // Write a Flink config
    File file = TEMP_FOLDER.newFile("config.yaml");
    String config =
        String.format(
            "rest:\n  port: '%d'\njobmanager:\n  rpc:\n    address: %s\n    port: '%d'",
            flinkCluster.getRestPort(), "localhost", flinkCluster.getClusterPort());

    Files.write(file.toPath(), config.getBytes(StandardCharsets.UTF_8));

    // Create a new environment with the location of the Flink config for CliFrontend
    ImmutableMap<String, String> newEnv =
        ImmutableMap.<String, String>builder()
            .putAll(ENV.entrySet())
            .put(ConfigConstants.ENV_FLINK_CONF_DIR, file.getParent())
            .build();

    modifyEnv(newEnv);
  }

  private static void restoreEnvironment() throws Exception {
    modifyEnv(ENV);
  }

  /**
   * We modify the JVM's environment variables here. This is necessary for the end-to-end test
   * because Flink's CliFrontend requires a Flink configuration file for which the location can only
   * be set using the {@code ConfigConstants.ENV_FLINK_CONF_DIR} environment variable.
   */
  private static void modifyEnv(Map<String, String> env) throws Exception {
    Class processEnv = Class.forName("java.lang.ProcessEnvironment");
    Field envField = processEnv.getDeclaredField("theUnmodifiableEnvironment");

    Field modifiersField = Field.class.getDeclaredField("modifiers");
    modifiersField.setAccessible(true);
    modifiersField.setInt(envField, envField.getModifiers() & ~Modifier.FINAL);

    envField.setAccessible(true);
    envField.set(null, env);
    envField.setAccessible(false);

    modifiersField.setInt(envField, envField.getModifiers() & Modifier.FINAL);
    modifiersField.setAccessible(false);
  }

  /** Prevents the CliFrontend from calling System.exit. */
  private static void throwExceptionOnSystemExit() {
    System.setSecurityManager(
        new SecurityManager() {
          @Override
          public void checkPermission(Permission permission) {
            if (permission.getName().startsWith("exitVM")) {
              throw new SystemExitException();
            }
            if (SECURITY_MANAGER != null) {
              SECURITY_MANAGER.checkPermission(permission);
            }
          }
        });
  }

  private static void restoreDefaultSystemExitBehavior() {
    System.setSecurityManager(SECURITY_MANAGER);
  }

  private static class SystemExitException extends SecurityException {}
}
