import type { Procedure } from "../procedure.ts";
import { WebSocketClient, type WebSocketLike, type WebSocketOptions } from "../ws.ts";

interface AuthenticatedWebSocketOptions {
	/**
	 * Timeout in milliseconds before a handshake procedure is considered failed.
	 */
	timeoutMs: number;
	/**
	 * When connection is closed because the handshake procedure has failed,
	 * these arguments are used as status code and reason for closure.
	 */
	closeArgs: [code: number, reason: string];
}

/**
 * Creates a `WebSocketClient` and wraps it with
 * {@link makeAuthenticated `makeAuthenticated`}.
 */
export const createAuthenticatedWebSocketClient = <TSend, TRecv = TSend>(
	options: WebSocketOptions<TSend, TRecv> & AuthenticatedWebSocketOptions,
): (<R>(handshake: Procedure<TSend, TRecv, R>) => WebSocketLike<TSend, TRecv, R>) =>
	makeAuthenticated(new WebSocketClient(options), options);

/**
 * Wraps the provided `WebSocketClient` and adds a handshake procedure that
 * executes every time connection is opened. If the handshake procedure fails,
 * the connection is closed.
 *
 * The `authed.open()` method resolves with the result of the handshake
 * procedure.
 */
export const makeAuthenticated =
	<TSend, TRecv = TSend>(
		ws: WebSocketLike<TSend, TRecv>,
		options: AuthenticatedWebSocketOptions,
	) =>
	<R>(handshake: Procedure<TSend, TRecv, R>): WebSocketLike<TSend, TRecv, R> => {
		async function open(): Promise<R> {
			await ws.open();

			let timeout: NodeJS.Timeout | undefined = undefined;

			try {
				if (Number.isFinite(options.timeoutMs)) {
					timeout = setTimeout(() => ws.close(...options.closeArgs), options.timeoutMs);
				}

				return await ws.exec(handshake, {
					suppressMessageEvents: true,
				});
			} catch (err) {
				await ws.close(...options.closeArgs);
				throw err;
			} finally {
				clearTimeout(timeout);
			}
		}

		return {
			open,
			close: ws.close.bind(ws),
			send: ws.send.bind(ws),
			exec: ws.exec.bind(ws),
			on: ws.on.bind(ws),
		};
	};