/*
 * 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;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

public class IterationNodeTest {

	private static final Logger log = LoggerFactory.getLogger(IterationNodeTest.class);

	private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

	private String runGraph(String input, NodeAction action) throws Exception {
		// 配置子图：START -> iterator -> END
		KeyStrategyFactory subFactory = () -> {
			Map<String, KeyStrategy> map = new HashMap<>();
			map.put("iterator_item", new ReplaceStrategy());
			map.put("iterator_item_result", new ReplaceStrategy());
			return map;
		};
		StateGraph subGraph = new StateGraph("iteration_graph", subFactory).addNode("iterator", node_async(action))
			.addEdge(StateGraph.START, "iterator")
			.addEdge("iterator", StateGraph.END);

		// 配置主图：START -> generate -> IterationNodeGraph -> END
		KeyStrategyFactory mainFactory = () -> {
			Map<String, KeyStrategy> map = new HashMap<>();
			map.put("input_array", new ReplaceStrategy());
			map.put("input_json_array", new ReplaceStrategy());
			map.put("iterator_item", new ReplaceStrategy());
			map.put("output_start", new ReplaceStrategy());
			map.put("iterator_item_result", new ReplaceStrategy());
			map.put("result", new ReplaceStrategy());
			map.put("output_continue", new ReplaceStrategy());
			map.put("iteration_index", new ReplaceStrategy());
			return map;
		};
		CompiledGraph graph = new StateGraph("main", mainFactory)
			.addNode("generate_array", node_async((OverAllState state) -> Map.of("input_json_array", input)))
			.addNode("iteration_node",
					IterationNode.converter()
						.inputArrayJsonKey("input_json_array")
						.tempIndexKey("iteration_index")
						.outputArrayJsonKey("result")
						.iteratorItemKey("iterator_item")
						.iteratorResultKey("iterator_item_result")
						.subGraph(subGraph)
						.convertToStateGraph())
			.addEdge(StateGraph.START, "generate_array")
			.addEdge("generate_array", "iteration_node")
			.addEdge("iteration_node", StateGraph.END)
			.compile();
		OverAllState state = graph.invoke(Map.of()).orElseThrow();
		return state.value("result").orElseThrow().toString();
	}

	@Test
	@DisplayName("Test Integer Iteration")
	public void testInteger() throws Exception {
		String res = this.runGraph("[1, 2, 3, 4, 5]", (OverAllState state) -> {
			int x = state.value("iterator_item", Integer.class).orElseThrow();
			int y = x * x;
			return Map.of("iterator_item_result", Integer.toString(y));
		});
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, new TypeReference<List<String>>() {
		}), List.of("1", "4", "9", "16", "25"));
	}

	@Test
	@DisplayName("Test String Iteration")
	public void testString() throws Exception {
		String res = this.runGraph("[\"a\", \"aa\", \"aaa\", \"aaaa\", \"aaaaa\"]", (OverAllState state) -> {
			int len = state.value("iterator_item", String.class).orElseThrow().length();
			return Map.of("iterator_item_result", len);
		});
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, new TypeReference<List<Integer>>() {
		}), List.of(1, 2, 3, 4, 5));
	}

	@Test
	@DisplayName("Test POJO Iteration")
	public void testPOJO() throws Exception {
		record Person(String name, int age) {
		}
		String res = this.runGraph("[{\"name\": \"a\", \"age\": 1}, {\"name\": \"b\", \"age\": 2}]",
				(OverAllState state) -> {
					Person p = OBJECT_MAPPER.readValue(
							OBJECT_MAPPER.writeValueAsString(state.value("iterator_item").orElseThrow()), Person.class);
					return Map.of("iterator_item_result", p.name);
				});
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, new TypeReference<List<String>>() {
		}), List.of("a", "b"));
	}

	@Test
	@DisplayName("Test Empty Iteration")
	public void testEmpty() throws Exception {
		String res = this.runGraph("[]", (OverAllState state) -> Map.of());
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, List.class), List.of());
	}

	@Test
	@DisplayName("Test two IterationNodes")
	public void testTwoIterationNodes() throws Exception {
		// 配置子图：START -> iterator -> END
		KeyStrategyFactory subFactory1 = () -> {
			Map<String, KeyStrategy> map = new HashMap<>();
			map.put("iterator_item", new ReplaceStrategy());
			map.put("iterator_item_result", new ReplaceStrategy());
			return map;
		};
		StateGraph subGraph1 = new StateGraph("iteration_graph", subFactory1)
			.addNode("iterator", node_async((OverAllState state) -> {
				int x = state.value("iterator_item", Integer.class).orElseThrow();
				int y = x * x;
				return Map.of("iterator_item_result", Integer.toString(y));
			}))
			.addEdge(StateGraph.START, "iterator")
			.addEdge("iterator", StateGraph.END);

		KeyStrategyFactory subFactory2 = () -> {
			Map<String, KeyStrategy> map = new HashMap<>();
			map.put("iterator_item", new ReplaceStrategy());
			map.put("iterator_item_result", new ReplaceStrategy());
			return map;
		};
		StateGraph subGraph2 = new StateGraph("iteration_graph", subFactory2)
			.addNode("iterator", node_async((OverAllState state) -> {
				int len = state.value("iterator_item", String.class).orElseThrow().length();
				return Map.of("iterator_item_result", len);
			}))
			.addEdge(StateGraph.START, "iterator")
			.addEdge("iterator", StateGraph.END);

		// 配置主图：START -> generate -> IterationNode1 -> IterationNode2 -> END
		KeyStrategyFactory mainFactory = () -> {
			Map<String, KeyStrategy> map = new HashMap<>();
			map.put("input_array", new ReplaceStrategy());
			map.put("input_json_array1", new ReplaceStrategy());
			map.put("input_json_array2", new ReplaceStrategy());
			map.put("iterator_item", new ReplaceStrategy());
			map.put("output_start", new ReplaceStrategy());
			map.put("iterator_item_result", new ReplaceStrategy());
			map.put("result1", new ReplaceStrategy());
			map.put("result2", new ReplaceStrategy());
			map.put("output_continue", new ReplaceStrategy());
			map.put("test_temp_array1", new ReplaceStrategy());
			map.put("test_temp_start1", new ReplaceStrategy());
			map.put("test_temp_end1", new ReplaceStrategy());
			map.put("test_temp_array2", new ReplaceStrategy());
			map.put("test_temp_start2", new ReplaceStrategy());
			map.put("test_temp_end2", new ReplaceStrategy());
			map.put("iteration_index1", new ReplaceStrategy());
			map.put("iteration_index2", new ReplaceStrategy());
			return map;
		};
		CompiledGraph graph = new StateGraph("main", mainFactory)
			.addNode("generate_array", node_async((OverAllState state) -> Map.of("input_json_array1", "[1, 4, 10]")))
			.addNode("iteration_node1",
					IterationNode.converter()
						.inputArrayJsonKey("input_json_array1")
						.tempIndexKey("iteration_index1")
						.outputArrayJsonKey("result1")
						.iteratorItemKey("iterator_item")
						.iteratorResultKey("iterator_item_result")
						.tempArrayKey("test_temp_array1")
						.tempStartFlagKey("test_temp_start1")
						.tempEndFlagKey("test_temp_end1")
						.subGraph(subGraph1)
						.convertToStateGraph())
			.addNode("pass", node_async((OverAllState state) -> {
				return Map.of("input_json_array2", state.value("result1", String.class).orElse("[]"));
			}))
			.addNode("iteration_node2",
					IterationNode.converter()
						.inputArrayJsonKey("input_json_array2")
						.tempIndexKey("iteration_index2")
						.outputArrayJsonKey("result2")
						.iteratorItemKey("iterator_item")
						.iteratorResultKey("iterator_item_result")
						.tempArrayKey("test_temp_array2")
						.tempStartFlagKey("test_temp_start2")
						.tempEndFlagKey("test_temp_end2")
						.subGraph(subGraph2)
						.convertToStateGraph())
			.addEdge(StateGraph.START, "generate_array")
			.addEdge("generate_array", "iteration_node1")
			.addEdge("iteration_node1", "pass")
			.addEdge("pass", "iteration_node2")
			.addEdge("iteration_node2", StateGraph.END)
			.compile();
		OverAllState state = graph.invoke(Map.of()).orElseThrow();
		String res = state.value("result2").orElseThrow().toString();
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, new TypeReference<List<Integer>>() {
		}), List.of(1, 2, 3));
	}

	@Test
	@DisplayName("Test Converter.appendToStateGraph Method")
	public void testIterationAttach() throws Exception {
		StateGraph stateGraph = new StateGraph("graph",
				() -> Map.of("input_json_array", new ReplaceStrategy(), "item", new ReplaceStrategy(), "item_result",
						new ReplaceStrategy(), "result", new ReplaceStrategy(), "tv1", new ReplaceStrategy(), "tv2",
						new ReplaceStrategy(), "tv3", new ReplaceStrategy(), "tv4", new ReplaceStrategy()))
			.addNode("generate", node_async((OverAllState state) -> Map.of("input_json_array", "[1, 2, 3, 4, 5]")))
			.addNode("apply", node_async((OverAllState state) -> {
				int x = state.value("item", Integer.class).orElseThrow();
				return Map.of("item_result", x * x * x);
			}));
		// 构造迭代节点
		IterationNode.<Integer, Integer>converter()
			.subGraphStartNodeName("apply")
			.subGraphEndNodeName("apply")
			.tempArrayKey("tv1")
			.tempStartFlagKey("tv2")
			.tempEndFlagKey("tv3")
			.tempIndexKey("tv4")
			.iteratorItemKey("item")
			.iteratorResultKey("item_result")
			.inputArrayJsonKey("input_json_array")
			.outputArrayJsonKey("result")
			.appendToStateGraph(stateGraph, "iteration", "iteration_out");
		stateGraph.addNode("print", node_async((OverAllState state) -> {
			System.out.println(state.value("result", String.class).orElseThrow());
			return Map.of();
		}))
			.addEdge(StateGraph.START, "generate")
			.addEdge("generate", "iteration")
			.addEdge("iteration_out", "print")
			.addEdge("print", StateGraph.END);
		CompiledGraph compiledGraph = stateGraph.compile();
		log.info(compiledGraph.getGraph(GraphRepresentation.Type.PLANTUML, "workflow").content());
		String res = compiledGraph.invoke(Map.of()).orElseThrow().value("result", String.class).orElseThrow();
		log.info("result: {}", res);
		Assertions.assertEquals(OBJECT_MAPPER.readValue(res, new TypeReference<List<Integer>>() {
		}), List.of(1, 8, 27, 64, 125));
	}

}
