package com.cameleer.server.core.agent;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;

import java.time.Instant;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

import static org.assertj.core.api.Assertions.assertThat;

class AgentRegistryServiceTest {

    // Thresholds: 90s stale, 300s dead, 60s command expiry
    private AgentRegistryService registry;

    @BeforeEach
    void setUp() {
        registry = new AgentRegistryService(90_000, 300_000, 60_000);
    }

    @Nested
    class Registration {

        @Test
        void registerNewAgent_createsWithLiveState() {
            AgentInfo agent = registry.register("agent-1", "agent-1", "order-svc", "default",
                    "1.0.0", List.of("route1", "route2"), Map.of("feature", "tracing"));

            assertThat(agent).isNotNull();
            assertThat(agent.instanceId()).isEqualTo("agent-1");
            assertThat(agent.displayName()).isEqualTo("agent-1");
            assertThat(agent.applicationId()).isEqualTo("order-svc");
            assertThat(agent.version()).isEqualTo("1.0.0");
            assertThat(agent.routeIds()).containsExactly("route1", "route2");
            assertThat(agent.capabilities()).containsEntry("feature", "tracing");
            assertThat(agent.state()).isEqualTo(AgentState.LIVE);
            assertThat(agent.registeredAt()).isNotNull();
            assertThat(agent.lastHeartbeat()).isNotNull();
            assertThat(agent.staleTransitionTime()).isNull();
        }

        @Test
        void reRegisterSameId_updatesMetadataAndTransitionsToLive() {
            registry.register("agent-1", "agent-1", "old-group", "default",
                    "1.0.0", List.of("route1"), Map.of());

            AgentInfo updated = registry.register("agent-1", "agent-1", "new-group", "default",
                    "2.0.0", List.of("route1", "route2"), Map.of("new", "cap"));

            assertThat(updated.instanceId()).isEqualTo("agent-1");
            assertThat(updated.displayName()).isEqualTo("agent-1");
            assertThat(updated.applicationId()).isEqualTo("new-group");
            assertThat(updated.version()).isEqualTo("2.0.0");
            assertThat(updated.routeIds()).containsExactly("route1", "route2");
            assertThat(updated.capabilities()).containsEntry("new", "cap");
            assertThat(updated.state()).isEqualTo(AgentState.LIVE);
            assertThat(updated.staleTransitionTime()).isNull();
        }

        @Test
        void reRegisterSameId_updatesRegisteredAtAndLastHeartbeat() {
            AgentInfo first = registry.register("agent-1", "agent-1", "group", "default",
                    "1.0.0", List.of(), Map.of());
            Instant firstRegisteredAt = first.registeredAt();

            AgentInfo second = registry.register("agent-1", "agent-1", "group", "default",
                    "1.0.0", List.of(), Map.of());

            assertThat(second.registeredAt()).isAfterOrEqualTo(firstRegisteredAt);
            assertThat(second.lastHeartbeat()).isAfterOrEqualTo(first.lastHeartbeat());
        }
    }

    @Nested
    class Heartbeat {

        @Test
        void heartbeatKnownAgent_returnsTrue() {
            registry.register("agent-1", "agent-1", "group", "default", "1.0.0", List.of(), Map.of());

            boolean result = registry.heartbeat("agent-1");

            assertThat(result).isTrue();
        }

        @Test
        void heartbeatKnownAgent_updatesLastHeartbeat() {
            registry.register("agent-1", "agent-1", "group", "default", "1.0.0", List.of(), Map.of());
            Instant before = registry.findById("agent-1").lastHeartbeat();

            registry.heartbeat("agent-1");

            AgentInfo after = registry.findById("agent-1");
            assertThat(after.lastHeartbeat()).isAfterOrEqualTo(before);
        }

        @Test
        void heartbeatUnknownAgent_returnsFalse() {
            boolean result = registry.heartbeat("unknown-agent");

            assertThat(result).isFalse();
        }

        @Test
        void heartbeatStaleAgent_transitionsToLive() {
            registry.register("agent-1", "agent-1", "group", "default", "1.0.0", List.of(), Map.of());
            registry.transitionState("agent-1", AgentState.STALE);

            assertThat(registry.findById("agent-1").state()).isEqualTo(AgentState.STALE);

            registry.heartbeat("agent-1");

            assertThat(registry.findById("agent-1").state()).isEqualTo(AgentState.LIVE);
            assertThat(registry.findById("agent-1").staleTransitionTime()).isNull();
        }

