/*
 * Copyright 2024-2026 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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 com.alibaba.cloud.ai.graph.node.code;

import com.alibaba.cloud.ai.graph.node.code.entity.CodeBlock;
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionConfig;
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionResult;
import com.alibaba.cloud.ai.graph.utils.CodeUtils;
import com.alibaba.cloud.ai.graph.utils.FileUtils;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.exec.CommandLine;
import org.apache.commons.exec.DefaultExecutor;
import org.apache.commons.exec.ExecuteException;
import org.apache.commons.exec.ExecuteWatchdog;
import org.apache.commons.exec.PumpStreamHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @author HeYQ
 * @since 2024-12-02 17:23
 */
public class LocalCommandlineCodeExecutor implements CodeExecutor {

	private static final Logger logger = LoggerFactory.getLogger(LocalCommandlineCodeExecutor.class);

	@Override
	public CodeExecutionResult executeCodeBlocks(List<CodeBlock> codeBlockList, CodeExecutionConfig codeExecutionConfig)
			throws Exception {
		StringBuilder allLogs = new StringBuilder();
		CodeExecutionResult result;
		for (int i = 0; i < codeBlockList.size(); i++) {
			CodeBlock codeBlock = codeBlockList.get(i);
			String language = codeBlock.language();
			String code = codeBlock.code();
			logger.info("\n>>>>>>>> EXECUTING CODE BLOCK {} (inferred language is {})...", i + 1, language);
			// "bash", "shell", "sh", "python"
			result = executeCode(language, code, codeExecutionConfig);
			allLogs.append("\n").append(result.logs());
			if (result.exitCode() != 0) {
				return new CodeExecutionResult(result.exitCode(), allLogs.toString());
			}
		}
		return new CodeExecutionResult(0, allLogs.toString());
	}

	@Override
	public void restart() {

		logger.warn("Restarting local command line code executor is not supported. No action is taken.");
	}

	public CodeExecutionResult executeCode(String language, String code, CodeExecutionConfig config) throws Exception {
		if (Objects.isNull(language) || Objects.isNull(code)) {
			throw new Exception("Either language or code must be provided.");
		}
		String workDir = config.getWorkDir();
		String codeHash = DigestUtils.md5Hex(code);
		String fileExt = CodeUtils.getFileExtForLanguage(language);
		String filename = String.format("tmp_code_%s.%s", codeHash, fileExt);

		// write the code string to a file specified by the filename.
		FileUtils.writeCodeToFile(workDir, filename, code);

		// Copy required JAR files to workDir if language is Java
		if ("java".equals(language)) {
			FileUtils.copyResourceJarToWorkDir(workDir);
		}

		CodeExecutionResult executionResult = executeCodeLocally(language, workDir, filename, config);

		FileUtils.deleteFile(workDir, filename);

		// Delete JAR files if language is Java
		if ("java".equals(language)) {
			FileUtils.deleteResourceJarFromWorkDir(workDir);
		}
		return executionResult;
	}

	private CodeExecutionResult executeCodeLocally(String language, String workDir, String filename,
			CodeExecutionConfig config) throws Exception {
		// Set up command line based on language
		String executable = CodeUtils.getExecutableForLanguage(language);
		CommandLine commandLine = new CommandLine(executable);

		if ("java".equals(language)) {
			commandLine.addArgument("-cp");
			StringBuilder classPathBuilder = new StringBuilder();
			classPathBuilder.append(".").append(File.pathSeparator).append(workDir);

			// Add all JAR files in workDir to classpath
			try {
				Path workDirPath = Path.of(workDir);
				if (Files.exists(workDirPath)) {
					try (var stream = Files.walk(workDirPath)) {
						stream.filter(path -> path.toString().endsWith(".jar")).forEach(jarPath -> {
							classPathBuilder.append(File.pathSeparator).append(jarPath.toString());
						});
					}
				}
			}
			catch (IOException e) {
				logger.warn("Failed to scan JAR files in work directory", e);
			}

			if (config.getClassPath() != null && !config.getClassPath().isEmpty()) {
				classPathBuilder.append(File.pathSeparator).append(config.getClassPath());
			}

			String classPath = classPathBuilder.toString();
			commandLine.addArgument(classPath).addArgument(filename);
		}
		else {
			commandLine.addArgument(filename);
		}

		// Configure executor
		DefaultExecutor executor = new DefaultExecutor();
		executor.setWorkingDirectory(new File(workDir));
		executor.setExitValue(0);

		// Set up stream handling
		ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
		ByteArrayOutputStream errorStream = new ByteArrayOutputStream();
		executor.setStreamHandler(new PumpStreamHandler(outputStream, errorStream));

		// Set timeout
		executor.setWatchdog(new ExecuteWatchdog(TimeUnit.SECONDS.toMillis(config.getTimeout())));

		try {
			executor.execute(commandLine);
			return new CodeExecutionResult(0, outputStream.toString().trim());
		}
		catch (ExecuteException e) {
			String errorOutput = errorStream.toString()
				.replace(Path.of(workDir).toAbsolutePath() + File.separator, "")
				.trim();
			return new CodeExecutionResult(e.getExitValue(), errorOutput);
		}
		catch (IOException e) {
			throw new Exception("Failed to execute code", e);
		}
		finally {
			// Cleanup Java class files
			if ("java".equals(language)) {
				FileUtils.deleteFile(workDir, filename.replace(".java", ".class"));
			}
		}
	}

}
