import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import type { DaemonCommand, DaemonResponse } from "../src/modes/daemon/daemon-protocol.js";
import { DaemonSupervisor } from "../src/modes/daemon/daemon-supervisor.js";

type ClientFixture = {
	id: string;
	socket: { destroy: ReturnType<typeof vi.fn> };
	attachedActiveSessionIds: Set<string>;
	capabilities: Set<string>;
};

type WorkerFixture = {
	descriptor: { workerId: string };
	client?: { close: ReturnType<typeof vi.fn> };
};

type PauseEntry = { owner: ClientFixture; worker: WorkerFixture; pauseId: string };

type SupervisorInternals = {
	clients: Set<ClientFixture>;
	workers: Map<string, WorkerFixture>;
	connectionIds: WeakMap<ClientFixture, string>;
	protocolClientIds: WeakMap<ClientFixture, string>;
	sessionInputPauseEpochs: WeakMap<ClientFixture, number>;
	detachingInputPauseSessions: WeakMap<ClientFixture, Set<string>>;
	sessionInputPauses: Map<string, PauseEntry>;
	findWorkerForClient: ReturnType<typeof vi.fn>;
	forwardToWorker: ReturnType<typeof vi.fn>;
	attachClient: ReturnType<typeof vi.fn>;
	reserveSnapshotStream: ReturnType<typeof vi.fn>;
	write: ReturnType<typeof vi.fn>;
	detachClient: ReturnType<typeof vi.fn>;
	handleCommand(client: ClientFixture, command: DaemonCommand): Promise<DaemonResponse | undefined>;
	invalidateWorkerSessionInputPauses(worker: WorkerFixture, reason: string): void;
	handleWorkerClose(worker: WorkerFixture, client: { close: ReturnType<typeof vi.fn> }, error: Error): Promise<void>;
};

const tempDirs: string[] = [];

afterEach(() => {
	for (const directory of tempDirs.splice(0)) rmSync(directory, { recursive: true, force: true });
});

function createHarness() {
	const directory = mkdtempSync(join(tmpdir(), "prime-supervisor-input-pause-"));
	tempDirs.push(directory);
	const supervisor = new DaemonSupervisor(join(directory, "daemon.sock"), {
		defaultSessionConfig: { agentDir: directory, cwd: directory },
		descriptorDir: join(directory, "workers"),
	}) as unknown as SupervisorInternals;
	const worker: WorkerFixture = { descriptor: { workerId: "worker-1" }, client: { close: vi.fn() } };
	supervisor.workers.set(worker.descriptor.workerId, worker);
	supervisor.findWorkerForClient = vi.fn(async () => ({
		worker,
		summary: { id: "active-1", activeSessionId: "active-1" },
	}));
	let pauseSequence = 0;
	supervisor.forwardToWorker = vi.fn(async (_worker: WorkerFixture, command: DaemonCommand) => {
		if (command.type === "acquire_session_input_pause") {
			pauseSequence++;
			return {
				type: "response",
				command: command.type,
				success: true,
				data: { pauseId: `pause-${pauseSequence}` },
			} satisfies DaemonResponse;
		}
		return { type: "response", command: command.type, success: true } satisfies DaemonResponse;
	});
	supervisor.attachClient = vi.fn(async () => ({ worker, result: { activeSessionId: "active-1" } }));
	supervisor.reserveSnapshotStream = vi.fn(() => vi.fn());
	supervisor.write = vi.fn();
	supervisor.detachClient = vi.fn((client: ClientFixture, activeSessionId?: string) => {
		if (activeSessionId) client.attachedActiveSessionIds.delete(activeSessionId);
		else client.attachedActiveSessionIds.clear();
	});
	supervisor.handleWorkerClose = vi.fn(async (closedWorker: WorkerFixture, workerClient, error: Error) => {
		if (closedWorker.client !== workerClient) return;
		closedWorker.client = undefined;
		supervisor.invalidateWorkerSessionInputPauses(closedWorker, error.message);
	});
	return { supervisor, worker };
}

function addClient(supervisor: SupervisorInternals, socketId: string, protocolId: string): ClientFixture {
	const client = {
		id: protocolId,
		socket: { destroy: vi.fn() },
		attachedActiveSessionIds: new Set(["active-1"]),
		capabilities: new Set<string>(),
	};
	supervisor.clients.add(client);
	supervisor.connectionIds.set(client, socketId);
	supervisor.protocolClientIds.set(client, protocolId);
	supervisor.sessionInputPauseEpochs.set(client, 0);
	supervisor.detachingInputPauseSessions.set(client, new Set());
	return client;
}