        @Test
        void heartbeatDeadAgent_transitionsToLive() {
            // A paused docker container can exceed both the stale and dead thresholds
            // before resuming. When heartbeats arrive again the registry must revive
            // the agent — otherwise the UI shows DEAD forever despite an alive agent.
            registry.register("agent-1", "agent-1", "group", "default", "1.0.0", List.of(), Map.of());
            registry.transitionState("agent-1", AgentState.STALE);
            registry.transitionState("agent-1", AgentState.DEAD);

            assertThat(registry.findById("agent-1").state()).isEqualTo(AgentState.DEAD);

            registry.heartbeat("agent-1");

            assertThat(registry.findById("agent-1").state()).isEqualTo(AgentState.LIVE);
            assertThat(registry.findById("agent-1").staleTransitionTime()).isNull();
        }
    }

    @Nested
    class LifecycleTransitions {

        @Test
        void liveAgentBeyondStaleThreshold_transitionsToStale() {
            // Use very short thresholds for test
            AgentRegistryService shortRegistry = new AgentRegistryService(1, 300_000, 60_000);
            shortRegistry.register("agent-1", "Name", "group", "default", "1.0.0", List.of(), Map.of());

            // Wait briefly to exceed 1ms threshold
            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }

            shortRegistry.checkLifecycle();

            AgentInfo agent = shortRegistry.findById("agent-1");
            assertThat(agent.state()).isEqualTo(AgentState.STALE);
            assertThat(agent.staleTransitionTime()).isNotNull();
        }

        @Test
        void staleAgentBeyondDeadThreshold_transitionsToDead() {
            // Use very short thresholds for test: 1ms stale, 1ms dead
            AgentRegistryService shortRegistry = new AgentRegistryService(1, 1, 60_000);
            shortRegistry.register("agent-1", "Name", "group", "default", "1.0.0", List.of(), Map.of());

            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
            shortRegistry.checkLifecycle(); // LIVE -> STALE

            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
            shortRegistry.checkLifecycle(); // STALE -> DEAD

            assertThat(shortRegistry.findById("agent-1").state()).isEqualTo(AgentState.DEAD);
        }

        @Test
        void deadAgentRemainsDead() {
            AgentRegistryService shortRegistry = new AgentRegistryService(1, 1, 60_000);
            shortRegistry.register("agent-1", "Name", "group", "default", "1.0.0", List.of(), Map.of());

            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
            shortRegistry.checkLifecycle();
            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }
            shortRegistry.checkLifecycle();

            assertThat(shortRegistry.findById("agent-1").state()).isEqualTo(AgentState.DEAD);

