// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
import { Observable, Subscription, SubscriptionLike } from 'rxjs';
import { GraphQLError } from 'graphql';
import { ConsoleLogger, Hub, HubPayload } from '@aws-amplify/core';
import {
	CustomUserAgentDetails,
	DocumentType,
	NonRetryableError,
	amplifyUuid,
	base64Encoder,
	isNonRetryableError,
	jitteredExponentialRetry,
} from '@aws-amplify/core/internals/utils';

import {
	CONTROL_MSG,
	ConnectionState,
	PubSubContentObserver,
} from '../../types/PubSub';
import {
	AMPLIFY_SYMBOL,
	CONNECTION_INIT_TIMEOUT,
	CONNECTION_STATE_CHANGE,
	DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT,
	DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT,
	MAX_DELAY_MS,
	MESSAGE_TYPES,
	NON_RETRYABLE_CODES,
	NON_RETRYABLE_ERROR_TYPES,
	SOCKET_STATUS,
	START_ACK_TIMEOUT,
	SUBSCRIPTION_STATUS,
} from '../constants';
import {
	CONNECTION_CHANGE,
	ConnectionStateMonitor,
} from '../../utils/ConnectionStateMonitor';
import {
	ReconnectEvent,
	ReconnectionMonitor,
} from '../../utils/ReconnectionMonitor';
import type { AWSAppSyncRealTimeProviderOptions } from '../AWSAppSyncRealTimeProvider';

import {
	additionalHeadersFromOptions,
	queryParamsFromCustomHeaders,
	realtimeUrlWithQueryString,
} from './appsyncUrl';
import { awsRealTimeHeaderBasedAuth } from './authHeaders';

const dispatchApiEvent = (payload: HubPayload) => {
	Hub.dispatch('api', payload, 'PubSub', AMPLIFY_SYMBOL);
};

export interface ObserverQuery {
	observer: PubSubContentObserver;
	query: string;
	variables: DocumentType;
	subscriptionState: SUBSCRIPTION_STATUS;
	subscriptionReadyCallback?(): void;
	subscriptionFailedCallback?(reason?: any): void;
	startAckTimeoutId?: ReturnType<typeof setTimeout>;
}

interface ParsedMessagePayload {
	type: string;
	payload: {
		connectionTimeoutMs: number;
		errors?: [{ errorType: string; errorCode: number }];
	};
}

interface AWSWebSocketProviderArgs {
	providerName: string;
	wsProtocolName: string;
	connectUri: string;
}

export abstract class AWSWebSocketProvider {
	protected logger: ConsoleLogger;
	protected subscriptionObserverMap = new Map<string, ObserverQuery>();
	protected allowNoSubscriptions = false;

	protected awsRealTimeSocket?: WebSocket;
	private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
	private keepAliveTimestamp: number = Date.now();
	private keepAliveHeartbeatIntervalId?: ReturnType<typeof setInterval>;
	private promiseArray: { res(): void; rej(reason?: any): void }[] = [];
	private connectionState: ConnectionState | undefined;
	private readonly connectionStateMonitor = new ConnectionStateMonitor();
	private readonly reconnectionMonitor = new ReconnectionMonitor();
	private connectionStateMonitorSubscription: SubscriptionLike;
	private readonly wsProtocolName: string;
	private readonly wsConnectUri: string;

	constructor(args: AWSWebSocketProviderArgs) {
		this.logger = new ConsoleLogger(args.providerName);
		this.wsProtocolName = args.wsProtocolName;
		this.wsConnectUri = args.connectUri;

		this.connectionStateMonitorSubscription =
			this._startConnectionStateMonitoring();
	}

