import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { KillswitchService } from "./killswitch.service";
import { AgentLifecycleService } from "../spawner/agent-lifecycle.service";
import { ValkeyService } from "../valkey/valkey.service";
import { CleanupService } from "./cleanup.service";
import type { AgentState } from "../valkey/types";

describe("KillswitchService", () => {
  let service: KillswitchService;
  let mockLifecycleService: {
    transitionToKilled: ReturnType<typeof vi.fn>;
    getAgentLifecycleState: ReturnType<typeof vi.fn>;
    listAgentLifecycleStates: ReturnType<typeof vi.fn>;
  };
  let mockValkeyService: {
    getAgentState: ReturnType<typeof vi.fn>;
    listAgents: ReturnType<typeof vi.fn>;
  };
  let mockCleanupService: {
    cleanup: ReturnType<typeof vi.fn>;
  };

  const mockAgentState: AgentState = {
    agentId: "agent-123",
    status: "running",
    taskId: "task-456",
    startedAt: new Date().toISOString(),
    metadata: {
      containerId: "container-abc",
      repository: "/path/to/repo",
    },
  };

  beforeEach(() => {
    // Create mocks
    mockLifecycleService = {
      transitionToKilled: vi.fn(),
      getAgentLifecycleState: vi.fn(),
      listAgentLifecycleStates: vi.fn(),
    };

    mockValkeyService = {
      getAgentState: vi.fn(),
      listAgents: vi.fn(),
    };

    mockCleanupService = {
      cleanup: vi.fn(),
    };

    service = new KillswitchService(
      mockLifecycleService as unknown as AgentLifecycleService,
      mockValkeyService as unknown as ValkeyService,
      mockCleanupService as unknown as CleanupService
    );
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe("killAgent", () => {
    it("should kill single agent with full cleanup", async () => {
      // Arrange
      mockValkeyService.getAgentState.mockResolvedValue(mockAgentState);
      mockLifecycleService.transitionToKilled.mockResolvedValue({
        ...mockAgentState,
        status: "killed",
        completedAt: new Date().toISOString(),
      });
      mockCleanupService.cleanup.mockResolvedValue({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });

      // Act
      await service.killAgent("agent-123");

      // Assert
      expect(mockValkeyService.getAgentState).toHaveBeenCalledWith("agent-123");
      expect(mockLifecycleService.transitionToKilled).toHaveBeenCalledWith("agent-123");
      expect(mockCleanupService.cleanup).toHaveBeenCalledWith(mockAgentState);
    });

    it("should throw error if agent not found", async () => {
      // Arrange
      mockValkeyService.getAgentState.mockResolvedValue(null);

      // Act & Assert
      await expect(service.killAgent("agent-999")).rejects.toThrow("Agent agent-999 not found");

      expect(mockLifecycleService.transitionToKilled).not.toHaveBeenCalled();
      expect(mockCleanupService.cleanup).not.toHaveBeenCalled();
    });

    it("should handle agent already in killed state", async () => {
      // Arrange
      const killedState: AgentState = {
        ...mockAgentState,
        status: "killed",
        completedAt: new Date().toISOString(),
      };
      mockValkeyService.getAgentState.mockResolvedValue(killedState);
      mockLifecycleService.transitionToKilled.mockRejectedValue(
        new Error("Invalid state transition from killed to killed")
      );

      // Act & Assert
      await expect(service.killAgent("agent-123")).rejects.toThrow("Invalid state transition");

      // Cleanup should not be attempted
      expect(mockCleanupService.cleanup).not.toHaveBeenCalled();
    });
  });

  describe("killAllAgents", () => {
    it("should kill all running agents", async () => {
      // Arrange
      const agent1: AgentState = {
        ...mockAgentState,
        agentId: "agent-1",
        taskId: "task-1",
        metadata: { containerId: "container-1", repository: "/repo1" },
      };
      const agent2: AgentState = {
        ...mockAgentState,
        agentId: "agent-2",
        taskId: "task-2",
        metadata: { containerId: "container-2", repository: "/repo2" },
      };

      mockValkeyService.listAgents.mockResolvedValue([agent1, agent2]);
      mockValkeyService.getAgentState.mockResolvedValueOnce(agent1).mockResolvedValueOnce(agent2);
      mockLifecycleService.transitionToKilled
        .mockResolvedValueOnce({ ...agent1, status: "killed" })
        .mockResolvedValueOnce({ ...agent2, status: "killed" });
      mockCleanupService.cleanup.mockResolvedValue({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });

      // Act
      const result = await service.killAllAgents();

      // Assert
      expect(mockValkeyService.listAgents).toHaveBeenCalled();
      expect(result.total).toBe(2);
      expect(result.killed).toBe(2);
      expect(result.failed).toBe(0);
      expect(mockLifecycleService.transitionToKilled).toHaveBeenCalledTimes(2);
      expect(mockCleanupService.cleanup).toHaveBeenCalledTimes(2);
    });

    it("should only kill active agents (spawning or running)", async () => {
      // Arrange
      const runningAgent: AgentState = {
        ...mockAgentState,
        agentId: "agent-1",
        status: "running",
        metadata: { containerId: "container-1", repository: "/repo1" },
      };
      const completedAgent: AgentState = {
        ...mockAgentState,
        agentId: "agent-2",
        status: "completed",
      };
      const failedAgent: AgentState = {
        ...mockAgentState,
        agentId: "agent-3",
        status: "failed",
      };

      mockValkeyService.listAgents.mockResolvedValue([runningAgent, completedAgent, failedAgent]);
      mockValkeyService.getAgentState.mockResolvedValueOnce(runningAgent);
      mockLifecycleService.transitionToKilled.mockResolvedValueOnce({
        ...runningAgent,
        status: "killed",
      });
      mockCleanupService.cleanup.mockResolvedValue({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });

      // Act
      const result = await service.killAllAgents();

      // Assert
      expect(result.total).toBe(1);
      expect(result.killed).toBe(1);
      expect(result.failed).toBe(0);
      expect(mockLifecycleService.transitionToKilled).toHaveBeenCalledTimes(1);
      expect(mockLifecycleService.transitionToKilled).toHaveBeenCalledWith("agent-1");
    });

    it("should return zero results when no agents exist", async () => {
      // Arrange
      mockValkeyService.listAgents.mockResolvedValue([]);

      // Act
      const result = await service.killAllAgents();

      // Assert
      expect(result.total).toBe(0);
      expect(result.killed).toBe(0);
      expect(result.failed).toBe(0);
      expect(mockLifecycleService.transitionToKilled).not.toHaveBeenCalled();
    });

    it("should track failures when some agents fail to kill", async () => {
      // Arrange
      const agent1: AgentState = {
        ...mockAgentState,
        agentId: "agent-1",
        taskId: "task-1",
        metadata: { containerId: "container-1", repository: "/repo1" },
      };
      const agent2: AgentState = {
        ...mockAgentState,
        agentId: "agent-2",
        taskId: "task-2",
        metadata: { containerId: "container-2", repository: "/repo2" },
      };

      mockValkeyService.listAgents.mockResolvedValue([agent1, agent2]);
      mockValkeyService.getAgentState.mockResolvedValueOnce(agent1).mockResolvedValueOnce(agent2);
      mockLifecycleService.transitionToKilled
        .mockResolvedValueOnce({ ...agent1, status: "killed" })
        .mockRejectedValueOnce(new Error("State transition failed"));
      mockCleanupService.cleanup.mockResolvedValue({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });

      // Act
      const result = await service.killAllAgents();

      // Assert
      expect(result.total).toBe(2);
      expect(result.killed).toBe(1);
      expect(result.failed).toBe(1);
      expect(result.errors).toHaveLength(1);
      expect(result.errors?.[0]).toContain("agent-2");
    });

    it("should continue killing other agents even if one fails", async () => {
      // Arrange
      const agent1: AgentState = {
        ...mockAgentState,
        agentId: "agent-1",
        taskId: "task-1",
        metadata: { containerId: "container-1", repository: "/repo1" },
      };
      const agent2: AgentState = {
        ...mockAgentState,
        agentId: "agent-2",
        taskId: "task-2",
        metadata: { containerId: "container-2", repository: "/repo2" },
      };
      const agent3: AgentState = {
        ...mockAgentState,
        agentId: "agent-3",
        taskId: "task-3",
        metadata: { containerId: "container-3", repository: "/repo3" },
      };

      mockValkeyService.listAgents.mockResolvedValue([agent1, agent2, agent3]);
      mockValkeyService.getAgentState
        .mockResolvedValueOnce(agent1)
        .mockResolvedValueOnce(agent2)
        .mockResolvedValueOnce(agent3);
      mockLifecycleService.transitionToKilled
        .mockResolvedValueOnce({ ...agent1, status: "killed" })
        .mockRejectedValueOnce(new Error("Failed"))
        .mockResolvedValueOnce({ ...agent3, status: "killed" });
      mockCleanupService.cleanup.mockResolvedValue({
        docker: { success: true },
        worktree: { success: true },
        state: { success: true },
      });

      // Act
      const result = await service.killAllAgents();

      // Assert
      expect(result.total).toBe(3);
      expect(result.killed).toBe(2);
      expect(result.failed).toBe(1);
      expect(mockLifecycleService.transitionToKilled).toHaveBeenCalledTimes(3);
    });
  });
});
