import { AbortedError, UnexpectedMessageError, WebSocketClosedError } from "./errors.ts";
import { promiseWithResolvers } from "./promise.ts";
import { type MessageFormat, recv, send } from "./sendrecv.ts";
import { closeTag, type TaggedWebSocket } from "./taggedws.ts";

export interface Procedure<TSend, TRecv = TSend, R = void> {
	(
		helpers: ProcedureHelpers<TSend, TRecv>,
	): Generator<ProcedureCommand<TSend, TRecv>, R, TRecv>;
}

export interface ProcedureHelpers<TSend, TRecv = TSend> {
	send: (payload: TSend) => SendCommand<TSend>;
	recv: () => RecvCommand<TRecv>;
	expect: (predicate: (message: TRecv) => boolean) => RecvCommand<TRecv>;
	settle: (promise: Promise<unknown>) => SettleCommand;
}

export type ProcedureCommand<TSend, TRecv = TSend> =
	| SendCommand<TSend>
	| RecvCommand<TRecv>
	| SettleCommand;

export interface SendCommand<TSend> {
	action: "send";
	payload: TSend;
}

export interface RecvCommand<TRecv> {
	action: "recv";
	predicate?: (payload: TRecv) => boolean;
	strict: boolean;
}

export interface SettleCommand {
	action: "settle";
	promise: Promise<unknown>;
}

export interface ExecProcedureOptions<TSend, TRecv = TSend, R = void> {
	ws: TaggedWebSocket;
	procedure: Procedure<TSend, TRecv, R>;
	format: MessageFormat<TSend, TRecv>;
}

type ProcedureCont<TRecv> =
	| { type: "next"; args: [] | [payload: TRecv] }
	| { type: "throw"; args: [error: Error] };

export async function execProcedure<TSend, TRecv = TSend, R = void>(
	options: ExecProcedureOptions<TSend, TRecv, R>,
): Promise<R> {
	const steps = options.procedure(getProcedureHelpers());
	let stepResult: IteratorResult<ProcedureCommand<TSend, TRecv>, R>;
	let cont: ProcedureCont<TRecv> = { type: "next", args: [] };

	const next = (cont: ProcedureCont<TRecv>) => {
		switch (cont.type) {
			case "next":
				return steps.next(...cont.args);

			case "throw":
				return steps.throw(...cont.args);
		}
	};

	while (!(stepResult = next(cont)).done) {
		const cmd = stepResult.value;
		const action = cmd.action;

		switch (action) {
			case "send":
				cont = execSend(cmd, options);
				break;

			case "recv":
				cont = await execRecv(cmd, options);
				break;

			case "settle":
				cont = await execSettle(cmd, options);
				break;

			default:
				throw new Error("Procedure yielded unexpected command");
		}
	}

	return stepResult.value;
}

function execSend<TSend, TRecv>(
	cmd: SendCommand<TSend>,
	options: ExecProcedureOptions<TSend, TRecv, unknown>,
): ProcedureCont<TRecv> {
	send(cmd.payload, options);
	return { type: "next", args: [] };
}

function execRecv<TSend, TRecv>(
	cmd: RecvCommand<TRecv>,
	options: ExecProcedureOptions<TSend, TRecv, unknown>,
): Promise<ProcedureCont<TRecv>> {
	const { ws } = options;
	const { promise, resolve, reject } = promiseWithResolvers<ProcedureCont<TRecv>>();
	const ac = new AbortController();

	ws.addEventListener(
		"message",
		(event: MessageEvent) => {
			const result = recv(event, options);

			switch (result.type) {
				case "error":
					return resolve({ type: "throw", args: [result.error] });

				case "message": {
					let ok = true;
					try {
						if (cmd.predicate) {
							ok = cmd.predicate(result.message);
						}
					} catch (err) {
						return reject(err);
					}

					return ok
						? resolve({ type: "next", args: [result.message] })
						: cmd.strict
							? resolve({
									type: "throw",
									args: [new UnexpectedMessageError(result.message)],
								})
							: undefined;
				}
			}
		},
		{ signal: ac.signal },
	);

	rejectOnClose(ws, reject, ac.signal);
	return promise.finally(() => ac.abort());
}

function execSettle<TSend, TRecv>(
	cmd: SettleCommand,
	options: ExecProcedureOptions<TSend, TRecv, unknown>,
): Promise<ProcedureCont<TRecv>> {
	const ac = new AbortController();

	const settle: Promise<ProcedureCont<TRecv>> = cmd.promise.then(
		() => ({ type: "next", args: [] }),
		(error) => ({ type: "throw", args: [error] }),
	);

	const close: Promise<never> = new Promise((_, reject) => {
		rejectOnClose(options.ws, reject, ac.signal);
	});

	return Promise.race([settle, close]).finally(() => ac.abort());
}

function getProcedureHelpers<TSend, TRecv = TSend>(): ProcedureHelpers<TSend, TRecv> {
	return {
		send: (payload) => ({ action: "send", payload }),
		recv: () => ({ action: "recv", strict: false }),
		expect: (predicate) => ({ action: "recv", predicate, strict: true }),
		settle: (promise) => ({ action: "settle", promise }),
	};
}

function rejectOnClose(
	ws: TaggedWebSocket,
	reject: (reason?: any) => void,
	signal: AbortSignal,
): void {
	ws.addEventListener(
		"close",
		(event) => {
			reject(
				ws[closeTag]
					? new AbortedError("Procedure aborted")
					: new WebSocketClosedError(
							"Connection was closed while executing procedure",
							event,
						),
			);
		},
		{ signal },
	);
}