NTSH-Control/src/Unity/UnityWebSocket.ts

import { RawData, WebSocket } from 'ws';
import { Main } from '../Main';
import { delay, ServiceState } from '../Utils';

const PREFIX = '[Unity]';
export class UnityWebSocket {
	private _Main: Main;

	state: ServiceState = 'DISCONNECTED';
	message?: string = 'Waiting for process...';
	error?: string;

	parameters: UnityParameters = {
		timelineWatching: false,
		timelineStanding: false,
		timelineProgress: 0,
		zedPath: '',
		zedReady: false,
		zedFPS: '-',
		outOfService: null,
		sliders: [],
		sensors: [],
	};

	socket: WebSocket;

	constructor(Main: Main) {
		this._Main = Main;
	}

	handle(command: string, ...args: any[]) {
		switch (command) {
			case 'parameterValue':
				const sliderIndex: number = args[0];
				const percentage: number = args[1];

				this.setSliderValue(sliderIndex, percentage);
				break;

			case 'enableOutOfService':
				const enableCallback: Function = args[0];
				if (typeof enableCallback !== 'function') return;

				this.setOutOfService(true);

				enableCallback({ succeed: true });
				break;

			case 'disableOutOfService':
				const disableCallback: Function = args[0];
				if (typeof disableCallback !== 'function') return;

				this.setOutOfService(false);

				disableCallback({ succeed: true });
				break;
		}
	}

	quitApplication() {
		if (this.socket == null || this.socket.readyState !== WebSocket.OPEN)
			return;

		this.socket.send(
			JSON.stringify({
				type: 'quit_application',
			})
		);
	}

	setSliderValue(sliderIndex: number, sliderValue: number) {
		if (this.socket == null || this.socket.readyState !== WebSocket.OPEN)
			return;

		this.socket.send(
			JSON.stringify({
				type: 'set_slider_value',
				sliderIndex,
				sliderValue,
			})
		);

		if (this.parameters.sliders[sliderIndex] == undefined) return;
		this.parameters.sliders[sliderIndex].outputValue = sliderValue;
		this.broadcastState();
	}

	setOutOfService(state: boolean) {
		if (this.socket == null || this.socket.readyState !== WebSocket.OPEN)
			return;

		this.socket.send(
			JSON.stringify({
				type: 'set_out_of_service',
				showOutOfService: state,
			})
		);

		this.parameters.outOfService = true;
		this.broadcastState();
	}

	broadcastState() {
		this._Main.WebServer.socket.emit(
			'unityWebSocketState',
			this.getState()
		);
	}

	setInfo(message: string, error: string, state: ServiceState = 'FAILED') {
		this.message = message;
		this.error = error;
		this.state = state;
		this.broadcastState();

		if (state == 'FAILED' || state == 'DISCONNECTED')
			console.warn(PREFIX, message ?? error);
		else console.log(PREFIX, message ?? error);
	}

	stopFetchClocks() {
		clearInterval(this.heartbeatClock);
		clearInterval(this.calibrationImageClock);
	}

	handleSocketMessage(data: RawData) {
		let message: UnitySocketMessage;
		try {
			message = JSON.parse(data.toString());
		} catch (error) {
			return;
		}

		switch (message.type) {
			case 'heartbeat_data':
				this.parameters.timelineWatching =
					message.heartbeat.dataTimeline.isWatching;
				this.parameters.timelineStanding =
					message.heartbeat.dataTimeline.isStanding;
				this.parameters.timelineProgress =
					message.heartbeat.dataTimeline.timelineProgress;
				this.parameters.zedPath = `${message.heartbeat.zedCamera.streamInputIP}:${message.heartbeat.zedCamera.streamInputPort}`;
				this.parameters.zedReady =
					message.heartbeat.zedCamera.isZedReady;
				this.parameters.zedFPS = message.heartbeat.zedCamera.cameraFPS;
				this.parameters.outOfService =
					message.heartbeat.showOutOfService ?? false;
				this.parameters.sensors = message.heartbeat.dataSensors;
				this.parameters.sliders = message.heartbeat.dataSliders.map(
					(slider) => {
						return {
							...slider,
							min: slider.min ?? 0,
							max: slider.max ?? 1,
							unit: slider.unit ?? '%',
							visualMultiplier:
								(slider.min ?? 0) == 0 && (slider.max ?? 1) == 1
									? 100
									: null,
							decimalPlaces:
								(slider.min ?? 0) == 0 && (slider.max ?? 1) == 1
									? 0
									: 2,
						} as UnityParameterSlider;
					}
				);

				this.broadcastState();
				break;

			case 'response_camera_frame':
				this._Main.WebServer.Calibration.writeCalibrationImage(
					message.imageBase64
				);
				break;
		}
	}

