content: /*
   * Copyright 2024-2025 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.generator.service.generator.workflow.sections;
  
  import com.alibaba.cloud.ai.studio.admin.generator.model.workflow.Edge;
  import com.alibaba.cloud.ai.studio.admin.generator.model.workflow.Node;
  import com.alibaba.cloud.ai.studio.admin.generator.model.workflow.NodeType;
  import com.alibaba.cloud.ai.studio.admin.generator.model.workflow.nodedata.IterationNodeData;
  import com.alibaba.cloud.ai.studio.admin.generator.service.dsl.DSLDialectType;
  import com.alibaba.cloud.ai.studio.admin.generator.service.generator.workflow.NodeSection;
  
  import org.springframework.stereotype.Component;
  
  import java.util.List;
  
  /**
   * @author vlsmb
   * @since 2025/7/23
   */
  // TODO: 支持并行模式、错误处理，支持Studio的默认输入值，支持Studio的多输入/多输出
  @Component
  public class IterationNodeSection implements NodeSection<IterationNodeData> {
  
  	@Override
  	public boolean support(NodeType nodeType) {
  		return NodeType.ITERATION.equals(nodeType);
  	}
  
  	@Override
  	public String render(Node node, String varName) {
  		// 迭代节点在转换为Workflow的时候已经拆分为多个节点，故本方法返回空
  		return "";
  	}
  
  	@Override
  	public String renderEdges(IterationNodeData nodeData, List<Edge> edges) {
  		return "";
  	}
  
  	@Override
  	public List<String> getImports() {
  		return List.of("java.util.ArrayList", "java.util.Arrays");
  	}
  
  	// 规定迭代节点的start为iterationVarName_start，end为iterationVarName_end
  
  	@Component
  	public static class IterationStartNodeSection implements NodeSection<IterationNodeData> {
  
  		@Override
  		public boolean support(NodeType nodeType) {
  			return NodeType.ITERATION_START.equals(nodeType);
  		}
  
  		@Override
  		public String render(Node node, String varName) {
  			IterationNodeData nodeData = ((IterationNodeData) node.getData());
  			return String.format("""
  					// Iteration [%s] Start Node
  					stateGraph.addNode("%s", AsyncNodeAction.node_async(
  					    createIterationStartAction("%s", "%s", "%s", "%s", "%s", %d)
  					));
  
  					""", node.getId(), varName, nodeData.getInputSelector().getNameInCode(),
  					nodeData.getSourceVarName() + "_state", nodeData.getItemKey(),
  					nodeData.getSourceVarName() + "_index", nodeData.getSourceVarName() + "_isFinished",
  					nodeData.getIndexOffset());
  		}
  
  		// TODO: 添加辅助节点以支持迭代起始节点并行
  		@Override
  		public String renderEdges(IterationNodeData nodeData, List<Edge> edges) {
  			Edge edge = edges.get(0);
  			return String.format("""
  					// Iteration [%s] Start Edge
  					stateGraph.addConditionalEdges("%s", AsyncEdgeAction.edge_async(
  					                state -> {
  					                    Boolean b = state.value("%s", false);
  					                    return b ? "end" : "iteration";
  					                }
  					        ), Map.of("end", "%s", "iteration", "%s"));
  
  					""", nodeData.getSourceVarName(), nodeData.getSourceVarName() + "_start",
  					nodeData.getSourceVarName() + "_isFinished", nodeData.getSourceVarName() + "_end",
  					edge.getTarget());
  		}
  
  		@Override
  		public String assistMethodCode(DSLDialectType dialectType) {
  			return """
  					private NodeAction createIterationStartAction(
  					        String arrayKey, String stateKey,
  					        String itemKey, String indexKey, String flagKey,
  					        int indexOffset) {
  					    return state -> {
  					        Object arrayObj = state.value(arrayKey).orElse(List.of());
  					        List<Integer> stateList = state.value(stateKey, List.class).orElse(null);
  
  					        List<?> arrayList;
  					        if (stateList == null) {
  					            // the first time in iteration
  					            if (arrayObj instanceof List<?>) {
  					                arrayList = new ArrayList<>((List<?>) arrayObj);
  					            } else if (arrayObj.getClass().isArray()) {
  					                arrayList = new ArrayList<>(Arrays.stream((Object[])arrayObj).toList());
  					            } else {
  					                throw new IllegalStateException("value {" + arrayKey + "} is not an array!");
  					            }
  					            int len = arrayList.size();
  					            stateList = new ArrayList<>();
  					            for (int i = 0; i < len; i++) {
  					                stateList.add(i);
  					            }
  					        } else {
  					            arrayList = (List<?>) arrayObj;
  					        }
  
  					        if(stateList.isEmpty()) {
  					            return Map.of(flagKey, true);
  					        }
  					        int index = stateList.get(0);
  					        Object item = arrayList.get(index);
  					        stateList.remove(0);
  					        return Map.of(arrayKey, arrayList, stateKey, stateList, itemKey, item,
  					                indexKey, index + indexOffset, flagKey, false);
  					    };
  					}
  					""";
  		}
  
  		@Override
  		public List<String> getImports() {
  			return List.of("java.util.ArrayList", "java.util.Arrays");
  		}
  
  	}
  
  	@Component
  	public static class IterationEndNodeSection implements NodeSection<IterationNodeData> {
  
  		@Override
  		public boolean support(NodeType nodeType) {
  			return NodeType.ITERATION_END.equals(nodeType);
  		}
  
  		@Override
  		public String render(Node node, String varName) {
  			IterationNodeData nodeData = ((IterationNodeData) node.getData());
  			return String.format("""
  					// Iteration [%s] End Node
  					stateGraph.addNode("%s", AsyncNodeAction.node_async(
  					    createIterationEndAction("%s", "%s", "%s")
  					));
  
  					""", nodeData.getSourceVarName(), varName, nodeData.getSourceVarName() + "_isFinished",
  					nodeData.getResultSelector().getNameInCode(), nodeData.getOutputKey());
  		}
  
  		// TODO: 添加辅助节点以支持迭代终止节点并行
  		@Override
  		public String renderEdges(IterationNodeData nodeData, List<Edge> edges) {
  			Edge edge = edges.get(0);
  			return String.format("""
  					// Iteration [%s] End Edge
  					stateGraph.addConditionalEdges("%s", AsyncEdgeAction.edge_async(
  					                state -> {
  					                    Boolean b = state.value("%s", false);
  					                    return b ? "finish" : "start";
  					                }
  					        ), Map.of("finish", "%s", "start", "%s"));
  
  					""", nodeData.getSourceVarName(), nodeData.getVarName(),
  					nodeData.getSourceVarName() + "_isFinished", edge.getTarget(),
  					nodeData.getSourceVarName() + "_start");
  		}
  
  		@Override
  		public String assistMethodCode(DSLDialectType dialectType) {
  			return """
  					 private NodeAction createIterationEndAction(String flagKey, String resultKey, String outputKey) {
  					     return state -> {
  					         boolean flag = state.value(flagKey, Boolean.class).orElse(true);
  					         List<Object> outputList = state.value(outputKey, List.class).orElse(new ArrayList<>());
  					         if(flag) {
  					             return Map.of(outputKey, outputList);
  					         }
  					         outputList.add(state.value(resultKey).orElseThrow());
  					         return Map.of(outputKey, outputList);
  					     };
  					 }
  					""";
  		}
  
  		@Override
  		public List<String> getImports() {
  			return List.of("java.util.ArrayList", "java.util.Arrays");
  		}
  
  	}
  
  }
  