import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { ModelMetadata, SessionSnapshot, TranscriptProgress } from "@earendil-works/pi-protocol";
import { afterEach, describe, expect, test } from "vitest";
import type { CreateSessionOptions, PiServer, PiSessionRuntime } from "../src/index.ts";
import {
	connectUnixTestClient,
	Deferred,
	type ProtocolTestClient,
	TEST_MODEL,
	TEST_TOKEN,
	TestSessionBackend,
} from "../src/testing/index.ts";
import { createUnixServer, type UnixServerOptions } from "../src/transports/unix/index.ts";

const TOKEN = TEST_TOKEN;
const MODEL = TEST_MODEL;
type Client = ProtocolTestClient;
class MemoryBackend extends TestSessionBackend {}

class OrderedSnapshotBackend extends MemoryBackend {
	readonly firstStarted = new Deferred<void>();
	readonly secondStarted = new Deferred<void>();
	readonly firstRelease = new Deferred<void>();
	readonly secondRelease = new Deferred<void>();
	controlled = false;
	startedCount = 0;

	override async listModels(): Promise<ModelMetadata[]> {
		if (!this.controlled) return super.listModels();
		this.startedCount += 1;
		if (this.startedCount === 1) {
			this.firstStarted.resolve(undefined);
			await this.firstRelease.promise;
		} else if (this.startedCount === 2) {
			this.secondStarted.resolve(undefined);
			await this.secondRelease.promise;
		}
		return [MODEL];
	}
}

const servers = new Set<PiServer>();
const clients = new Set<Client>();
const tempDirectories = new Set<string>();

async function startServer(backend = new MemoryBackend(), options: Partial<UnixServerOptions> = {}) {
	const directory = await mkdtemp(join(tmpdir(), "pst-"));
	tempDirectories.add(directory);
	const server = createUnixServer(backend, {
		token: TOKEN,
		path: join(directory, "server.sock"),
		...options,
	});
	servers.add(server);
	await server.start();
	return { server, backend };
}

async function connect(server: PiServer): Promise<Client> {
	const client = await connectUnixTestClient(server.addresses[0]!);
	clients.add(client);
	return client;
}

async function attach(client: Client, sessionId: string): Promise<SessionSnapshot> {
	const response = await client.request({ command: "attach", sessionId });
	if (!response.ok || response.result.command !== "attach") throw new Error("Attach failed");
	return response.result.session;
}

afterEach(async () => {
	await Promise.all([...clients].map((client) => client.close()));
	clients.clear();
	await Promise.all([...servers].map((server) => server.close()));
	servers.clear();
	await Promise.all([...tempDirectories].map((directory) => rm(directory, { recursive: true, force: true })));
	tempDirectories.clear();
});