            // Additional lifecycle check should not change state
            shortRegistry.checkLifecycle();
            assertThat(shortRegistry.findById("agent-1").state()).isEqualTo(AgentState.DEAD);
        }

        @Test
        void transitionState_setsStaleTransitionTimeWhenGoingStale() {
            registry.register("agent-1", "agent-1", "group", "default", "1.0.0", List.of(), Map.of());

            registry.transitionState("agent-1", AgentState.STALE);

            AgentInfo agent = registry.findById("agent-1");
            assertThat(agent.state()).isEqualTo(AgentState.STALE);
            assertThat(agent.staleTransitionTime()).isNotNull();
        }
    }

    @Nested
    class Queries {

        @Test
        void findAll_returnsAllAgents() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            registry.register("agent-2", "agent-2", "g", "default", "1.0", List.of(), Map.of());

            List<AgentInfo> all = registry.findAll();

            assertThat(all).hasSize(2);
            assertThat(all).extracting(AgentInfo::instanceId).containsExactlyInAnyOrder("agent-1", "agent-2");
        }

        @Test
        void findByState_filtersCorrectly() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            registry.register("agent-2", "agent-2", "g", "default", "1.0", List.of(), Map.of());
            registry.transitionState("agent-2", AgentState.STALE);

            List<AgentInfo> live = registry.findByState(AgentState.LIVE);
            List<AgentInfo> stale = registry.findByState(AgentState.STALE);

            assertThat(live).hasSize(1).extracting(AgentInfo::instanceId).containsExactly("agent-1");
            assertThat(stale).hasSize(1).extracting(AgentInfo::instanceId).containsExactly("agent-2");
        }

        @Test
        void findById_unknownReturnsNull() {
            AgentInfo result = registry.findById("nonexistent");

            assertThat(result).isNull();
        }

        @Test
        void findById_knownReturnsAgent() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());

            AgentInfo result = registry.findById("agent-1");

            assertThat(result).isNotNull();
            assertThat(result.instanceId()).isEqualTo("agent-1");
        }
    }

    @Nested
    class Commands {

        @Test
        void addCommand_createsPendingCommand() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());

            AgentCommand cmd = registry.addCommand("agent-1", CommandType.CONFIG_UPDATE, "{\"key\":\"val\"}");

            assertThat(cmd).isNotNull();
            assertThat(cmd.id()).isNotNull();
            assertThat(cmd.type()).isEqualTo(CommandType.CONFIG_UPDATE);
            assertThat(cmd.payload()).isEqualTo("{\"key\":\"val\"}");
            assertThat(cmd.targetInstanceId()).isEqualTo("agent-1");
            assertThat(cmd.status()).isEqualTo(CommandStatus.PENDING);
            assertThat(cmd.createdAt()).isNotNull();
        }

        @Test
        void addCommand_notifiesEventListener() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());

            AtomicReference<AgentCommand> received = new AtomicReference<>();
            registry.setEventListener((agentId, command) -> received.set(command));

            AgentCommand cmd = registry.addCommand("agent-1", CommandType.DEEP_TRACE, "{}");

            assertThat(received.get()).isNotNull();
            assertThat(received.get().id()).isEqualTo(cmd.id());
        }

        @Test
        void acknowledgeCommand_transitionsStatus() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            AgentCommand cmd = registry.addCommand("agent-1", CommandType.REPLAY, "{}");

            boolean acked = registry.acknowledgeCommand("agent-1", cmd.id());

            assertThat(acked).isTrue();
        }

        @Test
        void acknowledgeCommand_unknownReturnsFalse() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());

            boolean acked = registry.acknowledgeCommand("agent-1", "nonexistent-cmd");

            assertThat(acked).isFalse();
        }

        @Test
        void findPendingCommands_returnsOnlyPending() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            AgentCommand cmd1 = registry.addCommand("agent-1", CommandType.CONFIG_UPDATE, "{}");
            AgentCommand cmd2 = registry.addCommand("agent-1", CommandType.DEEP_TRACE, "{}");
            registry.acknowledgeCommand("agent-1", cmd1.id());

            List<AgentCommand> pending = registry.findPendingCommands("agent-1");

            assertThat(pending).hasSize(1);
            assertThat(pending.get(0).id()).isEqualTo(cmd2.id());
        }

        @Test
        void markDelivered_updatesStatus() {
            registry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            AgentCommand cmd = registry.addCommand("agent-1", CommandType.CONFIG_UPDATE, "{}");

            registry.markDelivered("agent-1", cmd.id());

            // After marking delivered, the command should no longer be PENDING
            List<AgentCommand> pending = registry.findPendingCommands("agent-1");
            assertThat(pending).isEmpty();
        }

        @Test
        void expireOldCommands_removesExpiredPendingCommands() {
            // Use 1ms expiry for test
            AgentRegistryService shortRegistry = new AgentRegistryService(90_000, 300_000, 1);
            shortRegistry.register("agent-1", "agent-1", "g", "default", "1.0", List.of(), Map.of());
            shortRegistry.addCommand("agent-1", CommandType.CONFIG_UPDATE, "{}");

            try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }

            shortRegistry.expireOldCommands();

            List<AgentCommand> pending = shortRegistry.findPendingCommands("agent-1");
            assertThat(pending).isEmpty();
        }

        @Test
        void findPendingCommands_emptyForUnknownAgent() {
            List<AgentCommand> pending = registry.findPendingCommands("unknown");

            assertThat(pending).isEmpty();
        }
    }
}
