package ai.intelliswarm.swarmai.swarm;

import ai.intelliswarm.swarmai.agent.Agent;
import ai.intelliswarm.swarmai.base.BaseSwarmTest;
import ai.intelliswarm.swarmai.base.MockChatClientFactory;
import ai.intelliswarm.swarmai.base.TestFixtures;
import ai.intelliswarm.swarmai.event.SwarmEvent;
import ai.intelliswarm.swarmai.memory.InMemoryMemory;
import ai.intelliswarm.swarmai.process.ProcessType;
import ai.intelliswarm.swarmai.task.Task;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicInteger;

import static org.junit.jupiter.api.Assertions.*;

@DisplayName("Swarm Kickoff Tests")
class SwarmKickoffTest extends BaseSwarmTest {

    @Nested
    @DisplayName("kickoff()")
    class KickoffTests {

        @Test
        @DisplayName("returns SwarmOutput with valid config")
        void kickoff_withValidConfig_returnsOutput() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            SwarmOutput output = swarm.kickoff(Map.of());

            assertNotNull(output);
            assertTrue(output.isSuccessful());
        }

        @Test
        @DisplayName("sets status to RUNNING during execution")
        void kickoff_setsStatusToRunning() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = Swarm.builder()
                    .agent(agent)
                    .task(task)
                    .process(ProcessType.SEQUENTIAL)
                    .eventPublisher(mockEventPublisher)
                    .build();

            assertEquals(Swarm.SwarmStatus.READY, swarm.getStatus());

            swarm.kickoff(Map.of());

