/*
 * 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.inlong.sort.tests.utils;

import org.apache.commons.io.IOUtils;
import org.apache.flink.api.common.JobStatus;
import org.apache.flink.api.common.time.Deadline;
import org.apache.flink.client.deployment.StandaloneClusterId;
import org.apache.flink.client.program.rest.RestClusterClient;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.RestOptions;
import org.apache.flink.runtime.client.JobStatusMessage;
import org.apache.flink.runtime.jobmaster.JobMaster;
import org.apache.flink.runtime.taskexecutor.TaskExecutor;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.util.TestLogger;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.rules.TemporaryFolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.Container.ExecResult;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.Network;
import org.testcontainers.containers.output.Slf4jLogConsumer;
import org.testcontainers.images.builder.Transferable;
import org.testcontainers.lifecycle.Startables;

import javax.annotation.Nullable;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.stream.Stream;

import static org.apache.flink.util.Preconditions.checkState;

/**
 * End to end base test environment for test sort-connectors.
 * Every link : MySQL -> Xxx (Test connector) -> MySQL
 */
public abstract class FlinkContainerTestEnv extends TestLogger {

    private static final Logger JM_LOG = LoggerFactory.getLogger(JobMaster.class);
    private static final Logger TM_LOG = LoggerFactory.getLogger(TaskExecutor.class);
    private static final Logger MYSQL_LOG = LoggerFactory.getLogger(MySqlContainer.class);
    private static final Logger LOG = LoggerFactory.getLogger(FlinkContainerTestEnv.class);

    private static final Path SORT_DIST_JAR = TestUtils.getResource("sort-dist.jar");
    // ------------------------------------------------------------------------------------------
    // Flink Variables
    // ------------------------------------------------------------------------------------------
    private static final int JOB_MANAGER_REST_PORT = 8081;
    private static final int DEBUG_PORT = 20000;
    private static final String FLINK_BIN = "bin";
    private static final String INTER_CONTAINER_JM_ALIAS = "jobmanager";
    private static final String INTER_CONTAINER_TM_ALIAS = "taskmanager";
    private static final String FLINK_PROPERTIES = String.join("\n", Arrays.asList(
            "jobmanager.rpc.address: jobmanager",
            "taskmanager.numberOfTaskSlots: 10",
            "parallelism.default: 4",
            "env.java.opts.jobmanager: -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=20000",
            "env.java.opts.taskmanager: -agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=20000",
            // this is needed for oracle-cdc tests.
            // see https://stackoverflow.com/a/47062742/4915129
            "env.java.opts: -Doracle.jdbc.timezoneAsRegion=false"));

    @ClassRule
    public static final Network NETWORK = Network.newNetwork();

    @Rule
    public final TemporaryFolder temporaryFolder = new TemporaryFolder();

    @Nullable
    private static RestClusterClient<StandaloneClusterId> restClusterClient;

    private static GenericContainer<?> jobManager;
    private static GenericContainer<?> taskManager;
    // ----------------------------------------------------------------------------------------
    // MYSQL Variables
    // ----------------------------------------------------------------------------------------
    protected static final String MYSQL_DRIVER_CLASS = "com.mysql.cj.jdbc.Driver";
    private static final String INTER_CONTAINER_MYSQL_ALIAS = "mysql";
    @ClassRule
    public static final MySqlContainer MYSQL =
            (MySqlContainer) new MySqlContainer()
                    .withConfigurationOverride("docker/mysql/my.cnf")
                    .withSetupSQL("docker/mysql/setup.sql")
                    .withDatabaseName("test")
                    .withUsername("flinkuser")
                    .withPassword("flinkpw")
                    .withUrlParam("allowMultiQueries", "true")
                    .withNetwork(NETWORK)
                    .withNetworkAliases(INTER_CONTAINER_MYSQL_ALIAS)
                    .withLogConsumer(new Slf4jLogConsumer(MYSQL_LOG));

    @BeforeClass
    public static void before() {
        LOG.info("Starting containers...");
        jobManager =
                new GenericContainer<>("flink:1.13.5-scala_2.11")
                        .withCommand("jobmanager")
                        .withNetwork(NETWORK)
                        .withNetworkAliases(INTER_CONTAINER_JM_ALIAS)
                        .withExposedPorts(JOB_MANAGER_REST_PORT, DEBUG_PORT)
                        .withEnv("FLINK_PROPERTIES", FLINK_PROPERTIES)
                        .withLogConsumer(new Slf4jLogConsumer(JM_LOG));
        taskManager =
                new GenericContainer<>("flink:1.13.5-scala_2.11")
                        .withCommand("taskmanager")
                        .withNetwork(NETWORK)
                        .withNetworkAliases(INTER_CONTAINER_TM_ALIAS)
                        .withExposedPorts(DEBUG_PORT)
                        .withEnv("FLINK_PROPERTIES", FLINK_PROPERTIES)
                        .dependsOn(jobManager)
                        .withLogConsumer(new Slf4jLogConsumer(TM_LOG));

        Startables.deepStart(Stream.of(jobManager)).join();
        Startables.deepStart(Stream.of(taskManager)).join();
        LOG.info("Containers are started.");
    }

    @AfterClass
    public static void after() {
        if (restClusterClient != null) {
            restClusterClient.close();
        }
        if (jobManager != null) {
            jobManager.stop();
        }
        if (taskManager != null) {
            taskManager.stop();
        }
        if (MYSQL != null) {
            MYSQL.stop();
        }
    }

