import { beforeEach, describe, expect, it, vi } from "vitest";
import { ObjectId } from "mongodb";
import type { Workspace } from "$lib/types/Workspace";

const mocks = vi.hoisted(() => {
	class Rpc {
		static instances: Rpc[] = [];
		commands: Array<Record<string, unknown>> = [];
		lastSequence = 0;

		constructor() {
			Rpc.instances.push(this);
		}
		async connect() {}
		subscribe() {
			return () => undefined;
		}
		subscribeReplayGap() {
			return () => undefined;
		}
		async command(command: Record<string, unknown>) {
			this.commands.push(command);
			if (command.type === "get_state") {
				return { type: "response", success: true, data: { isStreaming: false } };
			}
			return { type: "response", success: true, data: {} };
		}
		close() {}
	}
	return { Rpc };
});

vi.mock("./PiRpcClient", () => ({ PiRpcClient: mocks.Rpc }));
vi.mock("./SandboxClient", () => ({
	SandboxClient: {
		connect: () => ({
			proxyHeaders: {},
			proxyUrlFor: () => "ws://bridge.test/ws",
		}),
	},
}));

import { AgentSessionManager } from "./AgentSessionManager";

beforeEach(() => {
	mocks.Rpc.instances.length = 0;
});

describe("AgentSessionManager dispatch fencing", () => {
	it("never sends prompt or abort when a new-turn signal was already aborted", async () => {
		const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" });
		const controller = new AbortController();
		controller.abort();

		await expect(
			manager
				.runBatches("do work", {
					signal: controller.signal,
					commandId: "prompt-a",
				})
				.next()
		).rejects.toMatchObject({ name: "AbortError" });
		expect(mocks.Rpc.instances[0]?.commands).toEqual([]);
	});

	it("runs the last ownership assertion before prompt and sends nothing when it fails", async () => {
		const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" });
		const assertion = vi.fn(async () => {
			throw new Error("lease lost");
		});

		await expect(
			manager
				.runBatches("do work", {
					commandId: "prompt-a",
					beforePromptDispatch: assertion,
				})
				.next()
		).rejects.toThrow("lease lost");
		expect(assertion).toHaveBeenCalledOnce();
		expect(mocks.Rpc.instances[0]?.commands).toEqual([]);
	});

	it("fences a resume abort and carries the expected prompt id", async () => {
		const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" });
		const controller = new AbortController();
		controller.abort();
		const assertion = vi.fn(async () => undefined);

		await manager
			.resumeBatches({
				signal: controller.signal,
				expectedPromptId: "prompt-a",
				beforeAbortDispatch: assertion,
			})
			.next();
		await vi.waitFor(() => expect(assertion).toHaveBeenCalledOnce());
		expect(mocks.Rpc.instances[0]?.commands).toContainEqual({
			type: "abort",
			_chatUiExpectedPromptId: "prompt-a",
		});
	});

	it("detaching a stale relay never aborts pi", async () => {
		const manager = AgentSessionManager.connect({ workspace: workspace(), hfToken: "test" });
		const detached = new AbortController();
		detached.abort();

		await expect(
			manager.resumeBatches({ detachSignal: detached.signal, expectedPromptId: "prompt-a" }).next()
		).rejects.toMatchObject({ name: "AgentSessionDetachedError" });
		expect(mocks.Rpc.instances[0]?.commands).toEqual([]);
	});
});

function workspace(): Workspace {
	const now = new Date();
	return {
		_id: new ObjectId(),
		conversationId: new ObjectId(),
		sessionId: "owner",
		repo: { kind: "blank" },
		flavor: "cpu-basic",
		model: "test/model",
		thinkingLevel: "medium",
		policy: "ask",
		sandbox: {
			status: "ready",
			flavor: "cpu-basic",
			id: "sandbox",
			generation: "incarnation",
			namespace: "owner",
			baseUrl: "https://sandbox.hf.jobs",
			nonce: "nonce",
			image: "node:22",
			createdAt: now,
		},
		spend: { accruedUsd: 0 },
		createdAt: now,
		updatedAt: now,
	};
}