diplom-frontend/app/lib/websocket/voice/client.ts

import { createPrefixedLogger } from "~/lib/utils";
import { ClientMessageType, ConnectionState, ServerMessageType, type ClientMessage, type ServerMessage } from "./types";

export class WebRTCClient {
    private socket: WebSocket | null = null;
    private peerConnection: RTCPeerConnection | null = null;

    private onStateChange: (state: ConnectionState) => void;
    private onError: (error: Error) => void;
    private onRemoteStream: (stream: MediaStream) => void;

    private state: ConnectionState = ConnectionState.DISCONNECTED;
    private url: string;

    private connectionLock = false;
    private disconnectPromise: Promise<void> | null = null;
    private disconnectResolve: (() => void) | null = null;

    constructor(
        url: string,
        onStateChange: (state: ConnectionState) => void,
        onError: (error: Error) => void,
        onRemoteStream: (stream: MediaStream) => void
    ) {
        this.url = url;
        this.onStateChange = onStateChange;
        this.onError = onError;
        this.onRemoteStream = onRemoteStream;
    }

    public connect = async (token: string) => {
        if (this.connectionLock) {
            warn('WebRTC: Connection already in progress');
            return;
        }

        this.connectionLock = true;

        if (this.state !== ConnectionState.DISCONNECTED && this.state !== ConnectionState.ERROR) {
            this.disconnect();
        }

        if (this.disconnectPromise) {
            warn('WebRTC: Waiting for previous disconnect to complete');
            try {
                await this.disconnectPromise;
            } catch (error) {
                console.error('WebRTC: Previous disconnect failed:', error);
            }
        }

        log('Connecting to %s', this.url);

        try {
            this.setState(ConnectionState.CONNECTING);

            this.socket = new WebSocket(this.url);

            this.socket.onopen = () => {
                log('Socket opened');

                this.connectionLock = false;

                this.setState(ConnectionState.AUTHENTICATING);
                this.sendMessage({
                    type: ClientMessageType.AUTHENTICATE,
                    data: { token }
                });
            };

            this.socket.onmessage = this.handleServerMessage;

            this.socket.onerror = (event) => {
                this.handleError(new Error('WebSocket error occurred'));
            };

            this.socket.onclose = (e) => {
                log('Socket closed', e);
                this.cleanupResources();
                if (this.state !== ConnectionState.ERROR) {
                    this.setState(ConnectionState.DISCONNECTED);
                }

                if (this.disconnectResolve) {
                    this.disconnectResolve();
                    this.disconnectResolve = null;
                    this.disconnectPromise = null;
                }
            };
        } catch (error) {
            this.handleError(error instanceof Error ? error : new Error('Unknown error'));
        }
    };

    public disconnect = (): void => {
        if (this.state === ConnectionState.DISCONNECTED) {
            return;
        }

        this.setState(ConnectionState.DISCONNECTING);
        this.connectionLock = false;

        if (this.socket && this.socket.readyState !== WebSocket.CLOSED) {
            // If we're already waiting for a disconnect, cancel it
            if (this.disconnectPromise) {
                this.disconnectResolve = null;
                this.disconnectPromise = null;
            }

            this.disconnectPromise = new Promise((resolve) => {
                this.disconnectResolve = resolve;
            });

            const onSocketClose = () => {
                this.socket?.removeEventListener('close', onSocketClose);
                this.disconnectResolve?.();
                this.disconnectResolve = null;
                this.disconnectPromise = null;
            };

            this.socket.addEventListener('close', onSocketClose);

            if (this.socket.readyState !== WebSocket.CLOSING) {
                this.socket.close(1000, 'WebRTC: Cleaning up resources');
            }
        } else {
            this.cleanupResources();
            this.setState(ConnectionState.DISCONNECTED);
        }
    };

    public createOffer = async (localStream?: MediaStream): Promise<void> => {
        if (this.state !== ConnectionState.CONNECTED) {
            this.handleError(new Error('Cannot create offer: not connected'));
            return;
        }

        try {
            // Create RTCPeerConnection with standard configuration
            const configuration: RTCConfiguration = {
                iceServers: []
            };

            this.peerConnection = new RTCPeerConnection(configuration);

            // Add local stream tracks if provided
            if (localStream) {
                localStream.getTracks().forEach(track => {
                    this.peerConnection!.addTrack(track, localStream);
                });
            }

            // Handle ICE candidates
            this.peerConnection.onicecandidate = (event) => {
                if (event.candidate === null) {
                    // ICE gathering completed
                }
            };

            // Handle remote stream
            this.peerConnection.ontrack = (event) => {
                const [remoteStream] = event.streams;
                if (remoteStream) {
                    this.onRemoteStream(remoteStream);
                }
            };

            // Create offer and set local description
            const offer = await this.peerConnection.createOffer();
            await this.peerConnection.setLocalDescription(offer);

            // Send offer to server
            if (this.peerConnection.localDescription) {
                this.sendMessage({
                    type: ClientMessageType.SDP_OFFER,
                    data: {
                        sdp: this.peerConnection.localDescription
                    }
                });
            }
        } catch (error) {
            this.handleError(error instanceof Error ? error : new Error('Error creating WebRTC offer'));
        }
    };

    private handleServerMessage = async (event: MessageEvent): Promise<void> => {
        try {
            const message: ServerMessage = JSON.parse(event.data);

            log('Received message: %o', message);

            switch (message.type) {
                case ServerMessageType.AUTHENTICATE_ACCEPTED:
                    this.setState(ConnectionState.CONNECTED);
                    break;

                case ServerMessageType.AUTHENTICATE_DENIED:
                    this.handleError(new Error('Authentication failed'));
                    break;

                case ServerMessageType.SDP_ANSWER:
                    await this.handleSdpAnswer(message.data.sdp);
                    break;

                default:
                    warn('Unhandled message type:', message);
            }
        } catch (error) {
            this.handleError(error instanceof Error ? error : new Error('Failed to process message'));
        }
    };

    private handleSdpAnswer = async (sdp: RTCSessionDescription): Promise<void> => {
        log('Received SDP answer: %o', sdp);

        if (!this.peerConnection) {
            this.handleError(new Error('No peer connection established'));
            return;
        }

        try {
            await this.peerConnection.setRemoteDescription(sdp);
        } catch (error) {
            this.handleError(error instanceof Error ? error : new Error('Error setting remote description'));
        }
    };

    private sendMessage = (message: ClientMessage): void => {
        log('Sending message: %o', message);

        if (this.socket && this.socket.readyState === WebSocket.OPEN) {
            this.socket.send(JSON.stringify(message));
        } else {
            this.handleError(new Error('Cannot send message: socket not connected'));
        }
    };

    private setState = (state: ConnectionState): void => {
        log('State changed to %s', state);

        this.state = state;
        this.onStateChange(state);
    };

    private handleError = (error: Error): void => {
        log('Error: %s', error.message);

        this.setState(ConnectionState.ERROR);
        this.onError(error);
    };

    private cleanupResources = (): void => {
        log('Cleaning up resources');

        if (this.peerConnection) {
            this.peerConnection.close();
            this.peerConnection = null;
        }

        if (this.socket) {
            this.socket.close(1000, 'WebRTC: Cleaning up resources');
            this.socket = null;
        }
    };
}

const {
    log,
    warn,
    ...other
} = createPrefixedLogger('%cWebRTC WS%c:', ['color: blue; font-weight: bold;', '']);