	/**
	 * Mark the socket closed and release all active listeners
	 */
	close() {
		// Mark the socket closed both in status and the connection monitor
		this.socketStatus = SOCKET_STATUS.CLOSED;
		this.connectionStateMonitor.record(CONNECTION_CHANGE.CONNECTION_FAILED);

		// Turn off the subscription monitor Hub publishing
		this.connectionStateMonitorSubscription.unsubscribe();
		// Complete all reconnect observers
		this.reconnectionMonitor.close();

		return new Promise<void>((resolve, reject) => {
			if (this.awsRealTimeSocket) {
				this.awsRealTimeSocket.onclose = (_: CloseEvent) => {
					this._closeSocket();
					this.subscriptionObserverMap = new Map();
					this.awsRealTimeSocket = undefined;
					resolve();
				};

				this.awsRealTimeSocket.onerror = (err: any) => {
					reject(err);
				};

				this.awsRealTimeSocket.close();
			} else {
				resolve();
			}
		});
	}

	subscribe(
		options?: AWSAppSyncRealTimeProviderOptions,
		customUserAgentDetails?: CustomUserAgentDetails,
	): Observable<Record<string, unknown>> {
		return new Observable(observer => {
			if (!options?.appSyncGraphqlEndpoint) {
				observer.error({
					errors: [
						{
							...new GraphQLError(
								`Subscribe only available for AWS AppSync endpoint`,
							),
						},
					],
				});
				observer.complete();

				return;
			}

			let subscriptionStartInProgress = false;
			const subscriptionId = amplifyUuid();

			const startSubscription = () => {
				if (!subscriptionStartInProgress) {
					subscriptionStartInProgress = true;

					this._startSubscriptionWithAWSAppSyncRealTime({
						options,
						observer,
						subscriptionId,
						customUserAgentDetails,
					})
						.catch(err => {
							this.logger.debug(
								`${CONTROL_MSG.REALTIME_SUBSCRIPTION_INIT_ERROR}: ${err}`,
							);
							this._closeSocket();
						})
						.finally(() => {
							subscriptionStartInProgress = false;
						});
				}
			};

			// Add an observable to the reconnection list to manage reconnection for this subscription
			const reconnectSubscription = new Observable(
				reconnectSubscriptionObserver => {
					this.reconnectionMonitor.addObserver(reconnectSubscriptionObserver);
				},
			).subscribe(() => {
				startSubscription();
			});

			startSubscription();

			return async () => {
				await this._cleanupSubscription(subscriptionId, reconnectSubscription);
			};
		});
	}

	protected async connect(
		options: AWSAppSyncRealTimeProviderOptions,
	): Promise<void> {
		if (this.socketStatus === SOCKET_STATUS.READY) {
			return;
		}

		await this._connectWebSocket(options);
	}

	protected async publish(
		options: AWSAppSyncRealTimeProviderOptions,
		customUserAgentDetails?: CustomUserAgentDetails,
	): Promise<void> {
		if (this.socketStatus !== SOCKET_STATUS.READY) {
			throw new Error('Subscription has not been initialized');
		}

		return this._publishMessage(options, customUserAgentDetails);
	}

	private async _connectWebSocket(options: AWSAppSyncRealTimeProviderOptions) {
		const { apiKey, appSyncGraphqlEndpoint, authenticationType, region } =
			options;

		const { additionalCustomHeaders } =
			await additionalHeadersFromOptions(options);

		this.connectionStateMonitor.record(CONNECTION_CHANGE.OPENING_CONNECTION);
		await this._initializeWebSocketConnection({
			apiKey,
			appSyncGraphqlEndpoint,
			authenticationType,
			region,
			additionalCustomHeaders,
		});
	}