describe("PiServer Unix integration", () => {
	test("serializes server snapshot revisions", async () => {
		const backend = new OrderedSnapshotBackend();
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		backend.controlled = true;
		const messageIndex = client.messages.length;

		const firstCreate = client.request({ command: "create", name: "first" });
		await backend.firstStarted.promise;
		const secondCreate = client.request({ command: "create", name: "second" });
		await Promise.resolve();
		expect(backend.startedCount).toBe(1);

		backend.firstRelease.resolve(undefined);
		await backend.secondStarted.promise;
		backend.secondRelease.resolve(undefined);
		await Promise.all([firstCreate, secondCreate]);
		await client.nextFrom(
			messageIndex,
			(message) =>
				message.type === "event" &&
				message.event.type === "server_snapshot" &&
				message.event.snapshot.revision === 2,
		);

		const revisions = client.messages
			.slice(messageIndex)
			.flatMap((message) =>
				message.type === "event" && message.event.type === "server_snapshot"
					? [message.event.snapshot.revision]
					: [],
			);
		expect(revisions).toEqual([1, 2]);
	});

	test("creates server-assigned durable IDs and supports list, attach, and detach", async () => {
		const { server, backend } = await startServer();
		const client = await connect(server);
		await client.hello();
		const created = await client.request({ command: "create", cwd: "/work", name: "Created" });
		expect(created.ok).toBe(true);
		if (!created.ok || created.result.command !== "create") throw new Error("Create failed");
		const createdId = created.result.session.id;
		expect(created.result.session).toMatchObject({
			id: backend.lastCreatedId,
			cwd: "/work",
			name: "Created",
			attached: true,
			locked: true,
		});

		const listed = await client.request({ command: "list" });
		expect(listed).toMatchObject({
			ok: true,
			result: { command: "list", sessions: [{ id: backend.lastCreatedId, attached: true, locked: true }] },
		});
		const detached = await client.request({ command: "detach", sessionId: createdId });
		expect(detached).toMatchObject({ ok: true, result: { command: "detach", sessionId: createdId } });
		expect(backend.latestRuntime(createdId).disposeCount).toBe(1);
		const detachedAgain = await client.request({ command: "detach", sessionId: createdId });
		expect(detachedAgain).toMatchObject({ ok: true, result: { command: "detach", sessionId: createdId } });

		const attached = await attach(client, createdId);
		expect(attached.id).toBe(backend.lastCreatedId);
		expect(backend.runtimes.get(createdId)?.length).toBe(2);
	});

	test("keeps multiple attachments on one connection independent", async () => {
		const backend = new MemoryBackend();
		backend.seed("first");
		backend.seed("second");
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		await attach(client, "first");
		await attach(client, "second");

		await client.request({ command: "detach", sessionId: "first" });
		expect(backend.latestRuntime("first").disposeCount).toBe(1);
		expect(backend.latestRuntime("second").disposeCount).toBe(0);
		const response = await client.request({
			command: "set_thinking",
			sessionId: "second",
			thinkingLevel: "medium",
		});
		expect(response).toMatchObject({ ok: true, result: { session: { id: "second", thinkingLevel: "medium" } } });
	});

	test("broadcasts full snapshots and progress only to clients attached to that session", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server } = await startServer(backend);
		const attachedClient = await connect(server);
		const unattachedClient = await connect(server);
		await attachedClient.hello();
		await unattachedClient.hello();
		await attach(attachedClient, "session-1");
		const runtime = backend.latestRuntime("session-1");
		const progress: TranscriptProgress = {
			type: "assistant_delta",
			messageId: "assistant-1",
			contentIndex: 0,
			kind: "text",
			delta: "hello",
		};
		runtime.emitProgress(progress);
		const progressMessage = await attachedClient.next(
			(message) => message.type === "event" && message.event.type === "session_progress",
		);
		expect(progressMessage).toEqual({
			type: "event",
			event: { type: "session_progress", sessionId: "session-1", progress },
		});
		expect(
			unattachedClient.messages.some(
				(message) => message.type === "event" && message.event.type === "session_progress",
			),
		).toBe(false);

		const messageCount = attachedClient.messages.length;
		runtime.emitSnapshot();
		const snapshotMessage = await attachedClient.nextFrom(
			messageCount,
			(message) =>
				message.type === "event" &&
				message.event.type === "session_snapshot" &&
				message.event.snapshot.revision === runtime.snapshot().revision,
		);
		expect(snapshotMessage).toMatchObject({
			type: "event",
			event: { type: "session_snapshot", snapshot: { id: "session-1", attached: true, locked: true } },
		});
		expect(
			unattachedClient.messages.some(
				(message) => message.type === "event" && message.event.type === "session_snapshot",
			),
		).toBe(false);
	});

	test("allows every attached client to control a singleton live runtime", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server } = await startServer(backend);
		const first = await connect(server);
		const second = await connect(server);
		await first.hello();
		await second.hello();
		await attach(first, "session-1");
		const secondList = await second.request({ command: "list" });
		expect(secondList).toMatchObject({
			ok: true,
			result: { sessions: [{ id: "session-1", attached: false, locked: true }] },
		});
		await attach(second, "session-1");
		expect(backend.runtimes.get("session-1")?.length).toBe(1);

		const modelResponse = await second.request({
			command: "set_model",
			sessionId: "session-1",
			model: { provider: "test", id: "large" },
		});
		expect(modelResponse).toMatchObject({ ok: true, result: { session: { model: { id: "large" } } } });
		await first.next(
			(message) =>
				message.type === "event" &&
				message.event.type === "session_snapshot" &&
				message.event.snapshot.model.id === "large",
		);
		const thinkingResponse = await first.request({
			command: "set_thinking",
			sessionId: "session-1",
			thinkingLevel: "high",
		});
		expect(thinkingResponse).toMatchObject({ ok: true, result: { session: { thinkingLevel: "high" } } });
	});

	test("does not queue prompts and processes steer and abort while a prompt response is pending", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		await attach(client, "session-1");

		const prompt = client.request({ command: "prompt", sessionId: "session-1", text: "first" });
		await client.next(
			(message) =>
				message.type === "event" &&
				message.event.type === "session_snapshot" &&
				message.event.snapshot.phase === "turn",
		);
		const busy = await client.request({ command: "prompt", sessionId: "session-1", text: "second" });
		expect(busy).toMatchObject({ ok: false, error: { code: "busy" } });

		const steer = await client.request({ command: "steer", sessionId: "session-1", text: "adjust" });
		expect(steer).toMatchObject({ ok: true, result: { command: "steer" } });
		expect(backend.latestRuntime("session-1").steers).toEqual([{ text: "adjust" }]);
		const abort = await client.request({ command: "abort", sessionId: "session-1" });
		expect(abort).toMatchObject({ ok: true, result: { command: "abort" } });
		expect(await prompt).toMatchObject({ ok: true, result: { command: "prompt", session: { phase: "idle" } } });
	});

	test("returns operation attachment state relative to the requesting connection", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server } = await startServer(backend);
		const first = await connect(server);
		const second = await connect(server);
		await first.hello();
		await second.hello();
		await attach(first, "session-1");
		await attach(second, "session-1");

		const prompt = first.request({ command: "prompt", sessionId: "session-1", text: "hello" });
		await first.next(
			(message) =>
				message.type === "event" &&
				message.event.type === "session_snapshot" &&
				message.event.snapshot.phase === "turn",
		);
		await first.request({ command: "detach", sessionId: "session-1" });
		backend.latestRuntime("session-1").finishPrompt();

		expect(await prompt).toMatchObject({
			ok: true,
			result: { command: "prompt", session: { id: "session-1", attached: false } },
		});
	});

	test("keeps busy work alive after disconnect and disposes when it next becomes idle", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		await attach(client, "session-1");
		const prompt = client.request({ command: "prompt", sessionId: "session-1", text: "survive" });
		await client.next(
			(message) =>
				message.type === "event" &&
				message.event.type === "session_snapshot" &&
				message.event.snapshot.phase === "turn",
		);
		const runtime = backend.latestRuntime("session-1");
		await client.close();
		await expect(prompt).rejects.toThrow(Error);
		expect(runtime.disposeCount).toBe(0);
		runtime.finishPrompt();
		await runtime.disposed.promise;
		expect(runtime.disposeCount).toBe(1);

		const reconnect = await connect(server);
		await reconnect.hello();
		const snapshot = await attach(reconnect, "session-1");
		expect(snapshot.transcript).toHaveLength(2);
		expect(snapshot.transcript[1]).toMatchObject({ role: "assistant", content: [{ text: "reply:survive" }] });
	});

	test("restores persisted sessions lazily after a server restart", async () => {
		const backend = new MemoryBackend();
		backend.seed();
		const { server: firstServer } = await startServer(backend);
		const firstClient = await connect(firstServer);
		await firstClient.hello();
		await attach(firstClient, "session-1");
		await firstClient.request({ command: "set_thinking", sessionId: "session-1", thinkingLevel: "high" });
		await firstClient.close();
		await firstServer.close();

		const { server: secondServer } = await startServer(backend);
		expect(backend.runtimes.get("session-1")).toHaveLength(1);
		const secondClient = await connect(secondServer);
		await secondClient.hello();
		const restored = await attach(secondClient, "session-1");
		expect(restored.thinkingLevel).toBe("high");
		expect(backend.runtimes.get("session-1")).toHaveLength(2);
	});

	test("rejects and disposes a backend runtime with the wrong server-assigned ID", async () => {
		class WrongIdBackend extends MemoryBackend {
			override async createSession(options: CreateSessionOptions): Promise<PiSessionRuntime> {
				return super.createSession({ ...options, id: "wrong-id" });
			}
		}
		const backend = new WrongIdBackend();
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		const response = await client.request({ command: "create" });
		expect(response).toMatchObject({ ok: false, error: { code: "invalid_request" } });
		expect(backend.latestRuntime("wrong-id").disposeCount).toBe(1);
	});

	test("maps backend lock errors and rejects control from unattached clients", async () => {
		const backend = new MemoryBackend();
		backend.seed("locked");
		backend.locked.add("locked");
		const { server } = await startServer(backend);
		const client = await connect(server);
		await client.hello();
		const locked = await client.request({ command: "attach", sessionId: "locked" });
		expect(locked).toMatchObject({ ok: false, error: { code: "session_locked" } });
		const unattached = await client.request({ command: "abort", sessionId: "locked" });
		expect(unattached).toMatchObject({ ok: false, error: { code: "invalid_request" } });
	});
});