/*
 * 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.agent.flow;

import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.agent.flow.agent.LlmRoutingAgent;
import com.alibaba.cloud.ai.graph.agent.flow.agent.ParallelAgent;
import com.alibaba.cloud.ai.graph.agent.flow.agent.SequentialAgent;
import com.alibaba.cloud.ai.graph.agent.flow.builder.FlowAgentBuilder;
import com.alibaba.cloud.ai.graph.agent.flow.builder.FlowGraphBuilder;
import com.alibaba.cloud.ai.graph.agent.flow.enums.FlowAgentEnum;
import com.alibaba.cloud.ai.graph.agent.flow.strategy.FlowGraphBuildingStrategy;
import com.alibaba.cloud.ai.graph.agent.flow.strategy.FlowGraphBuildingStrategyRegistry;
import com.alibaba.cloud.ai.graph.agent.flow.strategy.SequentialGraphBuildingStrategy;
import com.alibaba.cloud.ai.graph.agent.utils.ChatClientTestUtils;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.tool.resolution.ToolCallbackResolver;

import java.util.List;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * Unit tests for the refactored FlowAgent architecture demonstrating the improved design
 * patterns. These tests verify that the new architecture provides extensibility,
 * consistency, and type safety.
 */
class FlowAgentArchitectureTest {

	@Mock
	private ChatClient chatClient;

	@Mock
	private ChatModel chatModel;

	@Mock
	private ToolCallbackResolver resolver;

	@BeforeEach
	void setUp() {
		MockitoAnnotations.openMocks(this);
		ChatClientTestUtils.stubChatClientMutateChain(chatClient);
	}

	@Test
	void testSequentialAgentBuilderPattern() throws Exception {
		// Create a sub-agent
		ReactAgent subAgent = createMockReactAgent("dataProcessor", "processed_data");

		// Test that SequentialAgent uses the unified builder pattern
		SequentialAgent agent = SequentialAgent.builder()
			.name("sequentialWorkflow")
			.description("A sequential workflow")
			.subAgents(List.of(subAgent))
			.build();

		// Verify agent properties
		assertNotNull(agent);
		assertEquals("sequentialWorkflow", agent.name());
		assertEquals("A sequential workflow", agent.description());
		assertEquals(1, agent.subAgents().size());
	}

	@Test
	void testLlmRoutingAgentBuilderPattern() throws Exception {
		// Create sub-agents
		ReactAgent agent1 = createMockReactAgent("analysisAgent", "analysis_result");
		ReactAgent agent2 = createMockReactAgent("reportAgent", "report_result");

		// Test that LlmRoutingAgent uses the unified builder pattern with LLM-specific
		// features
		LlmRoutingAgent agent = LlmRoutingAgent.builder()
			.name("intelligentRouter")
			.description("Routes tasks intelligently")
			.subAgents(List.of(agent1, agent2))
			.model(chatModel) // LLM-specific configuration
			.build();

		// Verify agent properties
		assertNotNull(agent);
		assertEquals("intelligentRouter", agent.name());
		assertEquals("Routes tasks intelligently", agent.description());
		assertEquals(2, agent.subAgents().size());
	}

	@Test
	void testParallelAgentBuilderPattern() throws Exception {
		// Create sub-agents
		ReactAgent agent1 = createMockReactAgent("dataAnalyzer", "analysis_result");
		ReactAgent agent2 = createMockReactAgent("dataValidator", "validation_result");
		ReactAgent agent3 = createMockReactAgent("dataCleaner", "cleaning_result");

		// Test that ParallelAgent uses the unified builder pattern with parallel-specific
		// features
		ParallelAgent agent = ParallelAgent.builder()
			.name("dataProcessingPipeline")
			.description("Processes data in parallel")
			.mergeOutputKey("processing_result")
			.subAgents(List.of(agent1, agent2, agent3))
			.mergeStrategy(new ParallelAgent.DefaultMergeStrategy())
			.maxConcurrency(3)
			.build();

		// Verify agent properties
		assertNotNull(agent);
		assertEquals("dataProcessingPipeline", agent.name());
		assertEquals("Processes data in parallel", agent.description());
		assertEquals(3, agent.subAgents().size());
		assertEquals(3, agent.maxConcurrency());
		assertTrue(agent.mergeStrategy() instanceof ParallelAgent.DefaultMergeStrategy);
	}