	private async _publishMessage(
		options: AWSAppSyncRealTimeProviderOptions,
		customUserAgentDetails?: CustomUserAgentDetails,
	): Promise<void> {
		const subscriptionId = amplifyUuid();

		const { additionalCustomHeaders, libraryConfigHeaders } =
			await additionalHeadersFromOptions(options);

		const serializedSubscriptionMessage =
			await this._prepareSubscriptionPayload({
				options,
				subscriptionId,
				customUserAgentDetails,
				additionalCustomHeaders,
				libraryConfigHeaders,
				publish: true,
			});

		return new Promise((resolve, reject) => {
			if (this.awsRealTimeSocket) {
				const publishListener = (event: MessageEvent) => {
					const data = JSON.parse(event.data);
					if (data.id === subscriptionId && data.type === 'publish_success') {
						this.awsRealTimeSocket &&
							this.awsRealTimeSocket.removeEventListener(
								'message',
								publishListener,
							);
						cleanup();
						resolve();
					}

					if (data.errors && data.errors.length > 0) {
						const errorTypes = data.errors.map((error: any) => error.errorType);
						cleanup();
						reject(new Error(`Publish errors: ${errorTypes.join(', ')}`));
					}
				};

				const errorListener = (error: Event) => {
					cleanup();
					reject(new Error(`WebSocket error: ${error}`));
				};

				const closeListener = () => {
					cleanup();
					reject(new Error('WebSocket is closed'));
				};

				const cleanup = () => {
					this.awsRealTimeSocket?.removeEventListener(
						'message',
						publishListener,
					);
					this.awsRealTimeSocket?.removeEventListener('error', errorListener);
					this.awsRealTimeSocket?.removeEventListener('close', closeListener);
				};

				this.awsRealTimeSocket.addEventListener('message', publishListener);
				this.awsRealTimeSocket.addEventListener('error', errorListener);
				this.awsRealTimeSocket.addEventListener('close', closeListener);

				this.awsRealTimeSocket.send(serializedSubscriptionMessage);
			} else {
				reject(new Error('WebSocket is not connected'));
			}
		});
	}

	private async _cleanupSubscription(
		subscriptionId: string,
		reconnectSubscription: Subscription,
	) {
		// Cleanup reconnection subscription
		reconnectSubscription?.unsubscribe();

		// Cleanup after unsubscribing or observer.complete was called after _startSubscriptionWithAWSAppSyncRealTime
		try {
			// Waiting that subscription has been connected before trying to unsubscribe
			await this._waitForSubscriptionToBeConnected(subscriptionId);

			const { subscriptionState } =
				this.subscriptionObserverMap.get(subscriptionId) || {};

			if (!subscriptionState) {
				// subscription already unsubscribed
				return;
			}

			if (subscriptionState === SUBSCRIPTION_STATUS.CONNECTED) {
				this._sendUnsubscriptionMessage(subscriptionId);
			} else {
				throw new Error('Subscription never connected');
			}
		} catch (err) {
			this.logger.debug(`Error while unsubscribing ${err}`);
		} finally {
			this._removeSubscriptionObserver(subscriptionId);
		}
	}

	// Monitor the connection state and pass changes along to Hub
	private _startConnectionStateMonitoring() {
		return this.connectionStateMonitor.connectionStateObservable.subscribe(
			connectionState => {
				dispatchApiEvent({
					event: CONNECTION_STATE_CHANGE,
					data: {
						provider: this,
						connectionState,
					},
					message: `Connection state is ${connectionState}`,
				});
				this.connectionState = connectionState;

				// Trigger START_RECONNECT when the connection is disrupted
				if (connectionState === ConnectionState.ConnectionDisrupted) {
					this.reconnectionMonitor.record(ReconnectEvent.START_RECONNECT);
				}

				// Trigger HALT_RECONNECT to halt reconnection attempts when the state is anything other than
				// ConnectionDisrupted or Connecting
				if (
					[
						ConnectionState.Connected,
						ConnectionState.ConnectedPendingDisconnect,
						ConnectionState.ConnectedPendingKeepAlive,
						ConnectionState.ConnectedPendingNetwork,
						ConnectionState.ConnectionDisruptedPendingNetwork,
						ConnectionState.Disconnected,
					].includes(connectionState)
				) {
					this.reconnectionMonitor.record(ReconnectEvent.HALT_RECONNECT);
				}
			},
		);
	}

	protected abstract _prepareSubscriptionPayload(param: {
		options: AWSAppSyncRealTimeProviderOptions;
		subscriptionId: string;
		customUserAgentDetails: CustomUserAgentDetails | undefined;
		additionalCustomHeaders: Record<string, string>;
		libraryConfigHeaders: Record<string, string>;
		publish?: boolean;
	}): Promise<string>;

