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

import com.alibaba.cloud.ai.graph.action.EdgeAction;
import com.alibaba.cloud.ai.graph.action.InterruptionMetadata;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.checkpoint.Checkpoint;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.checkpoint.savers.VersionedMemorySaver;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.logging.LogManager;

import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;
import static java.lang.String.format;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertIterableEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrowsExactly;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class StateGraphMemorySaverTest {

	private static final org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(StateGraphMemorySaverTest.class);

	NodeAction agent_whether = state -> {
		List<String> messages = (List<String>) state.value("messages").get();
		String lastMessage = messages.get(messages.size() - 1);

		if (lastMessage.contains("temperature")) {
			return Map.of("messages", "whether in Naples is sunny");
		}
		if (lastMessage.contains("whether")) {
			return Map.of("messages", "tool_calls");
		}
		if (lastMessage.contains("bartolo")) {
			return Map.of("messages", "Hi bartolo, nice to meet you too! How can I assist you today?");
		}
		if (messages.stream().anyMatch(m -> m.contains("bartolo"))) {
			return Map.of("messages", "Hi, bartolo welcome back?");
		}
		throw new IllegalStateException("unknown message!");
	};

	// Simulate LLM agent
	NodeAction tool_whether = state -> Map.of("messages", "temperature in Napoli is 30 degree");

	EdgeAction shouldContinue_whether = state -> {
		List<String> messages = (List<String>) state.value("messages").get();
		return messages.get(messages.size() - 1).equals("tool_calls") ? "tools" : END;
	};

	private KeyStrategyFactory keyStrategyFactory = () -> {
		Map<String, KeyStrategy> keyStrategyMap = new HashMap<>();
		keyStrategyMap.put("agent_1:prop1", new ReplaceStrategy());
		keyStrategyMap.put("messages", new AppendStrategy());
		keyStrategyMap.put("steps", new ReplaceStrategy());
		return keyStrategyMap;
	};

	@BeforeAll
	public static void initLogging() throws IOException {
		try (var is = StateGraphMemorySaverTest.class.getResourceAsStream("/logging.properties")) {
			LogManager.getLogManager().readConfiguration(is);
		}
	}

	@Test
	public void testCheckpointInitialState() throws Exception {
		NodeAction agent_1 = state -> {
			log.info("agent_1");
			return Map.of("agent_1:prop1", "agent_1:test");
		};

		var workflow = new StateGraph(keyStrategyFactory).addNode("agent_1", node_async(agent_1))
			.addEdge(START, "agent_1")
			.addEdge("agent_1", END);

		var saver = MemorySaver.builder().build();

		var compileConfig = CompileConfig.builder()
			.saverConfig(SaverConfig.builder()
				.register(saver)
				.build())
			.build();

		var runnableConfig = RunnableConfig.builder().build();
		var app = workflow.compile(compileConfig);

		Map<String, Object> inputs = Map.of("input", "test1");

		var initState = app.cloneState(app.getInitialState(inputs, runnableConfig));

		assertEquals(1, initState.data().size());
		assertTrue(initState.value("input").isPresent());
		assertEquals("test1", initState.value("input").get());

		//
		// Test checkpoint not override inputs
		//
		var newState = new OverAllState(Map.of("input", "test2"));

		saver.put(runnableConfig,
				Checkpoint.builder().state(newState.data()).nodeId(START).nextNodeId("agent_1").build());

		app = workflow.compile(compileConfig);
		initState = app.cloneState(app.getInitialState(inputs, runnableConfig));

		assertEquals(1, initState.data().size());
		assertTrue(initState.value("input").isPresent());
		assertEquals("test1", initState.value("input").get());

		// Test checkpoints are saved
		newState = new OverAllState(Map.of("input", "test2", "agent_1:prop1", "agent_1:test"));
		saver.put(runnableConfig, Checkpoint.builder().state(newState).nodeId("agent_1").nextNodeId(END).build());
		app = workflow.compile(compileConfig);
		initState = app.cloneState(app.getInitialState(inputs, runnableConfig));

		assertEquals(2, initState.data().size());
		assertTrue(initState.value("input").isPresent());
		assertEquals("test1", initState.value("input").get());
		assertTrue(initState.value("agent_1:prop1").isPresent());
		assertEquals("agent_1:test", initState.value("agent_1:prop1").get());

		var checkpoints = saver.list(runnableConfig);
		assertEquals(2, checkpoints.size());
		Optional<Checkpoint> last = saver.get(runnableConfig);
		assertTrue(last.isPresent());
		assertEquals("agent_1", last.get().getNodeId());
		assertNotNull(last.get().getState().get("agent_1:prop1"));
		assertEquals("agent_1:test", last.get().getState().get("agent_1:prop1"));

		var tag = saver.release(runnableConfig);

		assertIterableEquals(checkpoints, tag.checkpoints());

		var checkpointsAfterTag = saver.list(runnableConfig);
		assertTrue(checkpointsAfterTag.isEmpty());
	}

	@Test
	public void testCheckpointNumRetained() throws Exception {
		var saver = MemorySaver.builder().build();
		var config = RunnableConfig.builder()
				.threadId("retained-thread")
				.checkpointsNumRetained(2)
				.build();

		var first = Checkpoint.builder().id("cp1").state(Map.of("value", "1")).nodeId("node1").nextNodeId(END).build();
		var second = Checkpoint.builder().id("cp2").state(Map.of("value", "2")).nodeId("node2").nextNodeId(END).build();
		var third = Checkpoint.builder().id("cp3").state(Map.of("value", "3")).nodeId("node3").nextNodeId(END).build();

		saver.put(config, first);
		saver.put(config, second);
		saver.put(config, third);

		var checkpoints = List.copyOf(saver.list(config));
		assertEquals(2, checkpoints.size());
		assertEquals("cp3", checkpoints.get(0).getId());
		assertEquals("cp2", checkpoints.get(1).getId());
		assertTrue(saver.get(RunnableConfig.builder(config).checkPointId("cp1").build()).isEmpty());
	}

	@Test
	public void testCheckpointSaverResubmit() throws Exception {
		int expectedSteps = 5;

		NodeAction agent_1 = state -> {
			Integer steps = (Integer) state.value("steps").get();
			steps = steps + 1;
			log.info("agent_1: step: {}", steps);
			return Map.of("steps", steps, "messages", format("agent_1:step %d", steps));
		};

		EdgeAction shouldContinue = state -> {
			Integer steps = (Integer) state.value("steps").get();
			if (steps >= expectedSteps) {
				return "exit";
			}
			return "next";
		};

		var workflow = new StateGraph(keyStrategyFactory).addEdge(START, "agent_1")
			.addNode("agent_1", node_async(agent_1))
			.addConditionalEdges("agent_1", edge_async(shouldContinue), Map.of("next", "agent_1", "exit", END));

		var saver = VersionedMemorySaver.builder().build();

		var compileConfig = CompileConfig.builder()
			.saverConfig(SaverConfig.builder()
				.register(saver)
				.build())
			.build();

		var app = workflow.compile(compileConfig);

		Map<String, Object> inputs = Map.of("steps", 0);

		var runnableConfig = RunnableConfig.builder().threadId("thread_1").build();

		var state = app.invoke(inputs, runnableConfig);

		assertTrue(state.isPresent());
		assertEquals(expectedSteps, state.get().value("steps").get());

		List<String> messages = (List) state.get().value("messages").get();
		assertFalse(messages.isEmpty());

		log.info("{}", messages);

		assertEquals(expectedSteps, messages.size());
		for (int i = 0; i < messages.size(); i++) {
			assertEquals(format("agent_1:step %d", i + 1), messages.get(i));
		}

		assertTrue(saver.lastVersionByThreadId(runnableConfig).isEmpty());

		var snapshot = app.getState(runnableConfig);

		assertNotNull(snapshot);
		log.info("SNAPSHOT:\n{}\n", snapshot);

		// SUBMIT NEW THREAD 2
		runnableConfig = RunnableConfig.builder().threadId("thread_2").build();

		state = app.invoke(inputs, runnableConfig);

		assertTrue(state.isPresent());
		assertEquals(expectedSteps, state.get().value("steps").get());
		messages = (List) state.get().value("messages").get();

		log.info("{}", messages);

		assertEquals(expectedSteps, messages.size());
		for (int i = 0; i < messages.size(); i++) {
			assertEquals(format("agent_1:step %d", i + 1), messages.get(i));
		}

		// RE-SUBMIT THREAD 1
		state = app.invoke(Map.of(), runnableConfig);

		assertTrue(state.isPresent());
		assertEquals(expectedSteps + 1, state.get().value("steps").get());
		messages = (List) state.get().value("messages").get();

		log.info("{}", messages);

		assertEquals(expectedSteps + 1, messages.size());

	}

	@Test
	public void testViewAndUpdatePastGraphState() throws Exception {

		var workflow = new StateGraph(keyStrategyFactory).addNode("agent", node_async(agent_whether))
			.addNode("tools", node_async(tool_whether))
			.addEdge(START, "agent")
			.addConditionalEdges("agent", edge_async(shouldContinue_whether), Map.of("tools", "tools", END, END))
			.addEdge("tools", "agent");

		var saver = MemorySaver.builder().build();

		var compileConfig = CompileConfig.builder()
			.saverConfig(SaverConfig.builder()
				.register(saver)
				.build())
			.build();

		var app = workflow.compile(compileConfig);

		Map<String, Object> inputs = Map.of("messages", "whether in Naples?");

		var runnableConfig = RunnableConfig.builder().threadId("thread_1").build();

		var results = app.streamSnapshots(inputs, runnableConfig).collectList().block();

		results.forEach(r -> log.info("{}: Node: {} - {}", r.getClass().getSimpleName(), r.node(),
				r.state().value("messages").get()));

		assertEquals(5, results.size());
		assertInstanceOf(NodeOutput.class, results.get(0));
		assertInstanceOf(StateSnapshot.class, results.get(1));
		assertInstanceOf(StateSnapshot.class, results.get(2));
		assertInstanceOf(StateSnapshot.class, results.get(3));
		assertInstanceOf(NodeOutput.class, results.get(4));

		var snapshot = app.getState(runnableConfig);
		assertNotNull(snapshot);
		assertEquals(END, snapshot.next());

		log.info("LAST SNAPSHOT:\n{}\n", snapshot);

		var stateHistory = app.getStateHistory(runnableConfig);
		stateHistory.forEach(state -> log.info("SNAPSHOT HISTORY:\n{}\n", state));
		assertNotNull(stateHistory);
		assertEquals(4, stateHistory.size());

		for (StateSnapshot s : stateHistory) {
			log.info("SNAPSHOT HISTORY:\n{}\n", s);
		}

		RunnableConfig resumeConfig = RunnableConfig.builder().threadId("thread_1").resume().build();
		results = app.stream(null, resumeConfig).collectList().block();

		assertNotNull(results);
		assertFalse(results.isEmpty());
		assertEquals(1, results.size());
		assertTrue(results.get(0).state().value("messages").isPresent());
		List<String> messages = (List<String>) results.get(0).state().value("messages").get();
		assertEquals("whether in Naples is sunny", messages.get(messages.size() - 1));

		Optional<StateSnapshot> firstSnapshot = stateHistory.stream().reduce((first, second) -> second); // take
		// the
		// last
		assertTrue(firstSnapshot.isPresent());
		assertTrue(firstSnapshot.get().state().value("messages").isPresent());
		assertEquals("whether in Naples?", ((List<String>) firstSnapshot.get().state().value("messages").get()).get(0));

		var toReplay = firstSnapshot.get().config();

		toReplay = app.updateState(toReplay, Map.of("messages", "i'm bartolo"));
		RunnableConfig toReplayResumeConfig = toReplay.withResume();
		results = app.stream(null, toReplayResumeConfig).collectList().block();

		assertNotNull(results);
		assertFalse(results.isEmpty());
		assertEquals(2, results.size());
		assertEquals(END, results.get(1).node());
		assertTrue(results.get(1).state().value("messages").isPresent());
		messages = (List<String>) results.get(0).state().value("messages").get();
		assertEquals("Hi bartolo, nice to meet you too! How can I assist you today?",
				messages.get(messages.size() - 1));

	}

	@Test
	public void testPauseAndUpdatePastGraphState() throws Exception {

		var workflow = new StateGraph(keyStrategyFactory).addNode("agent", node_async(agent_whether))
			.addNode("tools", node_async(tool_whether))
			.addEdge(START, "agent")
			.addConditionalEdges("agent", edge_async(shouldContinue_whether), Map.of("tools", "tools", END, END))
			.addEdge("tools", "agent");

		var saver = MemorySaver.builder().build();

		var compileConfig = CompileConfig.builder()
			.saverConfig(SaverConfig.builder().register(saver).build())
			.interruptBefore("tools")
			.build();

		var app = workflow.compile(compileConfig);

		var runnableConfig = RunnableConfig.builder().threadId("thread_1").build();

		log.info("FIRST CALL WITH INTERRUPT BEFORE 'tools'");
		Map<String, Object> inputs = Map.of("messages", "whether in Naples?");
		var results = app.stream(inputs, runnableConfig)
			.doOnNext(n -> log.info("{}", n)).collectList().block();
		assertNotNull(results);
		assertEquals(3, results.size());
		assertEquals(START, results.get(0).node());
		assertEquals("agent", results.get(1).node());
		assertInstanceOf(InterruptionMetadata.class, results.get(2));
		List<String> messages = results.get(1).state().value("messages", List.class).get();
		messages.get(messages.size() - 1);
		assertTrue(messages.get(messages.size() - 1) != null);

		var state = app.getState(runnableConfig);

		assertNotNull(state);
		assertEquals("tools", state.next());

		log.info("RESUME CALL");
		RunnableConfig resumeConfig = runnableConfig.withResume();
		results = app.stream(null, resumeConfig).doOnNext(n -> log.info("{}", n)).collectList().block();

		assertNotNull(results);
		assertEquals(3, results.size());
		assertEquals("tools", results.get(0).node());
		assertEquals("agent", results.get(1).node());
		assertEquals(END, results.get(2).node());
		List<String> messages2 = results.get(2).state().value("messages", List.class).get();
		assertTrue(messages2.get(messages.size() - 1) != null);
		assertEquals("whether in Naples is sunny", messages2.get(messages2.size() - 1));

	}

	@Test
	public void testMemoryWithVersionsSaver() throws Exception {

		var threadId = "thread_1";

		var saver = VersionedMemorySaver.builder().build();

		// Check for error
		var configWithVersion = RunnableConfig.builder().threadId(threadId).build();

		// Create a new version of thread_1
		var configWithoutVersion = RunnableConfig.builder().threadId(threadId).build();

		var checkpoint = Checkpoint.builder().state(Map.of()).nodeId(START).nextNodeId(END).build();

		var newConfig = saver.put(configWithoutVersion, checkpoint);

		var list = saver.list(newConfig);

		assertEquals(1, list.size());

		var tag = saver.release(newConfig);

		assertEquals(1, tag.checkpoints().size());

		var versions = saver.versionsByThreadId(threadId);

		assertEquals(1, versions.size());

		// Check if checkpoints collection is immutable
		assertThrowsExactly(UnsupportedOperationException.class, () -> tag.checkpoints().remove(checkpoint));

		var configWithVersion1 = RunnableConfig.builder().threadId(threadId).build();

		assertEquals(1, tag.checkpoints().size());

		versions = saver.versionsByThreadId(configWithVersion);

		assertEquals(1, versions.size());
		assertEquals(checkpoint.getId(), list.stream().findFirst().map(Checkpoint::getId).orElseThrow());

		var checkpoint_1 = Checkpoint.builder().state(Map.of()).nodeId("test").nextNodeId(END).build();
		var checkpoint_2 = Checkpoint.builder().state(Map.of()).nodeId("test_1").nextNodeId(END).build();

		configWithVersion1 = saver.put(configWithVersion1, checkpoint_1);

		configWithVersion1 = saver.put(configWithVersion1.withCheckPointId(null), checkpoint_2);

		versions = saver.versionsByThreadId(threadId);

		assertEquals(1, versions.size());

		var tag2 = saver.release(configWithVersion1);

		assertEquals(2, tag2.checkpoints().size());

	}

}
