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", { cause: event }));
            };

            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 } = createPrefixedLogger("%cWebRTC WS%c:", ["color: blue; font-weight: bold;", ""]);