function acquireCommand(id: string): DaemonCommand {
	return { id, type: "acquire_session_input_pause", activeSessionId: "active-1", leaseKey: "shared-key" };
}

describe("daemon supervisor session input pause ownership", () => {
	it("keeps identical public lease keys independent across client sockets", async () => {
		const { supervisor } = createHarness();
		const first = addClient(supervisor, "socket-a", "protocol-a");
		const second = addClient(supervisor, "socket-b", "protocol-b");

		const firstResponse = await supervisor.handleCommand(first, acquireCommand("acquire-a"));
		const secondResponse = await supervisor.handleCommand(second, acquireCommand("acquire-b"));
		expect(firstResponse).toMatchObject({ success: true, data: { pauseId: "pause-1" } });
		expect(secondResponse).toMatchObject({ success: true, data: { pauseId: "pause-2" } });
		expect(
			supervisor.forwardToWorker.mock.calls
				.map((call) => call[1] as DaemonCommand)
				.filter((command) => command.type === "acquire_session_input_pause")
				.map((command) => command.leaseKey),
		).toEqual([
			JSON.stringify(["socket-a", "protocol-a", "shared-key"]),
			JSON.stringify(["socket-b", "protocol-b", "shared-key"]),
		]);

		await supervisor.handleCommand(first, {
			id: "release-a",
			type: "release_session_input_pause",
			activeSessionId: "active-1",
			pauseId: "pause-1",
		});
		expect(supervisor.sessionInputPauses.has("pause-1")).toBe(false);
		expect(supervisor.sessionInputPauses.has("pause-2")).toBe(true);
	});

	it("clears stable and live detach selectors after reattach", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		const detachingSessions = supervisor.detachingInputPauseSessions.get(client)!;
		detachingSessions.add("stable-id");

		await supervisor.handleCommand(client, {
			id: "reattach",
			type: "reattach",
			activeSessionId: "old-active",
			targetActiveSessionId: "stable-id",
		});

		expect(detachingSessions.has("stable-id")).toBe(false);
		await expect(
			supervisor.handleCommand(client, {
				id: "acquire",
				type: "acquire_session_input_pause",
				activeSessionId: "stable-id",
				leaseKey: "shared-key",
			}),
		).resolves.toMatchObject({ success: true, data: { pauseId: "pause-1" } });
	});

	it("reacquires a fresh worker pause while the prior pause is releasing", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		await supervisor.handleCommand(client, acquireCommand("acquire-1"));
		let finishRelease!: () => void;
		const releaseGate = new Promise<void>((resolve) => {
			finishRelease = resolve;
		});
		supervisor.forwardToWorker.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => {
			await releaseGate;
			return { type: "response", command: command.type, success: true } satisfies DaemonResponse;
		});

		const release = supervisor.handleCommand(client, {
			id: "release-1",
			type: "release_session_input_pause",
			activeSessionId: "active-1",
			pauseId: "pause-1",
		});
		const reacquired = await supervisor.handleCommand(client, acquireCommand("acquire-2"));
		expect(reacquired).toMatchObject({ success: true, data: { pauseId: "pause-2" } });

		finishRelease();
		await release;
		expect(supervisor.sessionInputPauses.has("pause-1")).toBe(false);
		expect(supervisor.sessionInputPauses.has("pause-2")).toBe(true);
	});

	it("deduplicates retries and releases the client lease on detach", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");

		const first = await supervisor.handleCommand(client, acquireCommand("acquire-1"));
		const retry = await supervisor.handleCommand(client, acquireCommand("acquire-2"));
		expect(first).toMatchObject({ success: true, data: { pauseId: "pause-1" } });
		expect(retry).toMatchObject({ success: true, data: { pauseId: "pause-1" } });
		expect(
			supervisor.forwardToWorker.mock.calls.filter(
				(call) => (call[1] as DaemonCommand).type === "acquire_session_input_pause",
			),
		).toHaveLength(1);

		await supervisor.handleCommand(client, { id: "detach", type: "detach", activeSessionId: "active-1" });
		expect(supervisor.sessionInputPauses.size).toBe(0);
		expect(supervisor.detachClient).toHaveBeenCalledWith(client, "active-1");
		const releaseCalls = supervisor.forwardToWorker.mock.calls.filter(
			(call) => (call[1] as DaemonCommand).type === "release_session_input_pause",
		);
		expect(releaseCalls).toHaveLength(1);
		expect(releaseCalls[0]?.[2]).toBe(5_000);
	});

	it("blocks new acquisitions before detach waits for cleanup", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		await supervisor.handleCommand(client, acquireCommand("acquire"));
		let releaseCleanup!: () => void;
		const cleanupGate = new Promise<void>((resolve) => {
			releaseCleanup = resolve;
		});
		supervisor.forwardToWorker.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => {
			await cleanupGate;
			return { type: "response", command: command.type, success: true } as DaemonResponse;
		});

		const detach = supervisor.handleCommand(client, {
			id: "detach",
			type: "detach",
			activeSessionId: "active-1",
		});
		await vi.waitFor(() => expect(supervisor.detachClient).toHaveBeenCalledWith(client, "active-1"));
		await expect(supervisor.handleCommand(client, acquireCommand("late-acquire"))).rejects.toThrow(
			"Session is detaching",
		);
		releaseCleanup();
		await expect(detach).resolves.toMatchObject({ success: true });
		expect(supervisor.sessionInputPauses.size).toBe(0);
	});

	it("fails closed for every owner when detach cannot release the worker pause", async () => {
		const { supervisor, worker } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		const other = addClient(supervisor, "socket-b", "protocol-b");
		await supervisor.handleCommand(client, acquireCommand("acquire-a"));
		await supervisor.handleCommand(other, acquireCommand("acquire-b"));
		supervisor.forwardToWorker.mockRejectedValueOnce(new Error("worker unavailable"));

		await expect(
			supervisor.handleCommand(client, { id: "detach", type: "detach", activeSessionId: "active-1" }),
		).rejects.toThrow("worker unavailable");
		expect(supervisor.sessionInputPauses.size).toBe(0);
		expect(supervisor.handleWorkerClose).toHaveBeenCalledTimes(1);
		expect(worker.client).toBeUndefined();
		expect(client.socket.destroy).toHaveBeenCalled();
		expect(other.socket.destroy).toHaveBeenCalled();
	});

	it("releases all client pauses when detach omits a session id", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		await supervisor.handleCommand(client, acquireCommand("acquire"));

		await supervisor.handleCommand(client, { id: "detach-all", type: "detach" });

		expect(supervisor.sessionInputPauses.size).toBe(0);
		expect(client.attachedActiveSessionIds.size).toBe(0);
		expect(supervisor.detachClient).toHaveBeenCalledWith(client, undefined);
	});

	it("rejects release from a different public client", async () => {
		const { supervisor } = createHarness();
		const owner = addClient(supervisor, "socket-a", "protocol-a");
		const other = addClient(supervisor, "socket-b", "protocol-b");
		await supervisor.handleCommand(owner, acquireCommand("acquire"));

		await expect(
			supervisor.handleCommand(other, {
				id: "release",
				type: "release_session_input_pause",
				activeSessionId: "active-1",
				pauseId: "pause-1",
			}),
		).rejects.toThrow("owned by another client");
	});

	it("invalidates every public owner when the worker connection closes", async () => {
		const { supervisor, worker } = createHarness();
		const first = addClient(supervisor, "socket-a", "protocol-a");
		const second = addClient(supervisor, "socket-b", "protocol-b");
		await supervisor.handleCommand(first, acquireCommand("acquire-a"));
		await supervisor.handleCommand(second, acquireCommand("acquire-b"));

		supervisor.invalidateWorkerSessionInputPauses(worker, "worker pause invalidated");

		expect(supervisor.sessionInputPauses.size).toBe(0);
		expect(first.socket.destroy).toHaveBeenCalledWith(
			expect.objectContaining({ message: "worker pause invalidated" }),
		);
		expect(second.socket.destroy).toHaveBeenCalledWith(
			expect.objectContaining({ message: "worker pause invalidated" }),
		);
	});

	it("releases an acquire that completes after its socket is invalidated", async () => {
		const { supervisor } = createHarness();
		const client = addClient(supervisor, "socket-a", "protocol-a");
		let releaseAcquire!: () => void;
		const gate = new Promise<void>((resolve) => {
			releaseAcquire = resolve;
		});
		const forward = supervisor.forwardToWorker;
		forward.mockImplementationOnce(async (_worker: WorkerFixture, command: DaemonCommand) => {
			await gate;
			return {
				type: "response",
				command: command.type,
				success: true,
				data: { pauseId: "pause-late" },
			} as DaemonResponse;
		});

		const acquisition = supervisor.handleCommand(client, acquireCommand("acquire"));
		await vi.waitFor(() => expect(forward).toHaveBeenCalledTimes(1));
		supervisor.sessionInputPauseEpochs.set(client, 1);
		supervisor.clients.delete(client);
		releaseAcquire();

		await expect(acquisition).rejects.toThrow("invalidated before completion");
		expect(forward.mock.calls.map((call) => (call[1] as DaemonCommand).type)).toEqual([
			"acquire_session_input_pause",
			"release_session_input_pause",
		]);
		expect(supervisor.sessionInputPauses.size).toBe(0);
	});
});