	private async _startSubscriptionWithAWSAppSyncRealTime({
		options,
		observer,
		subscriptionId,
		customUserAgentDetails,
	}: {
		options: AWSAppSyncRealTimeProviderOptions;
		observer: PubSubContentObserver;
		subscriptionId: string;
		customUserAgentDetails: CustomUserAgentDetails | undefined;
	}) {
		const { query, variables } = options;

		this.subscriptionObserverMap.set(subscriptionId, {
			observer,
			query: query ?? '',
			variables: variables ?? {},
			subscriptionState: SUBSCRIPTION_STATUS.PENDING,
			startAckTimeoutId: undefined,
		});

		const { additionalCustomHeaders, libraryConfigHeaders } =
			await additionalHeadersFromOptions(options);

		const serializedSubscriptionMessage =
			await this._prepareSubscriptionPayload({
				options,
				subscriptionId,
				customUserAgentDetails,
				additionalCustomHeaders,
				libraryConfigHeaders,
			});

		try {
			await this._connectWebSocket(options);
		} catch (err: any) {
			this._logStartSubscriptionError(subscriptionId, observer, err);

			return;
		}

		// Potential race condition can occur when unsubscribe is called during _initializeWebSocketConnection.
		// E.g.unsubscribe gets invoked prior to finishing WebSocket handshake or START_ACK.
		// Both subscriptionFailedCallback and subscriptionReadyCallback are used to synchronized this.
		const { subscriptionFailedCallback, subscriptionReadyCallback } =
			this.subscriptionObserverMap.get(subscriptionId) ?? {};

		// This must be done before sending the message in order to be listening immediately
		this.subscriptionObserverMap.set(subscriptionId, {
			observer,
			subscriptionState: SUBSCRIPTION_STATUS.PENDING,
			query: query ?? '',
			variables: variables ?? {},
			subscriptionReadyCallback,
			subscriptionFailedCallback,
			startAckTimeoutId: setTimeout(() => {
				this._timeoutStartSubscriptionAck(subscriptionId);
			}, START_ACK_TIMEOUT),
		});

		if (this.awsRealTimeSocket) {
			this.awsRealTimeSocket.send(serializedSubscriptionMessage);
		}
	}

	// Log logic for start subscription failures
	private _logStartSubscriptionError(
		subscriptionId: string,
		observer: PubSubContentObserver,
		err: { message?: string },
	) {
		this.logger.debug({ err });
		const message = String(err.message ?? '');
		// Resolving to give the state observer time to propogate the update
		this._closeSocket();

		// Capture the error only when the network didn't cause disruption
		if (
			this.connectionState !== ConnectionState.ConnectionDisruptedPendingNetwork
		) {
			// When the error is non-retriable, error out the observable
			if (isNonRetryableError(err)) {
				observer.error({
					errors: [
						{
							...new GraphQLError(
								`${CONTROL_MSG.CONNECTION_FAILED}: ${message}`,
							),
						},
					],
				});
			} else {
				this.logger.debug(`${CONTROL_MSG.CONNECTION_FAILED}: ${message}`);
			}

			const { subscriptionFailedCallback } =
				this.subscriptionObserverMap.get(subscriptionId) || {};

			// Notify concurrent unsubscription
			if (typeof subscriptionFailedCallback === 'function') {
				subscriptionFailedCallback();
			}
		}
	}

	// Waiting that subscription has been connected before trying to unsubscribe
	private async _waitForSubscriptionToBeConnected(subscriptionId: string) {
		const subscriptionObserver =
			this.subscriptionObserverMap.get(subscriptionId);

		if (subscriptionObserver) {
			const { subscriptionState } = subscriptionObserver;
			// This in case unsubscribe is invoked before sending start subscription message
			if (subscriptionState === SUBSCRIPTION_STATUS.PENDING) {
				return new Promise<void>((resolve, reject) => {
					const {
						observer,
						subscriptionState: observedSubscriptionState,
						variables,
						query,
					} = subscriptionObserver;
					this.subscriptionObserverMap.set(subscriptionId, {
						observer,
						subscriptionState: observedSubscriptionState,
						variables,
						query,
						subscriptionReadyCallback: resolve,
						subscriptionFailedCallback: reject,
					});
				});
			}
		}
	}

	protected abstract _unsubscribeMessage(subscriptionId: string): {
		id: string;
		type: string;
	};

