telemetry_visualization/frontend/src/composables/websocket.ts

import {
    computed,
    type MaybeRefOrGetter,
    onMounted,
    onUnmounted,
    onWatcherCleanup,
    ref,
    type Ref,
    shallowRef,
    toValue,
    watch,
} from 'vue';
import type { TelemetryDefinition } from '@/composables/telemetry';
import { onDocumentVisibilityChange } from '@/composables/document.ts';
import { parseJsonString, toJsonString } from '@/composables/json.ts';

export interface TelemetryDataItem {
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    value: any;
    timestamp: string;
}

interface TlmValue {
    uuid: string;
    value: TelemetryDataItem;
}

export class WebsocketHandle {
    websocket: WebSocket | null;
    should_be_connected: boolean;
    connected: Ref<boolean>;
    enabled: Ref<boolean>;
    on_telem_value: Map<string, Array<(value: TelemetryDataItem) => void>>;

    constructor() {
        this.websocket = null;
        this.should_be_connected = false;
        this.connected = ref(false);
        this.enabled = ref(true);
        this.on_telem_value = new Map();
    }

    connect() {
        this.should_be_connected = true;
        if (this.websocket != null) {
            return;
        }

        this.websocket = new WebSocket(
            location.protocol.replace('http', 'ws') +
                '//' +
                location.host +
                '/ws',
        );
        this.websocket.addEventListener('open', () => {
            this.connected.value = true;
        });
        this.websocket.addEventListener('close', () => {
            if (this.should_be_connected) {
                this.disconnect();
                setTimeout(() => {
                    this.connect();
                }, 1000);
            }
        });
        this.websocket.addEventListener('error', () => {
            if (this.should_be_connected) {
                this.disconnect();
                setTimeout(() => {
                    this.connect();
                }, 1000);
            }
        });
        this.websocket.addEventListener('message', (event) => {
            const message = parseJsonString(event.data);
            if (message['TlmValue']) {
                const tlm_value = message['TlmValue'] as TlmValue;
                const listeners = this.on_telem_value.get(tlm_value.uuid);
                if (listeners) {
                    listeners.forEach((listener) => {
                        listener(tlm_value.value);
                    });
                }
            }
        });
    }

    disconnect() {
        this.should_be_connected = false;
        if (this.websocket == null) {
            return;
        }

        this.connected.value = false;
        this.websocket.close();
        this.websocket = null;
        this.on_telem_value.clear();
    }

    listen_to_telemetry(
        telemetry: MaybeRefOrGetter<TelemetryDefinition | null>,
        minimum_separation_ms: MaybeRefOrGetter<number> | undefined,
        live: MaybeRefOrGetter<boolean>,
    ) {
        const value_result = ref<TelemetryDataItem | null>(null);

        const uuid = computed(() => {
            const tlm = toValue(telemetry);
            if (tlm) {
                return tlm.uuid;
            }
            return null;
        });

        const minimum_separation = computed(() => {
            const min_sep = toValue(minimum_separation_ms);
            if (min_sep) {
                return min_sep;
            }
            return 0;
        });

        const is_live = computed(() => toValue(live));

        watch(
            [uuid, this.connected, this.enabled, minimum_separation, is_live],
            ([uuid_value, connected, enabled, min_sep, live_value]) => {
                if (connected && enabled && uuid_value && live_value) {
                    this.websocket?.send(
                        toJsonString({
                            RegisterTlmListener: {
                                uuid: uuid_value,
                                minimum_separation_ms: min_sep,
                            },
                        }),
                    );
                    if (!this.on_telem_value.has(uuid_value)) {
                        this.on_telem_value.set(uuid_value, []);
                    }
                    const callback_fn = (value: TelemetryDataItem) => {
                        value_result.value = value;
                    };
                    this.on_telem_value.get(uuid_value)?.push(callback_fn);
                    onWatcherCleanup(() => {
                        this.websocket?.send(
                            toJsonString({
                                UnregisterTlmListener: {
                                    uuid: uuid_value,
                                },
                            }),
                        );
                        const index = this.on_telem_value
                            .get(uuid_value)
                            ?.indexOf(callback_fn);
                        if (index !== undefined && index >= 0) {
                            this.on_telem_value
                                .get(uuid_value)
                                ?.splice(index, 1);
                        }
                    });
                }
            },
        );

        return value_result;
    }

    pause() {
        this.enabled.value = false;
    }

    resume() {
        this.enabled.value = true;
    }
}

export const WEBSOCKET_SYMBOL = Symbol();

export function useWebsocket() {
    const handle = shallowRef<WebsocketHandle>(new WebsocketHandle());

    onMounted(() => {
        handle.value.connect();
    });
    onDocumentVisibilityChange((visible) => {
        if (visible) {
            handle.value.resume();
        } else {
            handle.value.pause();
        }
    });
    onUnmounted(() => {
        handle.value.disconnect();
    });

    return handle;
}