zep/src/useWebRTC.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { Identity } from 'spacetimedb';
import { tables } from './module_bindings';
import { useTable } from 'spacetimedb/react';

const [offers] = useTable(tables.sdp_offer);
const [answers] = useTable(tables.sdp_answer);
const [candidates] = useTable(tables.ice_candidate);

export type PeerConnection = {
  identity: Identity;
  connection: RTCPeerConnection;
  stream?: MediaStream;
  connectionState: RTCPeerConnectionState;
  iceConnectionState: RTCIceConnectionState;
};

export function useWebRTC(
  identity: Identity | null,
  currentChannelId: bigint | null,
  voiceStates: any[]
) {
  const [localStream, setLocalStream] = useState<MediaStream | null>(null);
  const [remotePeers, setRemotePeers] = useState<Map<string, PeerConnection>>(new Map());

  const peersRef = useRef<Map<string, PeerConnection>>(new Map());
  const localStreamRef = useRef<MediaStream | null>(null);

  const iceConfig = {
    iceServers: [
      { urls: 'stun:stun.l.google.com:19302' },
      { urls: 'stun:stun1.l.google.com:19302' },
    ],
  };

  const cleanupPeer = useCallback((peerId: string) => {
    const peer = peersRef.current.get(peerId);
    if (peer) {
      peer.connection.close();
      peersRef.current.delete(peerId);
      setRemotePeers(new Map(peersRef.current));
    }
  }, []);

  const createPeerConnection = useCallback((peerIdentity: Identity, isOfferer: boolean) => {
    const peerId = peerIdentity.toHexString();
    if (peersRef.current.has(peerId)) return peersRef.current.get(peerId)!.connection;

    const pc = new RTCPeerConnection(iceConfig);

    if (localStreamRef.current) {
      localStreamRef.current.getTracks().forEach(track => {
        pc.addTrack(track, localStreamRef.current!);
      });
    }

    pc.onicecandidate = (event) => {
      if (event.candidate && currentChannelId) {
        try {
          sendIceCandidate({
            receiver: peerIdentity,
            candidate: JSON.stringify(event.candidate),
            channelId: currentChannelId
          });
        } catch (e) {
          console.error("Error sending ICE candidate:", e);
        }
      }
    };

    pc.ontrack = (event) => {
      console.log("Received remote track from", peerId);
      const peer = peersRef.current.get(peerId);
      if (peer) {
        peer.stream = event.streams[0];
        setRemotePeers(new Map(peersRef.current));
      }
    };

    pc.onconnectionstatechange = () => {
      console.log(`Connection state change for ${peerId}: ${pc.connectionState}`);
      const updatedPeer = peersRef.current.get(peerId);
      if (updatedPeer) {
        updatedPeer.connectionState = pc.connectionState;
        setRemotePeers(new Map(peersRef.current));
      }
    };

    pc.oniceconnectionstatechange = () => {
      console.log(`ICE connection state change for ${peerId}: ${pc.iceConnectionState}`);
      const updatedPeer = peersRef.current.get(peerId);
      if (updatedPeer) {
        updatedPeer.iceConnectionState = pc.iceConnectionState;
        setRemotePeers(new Map(peersRef.current));
      }
    };

    const peerObj: PeerConnection = { identity: peerIdentity, connection: pc, connectionState: pc.connectionState, iceConnectionState: pc.iceConnectionState };
    peersRef.current.set(peerId, peerObj);
    setRemotePeers(new Map(peersRef.current));

    return pc;
  }, [currentChannelId]);

  useEffect(() => {
    if (currentChannelId) {
      navigator.mediaDevices.getUserMedia({ audio: true, video: false })
        .then(stream => {
          setLocalStream(stream);
          localStreamRef.current = stream;
        })
        .catch(err => console.error("Error getting audio stream:", err));
    } else {
      if (localStreamRef.current) {
        localStreamRef.current.getTracks().forEach(t => t.stop());
        localStreamRef.current = null;
        setLocalStream(null);
      }
      peersRef.current.forEach((_, id) => cleanupPeer(id));
    }
  }, [currentChannelId, cleanupPeer]);

  useEffect(() => {
    if (!currentChannelId || !identity || !localStream) { // Ensure localStream is available before proceeding
      return;
    }

    voiceStates.forEach(async (vs) => {
      if (vs.channelId === currentChannelId && !vs.identity.isEqual(identity)) {
        const peerId = vs.identity.toHexString();
        if (!peersRef.current.has(peerId)) {
          if (identity.toHexString() > peerId) { // Polite peer logic
            const pc = createPeerConnection(vs.identity, true);
            try {
              const offer = await pc.createOffer();
              await pc.setLocalDescription(offer);
              sendSdpOffer({
                receiver: vs.identity,
                sdp: JSON.stringify(offer),
                channelId: currentChannelId
              });
            } catch (e) {
              console.error("Error creating or sending offer:", e);
            }
          }
        }
      }
    });

    peersRef.current.forEach((_, id) => {
      if (!voiceStates.some(vs => vs.identity.toHexString() === id && vs.channelId === currentChannelId)) {
        cleanupPeer(id);
      }
    });
  }, [voiceStates, currentChannelId, identity, localStream, createPeerConnection, cleanupPeer]);

  useEffect(() => {
    if (!currentChannelId || !identity || !localStream) { // Ensure localStream is available before proceeding
      return;
    }

    offers.forEach(async (offerRow) => {
      if (offerRow.receiver.isEqual(identity) && offerRow.channelId === currentChannelId) {
        const peerId = offerRow.sender.toHexString();
        if (!peersRef.current.has(peerId)) {
          const pc = createPeerConnection(offerRow.sender, false);
          try {
            const offer = JSON.parse(offerRow.sdp);
            await pc.setRemoteDescription(new RTCSessionDescription(offer));
            const answer = await pc.createAnswer();
            await pc.setLocalDescription(answer);
            sendSdpAnswer({
              receiver: offerRow.sender,
              sdp: JSON.stringify(answer),
              channelId: currentChannelId
            });
          } catch (e) {
            console.error("Error creating or sending answer:", e);
          }
        }
      }
    });
  }, [offers, currentChannelId, identity, localStream, createPeerConnection]);

  useEffect(() => {
    if (!currentChannelId || !identity) return;

    candidates.forEach(async (candRow) => {
      if (candRow.receiver.isEqual(identity) && candRow.channelId === currentChannelId) {
        const peerId = candRow.sender.toHexString();
        const peer = peersRef.current.get(peerId);
        if (peer && peer.connection.remoteDescription) {
          try {
            const candidate = JSON.parse(candRow.candidate);
            await peer.connection.addIceCandidate(new RTCIceCandidate(candidate));
          } catch (e) {
            console.error("Error adding ICE candidate:", e);
          }
        }
      }
    });
  }, [candidates, currentChannelId, identity]);

  return { localStream, remotePeers };
}