	private _sendUnsubscriptionMessage(subscriptionId: string) {
		try {
			if (
				this.awsRealTimeSocket &&
				this.awsRealTimeSocket.readyState === WebSocket.OPEN &&
				this.socketStatus === SOCKET_STATUS.READY
			) {
				// Preparing unsubscribe message to stop receiving messages for that subscription
				const unsubscribeMessage = this._unsubscribeMessage(subscriptionId);
				const stringToAWSRealTime = JSON.stringify(unsubscribeMessage);
				this.awsRealTimeSocket.send(stringToAWSRealTime);
			}
		} catch (err) {
			// If GQL_STOP is not sent because of disconnection issue, then there is nothing the client can do
			this.logger.debug({ err });
		}
	}

	private _removeSubscriptionObserver(subscriptionId: string) {
		this.subscriptionObserverMap.delete(subscriptionId);

		// Verifying 1000ms after removing subscription in case there are new subscription unmount/mount
		if (!this.allowNoSubscriptions) {
			setTimeout(this._closeSocketIfRequired.bind(this), 1000);
		}
	}

	protected _closeSocketIfRequired() {
		if (this.subscriptionObserverMap.size > 0) {
			// Active subscriptions on the WebSocket
			return;
		}

		if (!this.awsRealTimeSocket) {
			this.socketStatus = SOCKET_STATUS.CLOSED;

			return;
		}

		this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSING_CONNECTION);