	@Test
	void testBuilderValidation() {
		// Test that builders properly validate required fields
		assertThrows(IllegalArgumentException.class, () -> {
			SequentialAgent.builder().description("Missing name").build();
		});

		assertThrows(IllegalArgumentException.class, () -> {
			LlmRoutingAgent.builder()
				.name("router")
				.subAgents(List.of(createMockReactAgent("agent", "output")))
				.build(); // Missing ChatModel
		});

		assertThrows(IllegalArgumentException.class, () -> {
			ParallelAgent.builder()
				.name("parallel")
				.subAgents(List.of(createMockReactAgent("agent", "output"))) // Need at
				// least 2
				// agents
				.build();
		});
	}

	@Test
	void testStrategyRegistryExtensibility() {
		// Test that the strategy registry allows for runtime extension
		FlowGraphBuildingStrategyRegistry registry = FlowGraphBuildingStrategyRegistry.getInstance();

		// Verify default strategies are registered
		assertTrue(registry.hasStrategy(FlowAgentEnum.SEQUENTIAL.getType()));
		assertTrue(registry.hasStrategy(FlowAgentEnum.ROUTING.getType()));
		assertTrue(registry.hasStrategy(FlowAgentEnum.PARALLEL.getType()));
		assertTrue(registry.hasStrategy(FlowAgentEnum.CONDITIONAL.getType()));

		// Test getting strategies
		FlowGraphBuildingStrategy sequentialStrategy = registry.getStrategy(FlowAgentEnum.SEQUENTIAL.getType());
		assertNotNull(sequentialStrategy);
		assertTrue(sequentialStrategy instanceof SequentialGraphBuildingStrategy);

		// Test that non-existent strategies throw exceptions
		assertThrows(IllegalArgumentException.class, () -> {
			registry.getStrategy("NON_EXISTENT_STRATEGY");
		});
	}

	@Test
	void testFlowGraphConfigCustomProperties() {
		// Test that FlowGraphConfig supports custom properties for extensibility
		FlowGraphBuilder.FlowGraphConfig config = FlowGraphBuilder.FlowGraphConfig.builder()
			.name("testConfig")
			.customProperty("maxRetries", 3)
			.customProperty("timeout", 5000L)
			.customProperty("enableDebug", true);

		assertEquals(3, config.getCustomProperty("maxRetries"));
		assertEquals(5000L, config.getCustomProperty("timeout"));
		assertEquals(true, config.getCustomProperty("enableDebug"));
		assertNull(config.getCustomProperty("nonExistentProperty"));
	}

	@Test
	void testMergeStrategies() {
		// Test different merge strategies for ParallelAgent
		ParallelAgent.DefaultMergeStrategy defaultStrategy = new ParallelAgent.DefaultMergeStrategy();
		ParallelAgent.ListMergeStrategy listStrategy = new ParallelAgent.ListMergeStrategy();
		ParallelAgent.ConcatenationMergeStrategy concatStrategy = new ParallelAgent.ConcatenationMergeStrategy(" | ");

		assertNotNull(defaultStrategy);
		assertNotNull(listStrategy);
		assertNotNull(concatStrategy);

		// Test strategy types
		assertTrue(defaultStrategy instanceof ParallelAgent.MergeStrategy);
		assertTrue(listStrategy instanceof ParallelAgent.MergeStrategy);
		assertTrue(concatStrategy instanceof ParallelAgent.MergeStrategy);
	}

	@Test
	void testFluentInterface() throws Exception {
		// Test that all builders support fluent interface
		ParallelAgent.ParallelAgentBuilder builder = ParallelAgent.builder();

		// All methods should return the same builder instance for method chaining
		assertSame(builder, builder.name("test"));
		assertSame(builder, builder.description("test description"));
		assertSame(builder, builder.maxConcurrency(5));
		assertSame(builder, builder.mergeStrategy(new ParallelAgent.DefaultMergeStrategy()));
	}

	@Test
	void testBuilderInheritance() {
		// Test that concrete builders properly inherit from FlowAgentBuilder
		assertTrue(SequentialAgent.builder() instanceof FlowAgentBuilder);
		assertTrue(LlmRoutingAgent.builder() instanceof FlowAgentBuilder);
		assertTrue(ParallelAgent.builder() instanceof FlowAgentBuilder);
	}

	private ReactAgent createMockReactAgent(String name, String outputKey) throws Exception {
		return ReactAgent.builder()
			.name(name)
			.description("Mock agent: " + name)
			.outputKey(outputKey)
			.chatClient(chatClient)
			.resolver(resolver)
			.build();
	}

}