	disconnected: boolean = false;
	async disconnect() {
		this.restartRequested = true;
		this.disconnected = true;

		if (this.socket != null) {
			this.socket.close();
			this.socket = null;
		}

		this.stopFetchClocks();

		this.setInfo('Waiting for process...', null, 'DISCONNECTED');
	}

	private restartRequested = false;
	async reconnect() {
		if (this.restartRequested) return;
		if (this.disconnected) return;
		this.restartRequested = true;

		this.stopFetchClocks();

		if (this.socket != null) {
			this.socket.close();
			this.socket = null;
		}

		await delay(2000);
		if (this.disconnected) return;

		this.message = `Reconnecting in 10 seconds...`;
		this.broadcastState();

		await delay(10000);
		if (this.disconnected) return;

		await this.connect();
	}

	async connect() {
		this.restartRequested = false;
		this.disconnected = false;

		this.stopFetchClocks();

		this.setInfo('Connecting...', null, 'CONNECTING');

		await delay(1000);

		this.socket = new WebSocket(
			`ws://${this._Main.Config.unity.webSocket.ip}:${this._Main.Config.unity.webSocket.port}`
		);

		this.socket.on('error', (error) => {
			if (this.restartRequested) return;
			this.setInfo(
				'Connection error',
				`Could not connect: ${error.message}`,
				'FAILED'
			);
			this.reconnect();
		});

		this.socket.on('open', () => {
			this.startFetchClocks();

			this.setInfo('Connected', null, 'CONNECTED');
		});

		this.socket.on('close', () => {
			if (this._Main.UnityRunner.restartTriggered) return;
			if (this.restartRequested) return;
			this.setInfo(
				'Disconnected',
				'Unity was disconnected unexpectedly',
				'FAILED'
			);
			this.reconnect();
		});

		this.socket.on('message', (data) => this.handleSocketMessage(data));
	}

	private heartbeatClock: NodeJS.Timeout;
	private calibrationImageClock: NodeJS.Timeout;
	startFetchClocks() {
		this.socket.send(
			JSON.stringify({ type: 'set_heartbeat_auto_send', autoSend: false })
		);
		this.heartbeatClock = setInterval(() => {
			if (
				this.socket == null ||
				this.socket.readyState !== WebSocket.OPEN
			)
				return;

			this.socket.send(JSON.stringify({ type: 'request_heartbeat' }));
		}, this._Main.Config.unity.heartbeatInterval);

		this.calibrationImageClock = setInterval(() => {
			if (
				this.socket == null ||
				this.socket.readyState !== WebSocket.OPEN
			)
				return;

			this.socket.send(JSON.stringify({ type: 'request_camera_frame' }));
		}, this._Main.Config.unity.calibrationImageInterval);
	}

	getState(): UnityWebSocketStatus {
		return {
			state: this.state,
			message: this.message,
			error: this.error,

			parameters: this.parameters,
		};
	}
}

interface UnityWebSocketStatus {
	state: ServiceState;
	message?: string;
	error?: string;

	parameters: UnityParameters;
}

interface UnityParameters {
	timelineWatching: boolean;
	timelineStanding: boolean;
	timelineProgress: number;
	zedPath: string;
	zedReady: boolean;
	zedFPS: string;
	outOfService: boolean;
	sliders: UnityParameterSlider[];
	sensors: UnitySocketMessageHeartbeat['heartbeat']['dataSensors'];
}

type UnityHeartbeatSlider =
	UnitySocketMessageHeartbeat['heartbeat']['dataSliders'][number];
interface UnityParameterSlider extends UnityHeartbeatSlider {
	visualMultiplier?: number;
	decimalPlaces?: number;
}

type UnitySocketMessage =
	| UnitySocketMessageHeartbeat
	| UnitySocketMessageCameraFrame;

interface UnitySocketMessageBase {
	type: string;
	timestamp: number;
}
interface UnitySocketMessageHeartbeat extends UnitySocketMessageBase {
	type: 'heartbeat_data';
	heartbeat: {
		dataSensors: {
			sensorIndex: number;
			deviceName: string;
			outputValue: number;
		}[];
		dataSliders: {
			sliderIndex: number;
			sliderName: string;
			outputValue: number;
			min: number;
			max: number;
			unit: string;
		}[];
		dataTimeline: {
			isStanding: boolean;
			isWatching: boolean;
			timelineProgress: number;
		};
		zedCamera: {
			cameraFPS: string;
			isZedReady: boolean;
			streamInputIP: string;
			streamInputPort: number;
			zedGrabError: number;
		};
		showOutOfService?: boolean;
	};
}
interface UnitySocketMessageCameraFrame extends UnitySocketMessageBase {
	type: 'response_camera_frame';
	imageBase64: string;
}