		if (this.awsRealTimeSocket.bufferedAmount > 0) {
			// Still data on the WebSocket
			setTimeout(this._closeSocketIfRequired.bind(this), 1000);
		} else {
			this.logger.debug('closing WebSocket...');

			const tempSocket = this.awsRealTimeSocket;
			// Cleaning callbacks to avoid race condition, socket still exists
			tempSocket.onclose = null;
			tempSocket.onerror = null;
			tempSocket.close(1000);
			this.awsRealTimeSocket = undefined;
			this.socketStatus = SOCKET_STATUS.CLOSED;
			this._closeSocket();
		}
	}

	protected abstract _handleSubscriptionData(
		message: MessageEvent,
	): [
		boolean,
		{ id: string; payload: string | Record<string, unknown>; type: string },
	];

	protected abstract _extractConnectionTimeout(
		data: Record<string, any>,
	): number;

	protected abstract _extractErrorCodeAndType(data: Record<string, any>): {
		errorCode: number;
		errorType: string;
	};

	private maintainKeepAlive() {
		this.keepAliveTimestamp = Date.now();
	}

	private keepAliveHeartbeat(connectionTimeoutMs: number) {
		const currentTime = Date.now();

		// Check for missed KA message
		if (
			currentTime - this.keepAliveTimestamp >
			DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT
		) {
			this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE_MISSED);
		} else {
			this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE);
		}

		// Recognize we are disconnected if we haven't seen messages in the keep alive timeout period
		if (currentTime - this.keepAliveTimestamp > connectionTimeoutMs) {
			this._errorDisconnect(CONTROL_MSG.TIMEOUT_DISCONNECT);
		}
	}

	private _handleIncomingSubscriptionMessage(message: MessageEvent) {
		if (typeof message.data !== 'string') {
			return;
		}

		const [isData, data] = this._handleSubscriptionData(message);
		if (isData) {
			this.maintainKeepAlive();

			return;
		}

		const { type, id, payload } = data;

		const {
			observer = null,
			query = '',
			variables = {},
			startAckTimeoutId,
			subscriptionReadyCallback,
			subscriptionFailedCallback,
		} = this.subscriptionObserverMap.get(id) || {};

		if (
			type === MESSAGE_TYPES.GQL_START_ACK ||
			type === MESSAGE_TYPES.EVENT_SUBSCRIBE_ACK
		) {
			this.logger.debug(
				`subscription ready for ${JSON.stringify({ query, variables })}`,
			);
			if (typeof subscriptionReadyCallback === 'function') {
				subscriptionReadyCallback();
			}
			if (startAckTimeoutId) clearTimeout(startAckTimeoutId);
			dispatchApiEvent({
				event: CONTROL_MSG.SUBSCRIPTION_ACK,
				data: { query, variables },
				message: 'Connection established for subscription',
			});
			const subscriptionState = SUBSCRIPTION_STATUS.CONNECTED;
			if (observer) {
				this.subscriptionObserverMap.set(id, {
					observer,
					query,
					variables,
					startAckTimeoutId: undefined,
					subscriptionState,
					subscriptionReadyCallback,
					subscriptionFailedCallback,
				});
			}
			this.connectionStateMonitor.record(
				CONNECTION_CHANGE.CONNECTION_ESTABLISHED,
			);

			return;
		}

		if (type === MESSAGE_TYPES.GQL_CONNECTION_KEEP_ALIVE) {
			this.maintainKeepAlive();

			return;
		}

		if (
			type === MESSAGE_TYPES.GQL_ERROR ||
			type === MESSAGE_TYPES.EVENT_SUBSCRIBE_ERROR
		) {
			const subscriptionState = SUBSCRIPTION_STATUS.FAILED;
			if (observer) {
				this.subscriptionObserverMap.set(id, {
					observer,
					query,
					variables,
					startAckTimeoutId,
					subscriptionReadyCallback,
					subscriptionFailedCallback,
					subscriptionState,
				});

				let errorMessage = JSON.stringify(payload ?? data);
				let isAuthError = false;

				const AUTH_ERROR_TYPES = [
					'UnauthorizedException',
					'Unauthorized',
					'NotAuthorizedException',
				];

				if (type === MESSAGE_TYPES.EVENT_SUBSCRIBE_ERROR) {
					const { errors } = JSON.parse(String(message.data));
					if (Array.isArray(errors) && errors.length > 0) {
						const error = errors[0];
						errorMessage = `${error.errorType}: ${error.message}`;
						isAuthError =
							AUTH_ERROR_TYPES.includes(error.errorType) ||
							error.message?.includes('Token expired');
					}
				} else if (
					type === MESSAGE_TYPES.GQL_ERROR &&
					payload &&
					typeof payload === 'object'
				) {
					const { errors } = payload as Record<string, unknown>;
					if (Array.isArray(errors) && errors.length > 0) {
						const error = errors[0] as Record<string, string>;
						isAuthError =
							AUTH_ERROR_TYPES.includes(error.errorType) ||
							error.message?.includes('Token expired');
					}
				}

				this.logger.debug(`${CONTROL_MSG.CONNECTION_FAILED}: ${errorMessage}`);

				// On auth errors, close the socket to trigger reconnection with fresh tokens.
				// The onclose handler fires _errorDisconnect → CONNECTION_CLOSED, which the
				// ConnectionStateMonitor picks up as ConnectionDisrupted and the
				// ReconnectionMonitor reconnects with refreshed credentials.
				if (isAuthError && this.awsRealTimeSocket) {
					this.logger.warn(
						'Subscription failed due to auth error, closing WebSocket to trigger reconnection with fresh tokens',
					);
					this.awsRealTimeSocket.close(1000, 'Auth error - reconnecting');
				}

				observer.error({
					errors: [
						{
							...new GraphQLError(
								`${CONTROL_MSG.CONNECTION_FAILED}: ${errorMessage}`,
							),
						},
					],
				});

				if (startAckTimeoutId) clearTimeout(startAckTimeoutId);

				if (typeof subscriptionFailedCallback === 'function') {
					subscriptionFailedCallback();
				}
			}
		}
	}

	private _errorDisconnect(msg: string) {
		this.logger.debug(`Disconnect error: ${msg}`);

		if (this.awsRealTimeSocket) {
			this._closeSocket();
			this.awsRealTimeSocket.close();
		}

		this.socketStatus = SOCKET_STATUS.CLOSED;
	}

	private _closeSocket() {
		if (this.keepAliveHeartbeatIntervalId) {
			clearInterval(this.keepAliveHeartbeatIntervalId);
			this.keepAliveHeartbeatIntervalId = undefined;
		}
		this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
	}

	private _timeoutStartSubscriptionAck(subscriptionId: string) {
		const subscriptionObserver =
			this.subscriptionObserverMap.get(subscriptionId);
		if (subscriptionObserver) {
			const { observer, query, variables } = subscriptionObserver;
			if (!observer) {
				return;
			}
			this.subscriptionObserverMap.set(subscriptionId, {
				observer,
				query,
				variables,
				subscriptionState: SUBSCRIPTION_STATUS.FAILED,
			});

			this._closeSocket();
			this.logger.debug(
				'timeoutStartSubscription',
				JSON.stringify({ query, variables }),
			);
		}
	}

	private _initializeWebSocketConnection({
		appSyncGraphqlEndpoint,
		authenticationType,
		apiKey,
		region,
		additionalCustomHeaders,
	}: AWSAppSyncRealTimeProviderOptions) {
		if (this.socketStatus === SOCKET_STATUS.READY) {
			return;
		}

		// TODO(Eslint): refactor to now use async function as the promise executor
		// eslint-disable-next-line no-async-promise-executor
		return new Promise<void>(async (resolve, reject) => {
			this.promiseArray.push({ res: resolve, rej: reject });

			if (this.socketStatus === SOCKET_STATUS.CLOSED) {
				try {
					this.socketStatus = SOCKET_STATUS.CONNECTING;

					// Empty payload on connect
					const payloadString = '{}';

					const authHeader = await awsRealTimeHeaderBasedAuth({
						authenticationType,
						payload: payloadString,
						canonicalUri: this.wsConnectUri,
						apiKey,
						appSyncGraphqlEndpoint,
						region,
						additionalCustomHeaders,
					});

					const headerString = authHeader ? JSON.stringify(authHeader) : '';
					// base64url-encoded string
					const encodedHeader = base64Encoder.convert(headerString, {
						urlSafe: true,
						skipPadding: true,
					});

					const authTokenSubprotocol = `header-${encodedHeader}`;

					const queryParams = queryParamsFromCustomHeaders(
						additionalCustomHeaders,
					);

					const awsRealTimeUrl = realtimeUrlWithQueryString(
						appSyncGraphqlEndpoint,
						queryParams,
					);

					await this._establishRetryableConnection(
						awsRealTimeUrl,
						authTokenSubprotocol,
					);

					this.promiseArray.forEach(({ res }) => {
						this.logger.debug('Notifying connection successful');
						res();
					});
					this.socketStatus = SOCKET_STATUS.READY;
					this.promiseArray = [];
				} catch (err) {
					this.logger.debug('Connection exited with', err);
					this.promiseArray.forEach(({ rej }) => {
						rej(err);
					});
					this.promiseArray = [];
					if (
						this.awsRealTimeSocket &&
						this.awsRealTimeSocket.readyState === WebSocket.OPEN
					) {
						this.awsRealTimeSocket.close(3001);
					}
					this.awsRealTimeSocket = undefined;
					this.socketStatus = SOCKET_STATUS.CLOSED;
				}
			}
		});
	}

	private async _establishRetryableConnection(
		awsRealTimeUrl: string,
		subprotocol: string,
	) {
		this.logger.debug(`Establishing retryable connection`);
		await jitteredExponentialRetry(
			this._establishConnection.bind(this),
			[awsRealTimeUrl, subprotocol],
			MAX_DELAY_MS,
		);
	}

	private async _openConnection(awsRealTimeUrl: string, subprotocol: string) {
		return new Promise<void>((resolve, reject) => {
			const newSocket = this._getNewWebSocket(awsRealTimeUrl, [
				this.wsProtocolName,
				subprotocol,
			]);

			newSocket.onerror = () => {
				this.logger.debug(`WebSocket connection error`);
			};
			newSocket.onclose = () => {
				this._closeSocket();
				reject(new Error('Connection handshake error'));
			};
			newSocket.onopen = () => {
				this.awsRealTimeSocket = newSocket;
				resolve();
			};
		});
	}

	private _getNewWebSocket(url: string, protocol: string[]) {
		return new WebSocket(url, protocol);
	}

	private async _initiateHandshake(): Promise<string> {
		return new Promise((resolve, reject) => {
			if (!this.awsRealTimeSocket) {
				reject(new Error('awsRealTimeSocket undefined'));

				return;
			}

			let ackOk = false;

			this.awsRealTimeSocket.onerror = error => {
				this.logger.debug(`WebSocket error ${JSON.stringify(error)}`);
			};

			this.awsRealTimeSocket.onclose = event => {
				this.logger.debug(`WebSocket closed ${event.reason}`);
				this._closeSocket();
				reject(new Error(JSON.stringify(event)));
			};

			this.awsRealTimeSocket.onmessage = (message: MessageEvent) => {
				if (typeof message.data !== 'string') {
					return;
				}
				this.logger.debug(
					`subscription message from AWS AppSyncRealTime: ${message.data} `,
				);

				const data = JSON.parse(message.data) as ParsedMessagePayload;

				const { type } = data;

				const connectionTimeoutMs = this._extractConnectionTimeout(data);

				if (type === MESSAGE_TYPES.GQL_CONNECTION_ACK) {
					ackOk = true;
					this._registerWebsocketHandlers(connectionTimeoutMs);
					resolve('Connected to AWS AppSyncRealTime');

					return;
				}

				if (type === MESSAGE_TYPES.GQL_CONNECTION_ERROR) {
					const { errorType, errorCode } = this._extractErrorCodeAndType(data);

					// TODO(Eslint): refactor to reject an Error object instead of a plain object
					// eslint-disable-next-line prefer-promise-reject-errors
					reject({ errorType, errorCode });
				}
			};

			const gqlInit = {
				type: MESSAGE_TYPES.GQL_CONNECTION_INIT,
			};
			this.awsRealTimeSocket.send(JSON.stringify(gqlInit));

			const checkAckOk = (targetAckOk: boolean) => {
				if (!targetAckOk) {
					this.connectionStateMonitor.record(
						CONNECTION_CHANGE.CONNECTION_FAILED,
					);
					reject(
						new Error(
							`Connection timeout: ack from AWSAppSyncRealTime was not received after ${CONNECTION_INIT_TIMEOUT} ms`,
						),
					);
				}
			};

			setTimeout(() => {
				checkAckOk(ackOk);
			}, CONNECTION_INIT_TIMEOUT);
		});
	}

	private _registerWebsocketHandlers(connectionTimeoutMs: number) {
		if (!this.awsRealTimeSocket) {
			return;
		}

		// Set up a keep alive heartbeat for this connection
		this.keepAliveHeartbeatIntervalId = setInterval(() => {
			this.keepAliveHeartbeat(connectionTimeoutMs);
		}, DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT);

		this.awsRealTimeSocket.onmessage =
			this._handleIncomingSubscriptionMessage.bind(this);

		this.awsRealTimeSocket.onerror = err => {
			this.logger.debug(err);
			this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
		};

		this.awsRealTimeSocket.onclose = event => {
			this.logger.debug(`WebSocket closed ${event.reason}`);
			this._closeSocket();
			this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
		};
	}

	/**
	 * Open WebSocket connection & perform handshake
	 * Ref: https://docs.aws.amazon.com/appsync/latest/devguide/real-time-websocket-client.html#appsynclong-real-time-websocket-client-implementation-guide-for-graphql-subscriptions
	 *
	 * @param subprotocol -
	 */
	private _establishConnection = async (
		awsRealTimeUrl: string,
		subprotocol: string,
	) => {
		this.logger.debug(`Establishing WebSocket connection to ${awsRealTimeUrl}`);
		try {
			await this._openConnection(awsRealTimeUrl, subprotocol);
			await this._initiateHandshake();
		} catch (err) {
			const { errorType, errorCode } = err as {
				errorType: string;
				errorCode: number;
			};

			if (
				NON_RETRYABLE_CODES.includes(errorCode) ||
				// Event API does not currently return `errorCode`. This may change in the future.
				// For now fall back to also checking known non-retryable error types
				NON_RETRYABLE_ERROR_TYPES.includes(errorType)
			) {
				throw new NonRetryableError(errorType);
			} else if (errorType) {
				throw new Error(errorType);
			} else {
				throw err;
			}
		}
	};
}