    /**
     * Submits a SQL job to the running cluster.
     *
     * <p><b>NOTE:</b> You should not use {@code '\t'}.
     */
    public void submitSQLJob(String sqlFile, Path... jars)
            throws IOException, InterruptedException {
        final List<String> commands = new ArrayList<>();
        String containerSqlFile = copyToContainerTmpPath(jobManager, sqlFile);
        commands.add(FLINK_BIN + "/flink run -d");
        commands.add("-c org.apache.inlong.sort.Entrance");
        commands.add(copyToContainerTmpPath(jobManager, constructDistJar(jars)));
        commands.add("--sql.script.file");
        commands.add(containerSqlFile);

        ExecResult execResult =
                jobManager.execInContainer("bash", "-c", String.join(" ", commands));
        LOG.info(execResult.getStdout());
        LOG.error(execResult.getStderr());
        if (execResult.getExitCode() != 0) {
            throw new AssertionError("Failed when submitting the SQL job.");
        }
    }

    public void submitGroupFileJob(String groupFile, Path... jars)
            throws IOException, InterruptedException {
        final List<String> commands = new ArrayList<>();
        String containerGroupFile = copyToContainerTmpPath(jobManager, groupFile);
        commands.add(FLINK_BIN + "/flink run -d");
        commands.add("-c org.apache.inlong.sort.Entrance");
        commands.add(copyToContainerTmpPath(jobManager, constructDistJar(jars)));
        commands.add("--group.info.file");
        commands.add(containerGroupFile);

        ExecResult execResult =
                jobManager.execInContainer("bash", "-c", String.join(" ", commands));
        LOG.info(execResult.getStdout());
        LOG.error(execResult.getStderr());
        if (execResult.getExitCode() != 0) {
            throw new AssertionError("Failed when submitting the SQL job.");
        }
    }

    /**
     * Get {@link RestClusterClient} connected to this FlinkContainer.
     *
     * <p>This method lazily initializes the REST client on-demand.
     */
    public RestClusterClient<StandaloneClusterId> getRestClusterClient() {
        if (restClusterClient != null) {
            return restClusterClient;
        }
        checkState(
                jobManager.isRunning(),
                "Cluster client should only be retrieved for a running cluster");
        try {
            final Configuration clientConfiguration = new Configuration();
            clientConfiguration.set(RestOptions.ADDRESS, jobManager.getHost());
            clientConfiguration.set(
                    RestOptions.PORT, jobManager.getMappedPort(JOB_MANAGER_REST_PORT));
            this.restClusterClient =
                    new RestClusterClient<>(clientConfiguration, StandaloneClusterId.getInstance());
        } catch (Exception e) {
            throw new IllegalStateException(
                    "Failed to create client for Flink container cluster", e);
        }
        return restClusterClient;
    }

    /**
     * Polling to detect task status until the task successfully into {@link JobStatus.RUNNING}
     *
     * @param timeout
     */
    public void waitUntilJobRunning(Duration timeout) {
        RestClusterClient<?> clusterClient = getRestClusterClient();
        Deadline deadline = Deadline.fromNow(timeout);
        while (deadline.hasTimeLeft()) {
            Collection<JobStatusMessage> jobStatusMessages;
            try {
                jobStatusMessages = clusterClient.listJobs().get(10, TimeUnit.SECONDS);
            } catch (Exception e) {
                LOG.warn("Error when fetching job status.", e);
                continue;
            }
            if (jobStatusMessages != null && !jobStatusMessages.isEmpty()) {
                JobStatusMessage message = jobStatusMessages.iterator().next();
                JobStatus jobStatus = message.getJobState();
                if (jobStatus.isTerminalState()) {
                    throw new ValidationException(
                            String.format(
                                    "Job has been terminated! JobName: %s, JobID: %s, Status: %s",
                                    message.getJobName(),
                                    message.getJobId(),
                                    message.getJobState()));
                } else if (jobStatus == JobStatus.RUNNING) {
                    return;
                }
            }
        }
    }

    /**
     * Copy all other dependencies into user jar 'lib/' entry.
     * Flink per-job mode only support upload one jar to cluster.
     */
    private String constructDistJar(Path... jars) throws IOException {

        File newJar = temporaryFolder.newFile("sort-dist.jar");
        try (
                JarFile jarFile = new JarFile(SORT_DIST_JAR.toFile());
                JarOutputStream jos = new JarOutputStream(new FileOutputStream(newJar))) {
            jarFile.stream().forEach(entry -> {
                try (InputStream is = jarFile.getInputStream(entry)) {
                    jos.putNextEntry(entry);
                    jos.write(IOUtils.toByteArray(is));
                    jos.closeEntry();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            });

            for (Path jar : jars) {
                try (InputStream is = new FileInputStream(jar.toFile())) {
                    jos.putNextEntry(new JarEntry("lib/" + jar.getFileName().toString()));
                    jos.write(IOUtils.toByteArray(is));
                    jos.closeEntry();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

        }
        return newJar.getAbsolutePath();
    }

    // Should not a big file, all file data will load into memory, then copy to container.
    private String copyToContainerTmpPath(GenericContainer<?> container, String filePath) throws IOException {
        Path path = Paths.get(filePath);
        byte[] fileData = Files.readAllBytes(path);
        String containerPath = "/tmp/" + path.getFileName();
        container.copyFileToContainer(Transferable.of(fileData), containerPath);
        return containerPath;
    }
}
