src/hooks/useWebRTC.ts

import { useState, useRef, useCallback, useEffect } from "react";

import { ICreateConversationResponse } from "../interfaces/conversation";
import { createConnection, createConversation } from "../services/api/chatService";
import { conversationWebSocket } from "../services/websockets/conversation";

export const useWebRTC = () => {
  const [isConnected, setIsConnected] = useState(false);
  const [transcript, setTranscript] = useState("");
  const [remoteStream, setRemoteStream] = useState<MediaStream | null>(null);
  const [callStatus, setCallStatus] = useState("");
  const peerConnectionRef = useRef<RTCPeerConnection | null>(null);
  const dataChannelRef = useRef<RTCDataChannel | null>(null);
  const socketRef = useRef<WebSocket | null>(null);
  const [error, setError] = useState<string | null>(null);
  const mediaStreamRef = useRef<MediaStream | null>(null);
  const isConnectedRef = useRef(isConnected);

  useEffect(() => {
    isConnectedRef.current = isConnected;
  }, [isConnected]);

  const endCall = useCallback(() => {
    if (peerConnectionRef.current) {
      peerConnectionRef.current.close();
      peerConnectionRef.current = null;
    }
    if (socketRef.current) {
      socketRef.current.close();
      socketRef.current = null;
    }
    if (mediaStreamRef.current) {
      mediaStreamRef.current.getTracks().forEach((track) => track.stop());
      mediaStreamRef.current = null;
    }
    if (dataChannelRef.current) {
      dataChannelRef.current.close();
      dataChannelRef.current = null;
    }
    setIsConnected(false);
    setCallStatus("Disconnected");
    setRemoteStream(null);
    setTranscript("");
  }, []);

  const startCall = useCallback(async () => {
    try {
      const peerConnection = new RTCPeerConnection();
      peerConnectionRef.current = peerConnection;
      peerConnection.onconnectionstatechange = () => {
        switch (peerConnection.connectionState) {
          case "connected":
            setIsConnected(true);
            setCallStatus("Connected");
            break;
          case "disconnected":
          case "failed":
          case "closed":
            endCall();
            break;
          default:
            break;
        }
      };
      peerConnection.ontrack = (event) => {
        setRemoteStream(event.streams[0]);
      };

      const mediaStream = await navigator.mediaDevices.getUserMedia({ audio: true });
      mediaStreamRef.current = mediaStream;
      peerConnection.addTrack(mediaStream.getTracks()[0], mediaStream);
      const dataChannel = peerConnection.createDataChannel("response");
      dataChannelRef.current = dataChannel;
      dataChannel.onmessage = (event) => {
        if (socketRef.current?.readyState === WebSocket.OPEN) {
          socketRef.current.send(event.data);
        }
      };

      const sessionResponse = await createConversation({ modality: "voice" });
      const sessionData: ICreateConversationResponse = await sessionResponse.data;
      const conversationId = sessionData.conversation_id;

      const offer = await peerConnection.createOffer();
      await peerConnection.setLocalDescription(offer);
      const webrtcResponse = await createConnection(conversationId, {
        conversation_id: conversationId,
        offer: { sdp: offer.sdp, type: offer.type },
      });

      const responseData = await webrtcResponse.data;
      const ephemeralKey = responseData.ephemeral_key;
      const socket = conversationWebSocket({
        conversationId,
        modality: "voice",
      });

      socketRef.current = socket;
      socket.onopen = () => {
        socket.send(
          JSON.stringify({
            type: "headers",
            headers: { Authorization: `Bearer ${ephemeralKey}` },
          })
        );
      };

      socket.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          if (data.transcript) {
            setTranscript((prev) => `${prev}\n${data.transcript}`);
          } else if (dataChannelRef.current?.readyState === "open") {
            dataChannelRef.current.send(JSON.stringify(data));
          }
        } catch {
          setError("Error handling message");
        }
      };

      socket.onclose = (_event) => {
        if (isConnectedRef.current) endCall();
      };
      await peerConnection.setRemoteDescription(new RTCSessionDescription(responseData.answer));
    } catch {
      setError("Call setup failed");
      endCall();
    }
  }, [endCall]);

  return {
    startCall,
    endCall,
    isConnected,
    transcript,
    remoteStream,
    callStatus,
    error,
  };
};