/*
 * 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 java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
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.node.code.entity.CodeLanguage;
import com.alibaba.cloud.ai.graph.node.code.entity.CodeParam;
import com.alibaba.cloud.ai.graph.node.code.entity.CodeStyle;
import com.alibaba.cloud.ai.graph.node.code.entity.RunnerAndPreload;
import com.alibaba.cloud.ai.graph.node.code.java.JavaTemplateTransformer;
import com.alibaba.cloud.ai.graph.node.code.javascript.NodeJsTemplateTransformer;
import com.alibaba.cloud.ai.graph.node.code.python3.Python3TemplateTransformer;
import org.springframework.util.StringUtils;

/**
 * @author HeYQ
 * @since 2024-11-28 11:47
 */
public class CodeExecutorNodeAction implements NodeAction {

	private final CodeExecutor codeExecutor;

	private final String codeLanguage;

	private final String code;

	private final CodeExecutionConfig codeExecutionConfig;

	private final List<CodeParam> params;

	private final String outputKey;

	private final CodeStyle style;

	private static final Map<CodeLanguage, TemplateTransformer> CODE_TEMPLATE_TRANSFORMERS = Map.of(
			CodeLanguage.PYTHON3, new Python3TemplateTransformer(), CodeLanguage.PYTHON,
			new Python3TemplateTransformer(), CodeLanguage.JAVASCRIPT, new NodeJsTemplateTransformer(),
			CodeLanguage.JAVA, new JavaTemplateTransformer());

	private static final Map<CodeLanguage, String> CODE_LANGUAGE_TO_RUNNING_LANGUAGE = Map.of(CodeLanguage.JAVASCRIPT,
			"nodejs", CodeLanguage.JINJA2, CodeLanguage.PYTHON3.getValue(), CodeLanguage.PYTHON3,
			CodeLanguage.PYTHON3.getValue(), CodeLanguage.PYTHON, CodeLanguage.PYTHON.getValue(), CodeLanguage.JAVA,
			CodeLanguage.JAVA.getValue());

	public CodeExecutorNodeAction(CodeExecutor codeExecutor, String codeLanguage, String code, CodeStyle style,
			CodeExecutionConfig config, List<CodeParam> params, String outputKey) {
		this.codeExecutor = codeExecutor;
		this.codeLanguage = codeLanguage;
		this.style = style;
		this.code = code;
		this.codeExecutionConfig = config;
		this.params = params;
		this.outputKey = outputKey;
	}

	private Map<String, Object> executeWorkflowCodeTemplate(CodeLanguage language, String code,
			Map<String, Object> inputs) throws Exception {
		TemplateTransformer templateTransformer = CODE_TEMPLATE_TRANSFORMERS.get(language);
		if (templateTransformer == null) {
			throw new RuntimeException("Unsupported language: " + language);
		}

		RunnerAndPreload runnerAndPreload = templateTransformer.transformCaller(code, inputs, style);
		String response = executeCode(language, runnerAndPreload.preloadScript(), runnerAndPreload.runnerScript());

		return templateTransformer.transformResponse(response);
	}

	private String executeCode(CodeLanguage language, String preloadScript, String code) throws Exception {
		List<CodeBlock> codeBlockList = new ArrayList<>(10);
		codeBlockList.add(new CodeBlock(CODE_LANGUAGE_TO_RUNNING_LANGUAGE.get(language), code));

		CodeExecutionResult codeExecutionResult = codeExecutor.executeCodeBlocks(codeBlockList,
				this.codeExecutionConfig);
		if (codeExecutionResult.exitCode() != 0) {
			throw new RuntimeException("code execution failed, exit code: " + codeExecutionResult.exitCode()
					+ ", logs: " + codeExecutionResult.logs());
		}

		return codeExecutionResult.logs();
	}

	@Override
	public Map<String, Object> apply(OverAllState state) throws Exception {
		Map<String, Object> inputs = Optional.ofNullable(params)
			.orElse(List.of())
			.stream()
			.collect(Collectors.toUnmodifiableMap(CodeParam::argName, param -> Optional.ofNullable(param.value())
				.or(() -> StringUtils.hasText(param.stateKey()) ? state.value(param.stateKey()) : Optional.empty())
				.orElseThrow(() -> new IllegalStateException("param has no value and legal key!"))));
		Map<String, Object> resultObjectMap = executeWorkflowCodeTemplate(CodeLanguage.fromValue(codeLanguage), code,
				inputs);
		Map<String, Object> updatedState = new HashMap<>();
		if (StringUtils.hasLength(this.outputKey)) {
			updatedState.put(this.outputKey, resultObjectMap);
		}
		return updatedState;
	}

	public static Builder builder() {
		return new Builder();
	}

	public static class Builder {

		private CodeExecutor codeExecutor;

		private String codeLanguage;

		private String code;

		private CodeStyle style;

		private CodeExecutionConfig config;

		private List<CodeParam> params;

		private String outputKey;

		public Builder() {
			style = CodeStyle.EXPLICIT_PARAMETERS;
		}

		public Builder codeExecutor(CodeExecutor codeExecutor) {
			this.codeExecutor = codeExecutor;
			return this;
		}

		public Builder codeLanguage(String codeLanguage) {
			this.codeLanguage = codeLanguage;
			return this;
		}

		public Builder code(String code) {
			this.code = code;
			return this;
		}

		public Builder codeStyle(CodeStyle style) {
			this.style = style;
			return this;
		}

		public Builder config(CodeExecutionConfig config) {
			this.config = config;
			return this;
		}

		public Builder params(List<CodeParam> params) {
			this.params = List.copyOf(params);
			return this;
		}

		public Builder outputKey(String outputKey) {
			this.outputKey = outputKey;
			return this;
		}

		public CodeExecutorNodeAction build() {
			return new CodeExecutorNodeAction(codeExecutor, codeLanguage, code, style, config, params, outputKey);
		}

	}

}
