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

import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.runtime.testutils.CommonTestUtils;
import org.apache.flink.runtime.testutils.CommonTestUtils.PipeForwarder;

import java.io.File;
import java.io.IOException;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.Map.Entry;

import static org.apache.flink.runtime.testutils.CommonTestUtils.getCurrentClasspath;
import static org.apache.flink.runtime.testutils.CommonTestUtils.getJavaCommandPath;
import static org.apache.flink.util.Preconditions.checkNotNull;

/**
 * Utility class wrapping {@link ProcessBuilder} and pre-configuring it with common options.
 */
public class TestProcessBuilder {
	private final String javaCommand = checkNotNull(getJavaCommandPath());

	private final ArrayList<String> jvmArgs = new ArrayList<>();
	private final ArrayList<String> mainClassArgs = new ArrayList<>();

	private final String mainClass;

	private MemorySize jvmMemory = MemorySize.parse("80mb");

	public TestProcessBuilder(String mainClass) throws IOException {
		File tempLogFile = File.createTempFile(getClass().getSimpleName() + "-", "-log4j.properties");
		tempLogFile.deleteOnExit();
		CommonTestUtils.printLog4jDebugConfig(tempLogFile);

		jvmArgs.add("-Dlog.level=DEBUG");
		jvmArgs.add("-Dlog4j.configurationFile=file:" + tempLogFile.getAbsolutePath());
		jvmArgs.add("-classpath");
		jvmArgs.add(getCurrentClasspath());

		this.mainClass = mainClass;
	}

	public TestProcess start() throws IOException {
		final ArrayList<String> commands = new ArrayList<>();

		commands.add(javaCommand);
		commands.add(String.format("-Xms%dm", jvmMemory.getMebiBytes()));
		commands.add(String.format("-Xmx%dm", jvmMemory.getMebiBytes()));
		commands.addAll(jvmArgs);
		commands.add(mainClass);
		commands.addAll(mainClassArgs);

		StringWriter processOutput = new StringWriter();
		StringWriter errorOutput = new StringWriter();
		Process process = new ProcessBuilder(commands).start();
		new PipeForwarder(process.getInputStream(), processOutput);
		new PipeForwarder(process.getErrorStream(), errorOutput);

		return new TestProcess(process, processOutput, errorOutput);
	}

	public TestProcessBuilder setJvmMemory(MemorySize jvmMemory) {
		this.jvmMemory = jvmMemory;
		return this;
	}

	public TestProcessBuilder addJvmArg(String arg) {
		jvmArgs.add(arg);
		return this;
	}

	public TestProcessBuilder addMainClassArg(String arg) {
		mainClassArgs.add(arg);
		return this;
	}

	public TestProcessBuilder addConfigAsMainClassArgs(Configuration config) {
		for (Entry<String, String> keyValue: config.toMap().entrySet()) {
			addMainClassArg("--" + keyValue.getKey());
			addMainClassArg(keyValue.getValue());
		}
		return this;
	}

	/**
	 * {@link Process} with it's {@code processOutput}.
	 */
	public static class TestProcess {
		private final Process process;
		private final StringWriter processOutput;
		private final StringWriter errorOutput;

		public TestProcess(Process process, StringWriter processOutput, StringWriter errorOutput) {
			this.process = process;
			this.processOutput = processOutput;
			this.errorOutput = errorOutput;
		}

		public Process getProcess() {
			return process;
		}

		public StringWriter getProcessOutput() {
			return processOutput;
		}

		public StringWriter getErrorOutput() {
			return errorOutput;
		}

		public void destroy() {
			process.destroy();
		}
	}
}