/**
 * Test helpers for stubbing the WebSocket used by `rpcCall`.
 *
 * `rpcCall` opens a transient WebSocket, sends an `rpc_request` frame, and
 * resolves once the matching `rpc_response` arrives.  `installRpcMock`
 * installs a `WebSocket` global that records sent frames and replies with
 * canned responses keyed by RPC method name.
 */

import { vi } from "vitest";

interface MockSocket {
	url: string;
	sent: string[];
	onopen: ((ev: Event) => void) | null;
	onmessage: ((ev: { data: string }) => void) | null;
	onerror: ((ev: Event) => void) | null;
	onclose: ((ev: CloseEvent) => void) | null;
	readyState: number;
	send(data: string): void;
	close(): void;
}

/**
 * Test handle returned by `installMockRpcWebSocket`: records sockets and calls,
 * lets the test register canned responses (or override responses for specific
 * methods), and restores the real `WebSocket` constructor on cleanup.
 */
export interface MockRpcInstaller {
	/** All sockets created during the test, in order. */
	instances: MockSocket[];
	/** All RPC method names that were called. */
	calls: { method: string; params: Record<string, unknown> }[];
	/**
	 * Register a result to be returned for `method`.  If the value is a
	 * function, it is invoked with the request params and its return value
	 * (or the resolved promise) is used as the result.
	 */
	respond(method: string, result: unknown): void;
	/** Make `method` reply with an `ok:false` response. */
	respondError(method: string, error: string, code?: string): void;
}

/**
 * Install a stub `WebSocket` global that synchronously resolves RPC calls
 * with results registered via the returned [`MockRpcInstaller`].
 */
export function installRpcMock(): MockRpcInstaller {
	const instances: MockSocket[] = [];
	const calls: { method: string; params: Record<string, unknown> }[] = [];
	const results = new Map<string, unknown>();
	const errors = new Map<string, { error: string; code?: string }>();

	class MockWebSocket implements MockSocket {
		static readonly CONNECTING = 0;
		static readonly OPEN = 1;
		static readonly CLOSING = 2;
		static readonly CLOSED = 3;

		url: string;
		sent: string[] = [];
		onopen: ((ev: Event) => void) | null = null;
		onmessage: ((ev: { data: string }) => void) | null = null;
		onerror: ((ev: Event) => void) | null = null;
		onclose: ((ev: CloseEvent) => void) | null = null;
		readyState = 0;

		constructor(url: string) {
			this.url = url;
			instances.push(this);
			queueMicrotask(() => {
				this.readyState = 1;
				this.onopen?.(new Event("open"));
			});
		}

		send(data: string) {
			this.sent.push(data);
			let frame: {
				correlation_id?: string;
				method?: string;
				params?: Record<string, unknown>;
			};
			try {
				frame = JSON.parse(data);
			} catch {
				return;
			}
			const { correlation_id, method, params } = frame;
			if (!correlation_id || !method) return;
			calls.push({ method, params: params ?? {} });
			queueMicrotask(() => {
				const err = errors.get(method);
				if (err) {
					this.onmessage?.({
						data: JSON.stringify({
							kind: "rpc_response",
							version: 1,
							correlation_id,
							ok: false,
							error: err.error,
							code: err.code,
						}),
					});
					return;
				}
				if (results.has(method)) {
					this.onmessage?.({
						data: JSON.stringify({
							kind: "rpc_response",
							version: 1,
							correlation_id,
							ok: true,
							result: results.get(method),
						}),
					});
					return;
				}
				// No registered response — synthesise NOT_FOUND so the test fails
				// loudly instead of timing out.
				this.onmessage?.({
					data: JSON.stringify({
						kind: "rpc_response",
						version: 1,
						correlation_id,
						ok: false,
						error: `no mock for ${method}`,
						code: "NOT_FOUND",
					}),
				});
			});
		}

		close() {
			this.readyState = 3;
		}
	}

	vi.stubGlobal("WebSocket", MockWebSocket);

	return {
		instances,
		calls,
		respond(method, result) {
			results.set(method, result);
		},
		respondError(method, error, code) {
			errors.set(method, { error, code });
		},
	};
}
