/*
 * 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.studio.admin.builder.generator.service.generator.workflow.sections;

import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.Node;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.NodeType;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.nodedata.CodeNodeData;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.dsl.DSLDialectType;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.generator.workflow.NodeSection;

import com.alibaba.cloud.ai.studio.admin.builder.generator.utils.ObjectToCodeUtil;
import org.springframework.stereotype.Component;

import java.util.List;

// TODO: 支持异常分支
@Component
public class CodeNodeSection implements NodeSection<CodeNodeData> {

	@Override
	public boolean support(NodeType nodeType) {
		return NodeType.CODE.equals(nodeType);
	}

	@Override
	public String render(Node node, String varName) {
		CodeNodeData nodeData = ((CodeNodeData) node.getData());
		return String.format("""
				// —— CodeNode [%s] ——
				CodeExecutorNodeAction %s = CodeExecutorNodeAction.builder()
				        .codeExecutor(codeExecutor)
				        .codeLanguage("%s")
				        .code(%s)
				        .codeStyle(%s)
				        .config(codeExecutionConfig)
				        .params(%s)
				        .outputKey("%s")
				        .build();
				stateGraph.addNode("%s", AsyncNodeAction.node_async(
				    wrapperCodeNodeAction(%s, "%s", "%s", %d, %d, %s)
				));

				""", node.getId(), varName, nodeData.getCodeLanguage(), ObjectToCodeUtil.toCode(nodeData.getCode()),
				ObjectToCodeUtil.toCode(nodeData.getCodeStyle()), ObjectToCodeUtil.toCode(nodeData.getInputParams()),
				nodeData.getOutputKey(), varName, varName, nodeData.getOutputKey(), varName,
				nodeData.getMaxRetryCount(), nodeData.getRetryIntervalMs(),
				ObjectToCodeUtil.toCode(nodeData.getDefaultValue()));
	}

	@Override
	public String assistMethodCode(DSLDialectType dialectType) {
		return switch (dialectType) {
			case DIFY, STUDIO ->
				"""
						   private static final CodeExecutionConfig codeExecutionConfig;
						   private static final CodeExecutor codeExecutor;

						   static {
						       // todo: configure your own code execution configuration
						       try {
						           Path tempDir = Files.createTempDirectory("code-execution-workdir-");
						           tempDir.toFile().deleteOnExit();
						           codeExecutionConfig = new CodeExecutionConfig().setWorkDir(tempDir.toString());
						           codeExecutor = new LocalCommandlineCodeExecutor();
						       } catch (Exception e) {
						           throw new RuntimeException(e);
						       }
						   }

						   private NodeAction wrapperCodeNodeAction(NodeAction codeNodeAction, String key,
						                    String nodeName, int maxRetryCount, int retryIntervalMs, Map<String, Object> defaultValue) {
						       return state -> {
						           int count = maxRetryCount;
						           while (count-- > 0) {
						               try {
						                   // 将代码运行的结果拆包
						                   Map<String, Object> result = codeNodeAction.apply(state);
						                   Object object = result.get(key);
						                   if(!(object instanceof Map)) {
						                       throw new RuntimeException("unexcepted result");
						                   }
						                   return ((Map<String, Object>) object).entrySet().stream()
						                           .collect(Collectors.toMap(
						                                   entry -> nodeName + "_" + entry.getKey(),
						                                   Map.Entry::getValue
						                           ));
						               } catch (Exception e) {
						                   Thread.sleep(retryIntervalMs);
						               }
						           }
						           if(defaultValue != null) {
						               return defaultValue;
						           } else {
						               throw new RuntimeException("code execution failed!");
						           }
						       };
						   }
						""";
			default -> "";
		};
	}

	@Override
	public List<String> getImports() {
		return List.of("com.alibaba.cloud.ai.graph.node.code.CodeExecutorNodeAction",
				"com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionConfig",
				"com.alibaba.cloud.ai.graph.node.code.CodeExecutor",
				"com.alibaba.cloud.ai.graph.node.code.LocalCommandlineCodeExecutor", "java.io.IOException",
				"java.nio.file.Files", "java.nio.file.Path", "java.util.stream.Collectors",
				"com.alibaba.cloud.ai.graph.node.code.entity.CodeParam",
				"com.alibaba.cloud.ai.graph.node.code.entity.CodeStyle");
	}

}