            // After completion
            assertEquals(Swarm.SwarmStatus.COMPLETED, swarm.getStatus());
        }

        @Test
        @DisplayName("sets status to COMPLETED after success")
        void kickoff_setsStatusToCompleted() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            swarm.kickoff(Map.of());

            assertEquals(Swarm.SwarmStatus.COMPLETED, swarm.getStatus());
        }

        @Test
        @DisplayName("publishes SWARM_STARTED event")
        void kickoff_publishesSwarmStartedEvent() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            swarm.kickoff(Map.of());

            assertEventPublished(SwarmEvent.Type.SWARM_STARTED);
        }

        @Test
        @DisplayName("publishes SWARM_COMPLETED event")
        void kickoff_publishesSwarmCompletedEvent() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            swarm.kickoff(Map.of());

            assertEventPublished(SwarmEvent.Type.SWARM_COMPLETED);
        }

        @Test
        @DisplayName("stores last output")
        void kickoff_storesLastOutput() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            assertNull(swarm.getLastOutput());

            SwarmOutput output = swarm.kickoff(Map.of());

            assertEquals(output, swarm.getLastOutput());
        }

        @Test
        @DisplayName("executes all tasks")
        void kickoff_executesAllTasks() {
            Agent agent = createAgent();
            // Use dependent task chain to avoid duplicate queue issue in orderTasks
            List<Task> tasks = TestFixtures.createDependentTaskChain(3, agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), tasks);

            SwarmOutput output = swarm.kickoff(Map.of());

            assertEquals(3, output.getTaskOutputs().size());
        }

        @Test
        @DisplayName("uses sequential process by default")
        void kickoff_withSequentialProcess_usesSequential() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = Swarm.builder()
                    .agent(agent)
                    .task(task)
                    .process(ProcessType.SEQUENTIAL)
                    .eventPublisher(mockEventPublisher)
                    .build();

            assertEquals(ProcessType.SEQUENTIAL, swarm.getProcessType());

            SwarmOutput output = swarm.kickoff(Map.of());

            assertTrue(output.isSuccessful());
        }

        @Test
        @DisplayName("uses hierarchical process when configured")
        void kickoff_withHierarchicalProcess_usesHierarchical() {
            Agent manager = TestFixtures.createManagerAgent(mockChatClient);
            Agent worker = createAgent();
            Task task = createTask(worker);

            Swarm swarm = Swarm.builder()
                    .agent(manager)
                    .agent(worker)
                    .task(task)
                    .process(ProcessType.HIERARCHICAL)
                    .managerAgent(manager)
                    .eventPublisher(mockEventPublisher)
                    .build();

            assertEquals(ProcessType.HIERARCHICAL, swarm.getProcessType());
        }

        @Test
        @DisplayName("fails fast when health checks are enabled and LLM is unavailable")
        void build_withHealthCheckEnabledAndUnavailableLlm_throws() {
            Agent failingAgent = createAgent(chatClientWithError("LLM unavailable"));
            Task task = createTask(failingAgent);

            IllegalStateException exception = assertThrows(IllegalStateException.class, () ->
                    Swarm.builder()
                            .agent(failingAgent)
                            .task(task)
                            .healthCheckEnabled(true)
                            .build());

            assertTrue(exception.getMessage().contains("LLM health check failed"));
        }

        @Test
        @DisplayName("builds when health checks are enabled and LLM is healthy")
        void build_withHealthCheckEnabledAndHealthyLlm_succeeds() {
            Agent agent = createAgent(chatClientWithResponse("pong"));
            Task task = createTask(agent);

            Swarm swarm = Swarm.builder()
                    .agent(agent)
                    .task(task)
                    .healthCheckEnabled(true)
                    .build();

            assertNotNull(swarm);
        }

        @Test
        @DisplayName("sets FAILED status when process fails")
        void kickoff_whenProcessFails_setsFailedStatus() {
            Agent failingAgent = TestFixtures.createTestAgent(chatClientWithError("Failure"));
            Task task = createTask(failingAgent);
            Swarm swarm = createSwarmWithPublisher(List.of(failingAgent), List.of(task));

            assertThrows(RuntimeException.class, () -> swarm.kickoff(Map.of()));

            assertEquals(Swarm.SwarmStatus.FAILED, swarm.getStatus());
        }

        @Test
        @DisplayName("publishes SWARM_FAILED event when process fails")
        void kickoff_whenProcessFails_publishesFailedEvent() {
            Agent failingAgent = TestFixtures.createTestAgent(chatClientWithError("Failure"));
            Task task = createTask(failingAgent);
            Swarm swarm = createSwarmWithPublisher(List.of(failingAgent), List.of(task));

            assertThrows(RuntimeException.class, () -> swarm.kickoff(Map.of()));

            assertEventPublished(SwarmEvent.Type.SWARM_FAILED);
        }

        @Test
        @DisplayName("throws exception when process fails")
        void kickoff_whenProcessFails_throwsException() {
            Agent failingAgent = TestFixtures.createTestAgent(chatClientWithError("Failure"));
            Task task = createTask(failingAgent);
            Swarm swarm = createSwarmWithPublisher(List.of(failingAgent), List.of(task));

            RuntimeException exception = assertThrows(RuntimeException.class, () ->
                    swarm.kickoff(Map.of()));

            assertTrue(exception.getMessage().contains("Swarm execution failed"));
        }

        @Test
        @DisplayName("re-evaluates evolved process each kickoff")
        void kickoff_reEvaluatesEvolvedProcessEachRun() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            AtomicInteger calls = new AtomicInteger();
            Swarm.setEvolutionAdvisor((configured, taskCount, maxDepth) ->
                    calls.getAndIncrement() == 0 ? ProcessType.PARALLEL : configured);

            try {
                invokeApplyEvolution(swarm);
                assertEquals(ProcessType.PARALLEL, readEvolvedProcessType(swarm));

                invokeApplyEvolution(swarm);
                assertNull(readEvolvedProcessType(swarm),
                        "Expected evolved process override to clear when advisor returns configured type");
            } finally {
                Swarm.setEvolutionAdvisor(null);
            }
        }
    }

    @Nested
    @DisplayName("Memory Operations")
    class MemoryOperationsTests {

        @Test
        @DisplayName("resetMemory clears memory")
        void resetMemory_clearsMemory() {
            InMemoryMemory memory = new InMemoryMemory();
            memory.save("agent-1", "Some memory", null);

            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = Swarm.builder()
                    .agent(agent)
                    .task(task)
                    .memory(memory)
                    .eventPublisher(mockEventPublisher)
                    .build();

            assertEquals(1, memory.size());

            swarm.resetMemory();

            assertEquals(0, memory.size());
        }

        @Test
        @DisplayName("resetMemory publishes MEMORY_RESET event")
        void resetMemory_publishesMemoryResetEvent() {
            InMemoryMemory memory = new InMemoryMemory();
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = Swarm.builder()
                    .agent(agent)
                    .task(task)
                    .memory(memory)
                    .eventPublisher(mockEventPublisher)
                    .build();

            swarm.resetMemory();

            assertEventPublished(SwarmEvent.Type.MEMORY_RESET);
        }
    }

    @Nested
    @DisplayName("kickoffAsync()")
    class KickoffAsyncTests {

        @Test
        @DisplayName("returns CompletableFuture")
        void kickoffAsync_returnsCompletableFuture() {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            CompletableFuture<SwarmOutput> future = swarm.kickoffAsync(Map.of());

            assertNotNull(future);
        }

        @Test
        @DisplayName("completes successfully")
        void kickoffAsync_completesSuccessfully() throws ExecutionException, InterruptedException {
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

            CompletableFuture<SwarmOutput> future = swarm.kickoffAsync(Map.of());
            SwarmOutput output = future.get();

            assertNotNull(output);
            assertTrue(output.isSuccessful());
        }

        @Test
        @DisplayName("propagates errors")
        void kickoffAsync_propagatesErrors() {
            Agent failingAgent = TestFixtures.createTestAgent(chatClientWithError("Async failure"));
            Task task = createTask(failingAgent);
            Swarm swarm = createSwarmWithPublisher(List.of(failingAgent), List.of(task));

            CompletableFuture<SwarmOutput> future = swarm.kickoffAsync(Map.of());

            assertThrows(ExecutionException.class, future::get);
        }
    }

    @Nested
    @DisplayName("kickoffForEach()")
    class KickoffForEachTests {

        @Test
        @DisplayName("executes first input successfully")
        void kickoffForEach_executesFirstInput() {
            // Note: Current implementation reuses tasks which can only be executed once
            // This test verifies at least the first input executes successfully
            Agent agent = createAgent();
            Task task = createTask(agent);
            Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));

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

            List<SwarmOutput> outputs = swarm.kickoffForEach(inputs);

            assertEquals(1, outputs.size());
            assertTrue(outputs.get(0).isSuccessful());
        }

        @Test
        @DisplayName("processes batch with new swarms per item")
        void kickoffForEach_withNewSwarmsPerItem() {
            // To process multiple inputs, create new swarm instances for each
            Agent agent = createAgent();

            List<Map<String, Object>> inputs = List.of(Map.of(), Map.of());
            List<SwarmOutput> outputs = new java.util.ArrayList<>();

            for (Map<String, Object> input : inputs) {
                Task task = createTask(agent);
                Swarm swarm = createSwarmWithPublisher(List.of(agent), List.of(task));
                outputs.add(swarm.kickoff(input));
            }

            assertEquals(2, outputs.size());
            assertTrue(outputs.stream().allMatch(SwarmOutput::isSuccessful));
        }
    }

    @Nested
    @DisplayName("Swarm Builder Validation")
    class SwarmValidationTests {

        @Test
        @DisplayName("throws exception with no agents")
        void swarm_withNoAgents_throwsException() {
            Task task = TestFixtures.createTestTaskWithoutAgent();

            assertThrows(IllegalStateException.class, () ->
                    Swarm.builder()
                            .task(task)
                            .build());
        }

        @Test
        @DisplayName("throws exception with no tasks")
        void swarm_withNoTasks_throwsException() {
            Agent agent = createAgent();

            assertThrows(IllegalStateException.class, () ->
                    Swarm.builder()
                            .agent(agent)
                            .build());
        }

        @Test
        @DisplayName("throws exception for hierarchical without manager")
        void swarm_hierarchicalWithoutManager_throwsException() {
            Agent agent = createAgent();
            Task task = createTask(agent);

            assertThrows(IllegalStateException.class, () ->
                    Swarm.builder()
                            .agent(agent)
                            .task(task)
                            .process(ProcessType.HIERARCHICAL)
                            .build());
        }
    }

    private void invokeApplyEvolution(Swarm swarm) {
        try {
            var method = Swarm.class.getDeclaredMethod("applyEvolution");
            method.setAccessible(true);
            method.invoke(swarm);
        } catch (Exception e) {
            fail("Failed to invoke applyEvolution via reflection: " + e.getMessage());
        }
    }

    private ProcessType readEvolvedProcessType(Swarm swarm) {
        try {
            var field = Swarm.class.getDeclaredField("evolvedProcessType");
            field.setAccessible(true);
            return (ProcessType) field.get(swarm);
        } catch (Exception e) {
            fail("Failed to read evolvedProcessType via reflection: " + e.getMessage());
            return null;
        }
    }
}
