sw-fe/src/components/Live/useWebRTCViewer.js

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

const STREAM_SERVER = "https://stream.sociowire.com";
const ICE_SERVERS = [
  { urls: "stun:stun.l.google.com:19302" },
  { urls: "stun:stun1.l.google.com:19302" },
];

/**
 * Hook for WebRTC viewing with join-as-participant capability
 */
export function useWebRTCViewer({ streamID, userID, username, token }) {
  const [isConnected, setIsConnected] = useState(false);
  const [isConnecting, setIsConnecting] = useState(false);
  const [error, setError] = useState(null);

  // Video streams received
  const [mainStream, setMainStream] = useState(null); // Host's main camera
  const [presenterStream, setPresenterStream] = useState(null); // Host's self-view
  const [participantStreams, setParticipantStreams] = useState([]); // Other participants

  // Stream metadata
  const [streamInfo, setStreamInfo] = useState(null);
  const [viewerCount, setViewerCount] = useState(0);

  // Join request state
  const [joinRequestStatus, setJoinRequestStatus] = useState(null); // null, 'pending', 'approved', 'denied'
  const [isParticipant, setIsParticipant] = useState(false);

  const pcRef = useRef(null);
  const wsRef = useRef(null);
  const localStreamRef = useRef(null); // Our camera when we join as participant

  // Connect to event WebSocket
  const connectEventSocket = useCallback(() => {
    if (!streamID) return;

    const protocol = window.location.protocol === "https:" ? "wss" : "ws";
    const wsUrl = `${protocol}://stream.sociowire.com/api/streams/${streamID}/events`;

    try {
      const ws = new WebSocket(wsUrl);

      ws.onopen = () => {
        console.log("[Viewer] Event socket connected");
        ws.send(JSON.stringify({ type: "auth", token, userID, role: "viewer" }));
      };

      ws.onmessage = (event) => {
        try {
          const msg = JSON.parse(event.data);
          handleEventMessage(msg);
        } catch (err) {
          console.error("[Viewer] Failed to parse event:", err);
        }
      };

      ws.onclose = () => {
        console.log("[Viewer] Event socket closed");
      };

      ws.onerror = (err) => {
        console.error("[Viewer] Event socket error:", err);
      };

      wsRef.current = ws;
    } catch (err) {
      console.error("[Viewer] Failed to connect event socket:", err);
    }
  }, [streamID, token, userID]);

  const handleEventMessage = useCallback((msg) => {
    switch (msg.type) {
      case "stream_info":
        setStreamInfo(msg.stream);
        break;

      case "viewer_count":
        setViewerCount(msg.count || 0);
        break;

      case "join_approved":
        setJoinRequestStatus("approved");
        // Now we can send our video
        startParticipating();
        break;

      case "join_denied":
        setJoinRequestStatus("denied");
        break;

      case "participant_joined":
        handleParticipantJoined(msg);
        break;

      case "participant_left":
        handleParticipantLeft(msg.userID);
        break;

      case "host_answer":
        // Host sent us an answer for our participant offer
        handleHostAnswer(msg);
        break;

      case "host_ice":
        handleHostICE(msg);
        break;

      case "stream_ended":
        setError("Stream has ended");
        disconnect();
        break;

      case "kicked":
        setError("You have been removed from the stream");
        setIsParticipant(false);
        stopLocalStream();
        break;

      default:
        console.log("[Viewer] Unknown event:", msg.type);
    }
  }, []);

  const handleParticipantJoined = useCallback((msg) => {
    const { userID: participantID, username: participantName } = msg;

    setParticipantStreams((prev) => [
      ...prev.filter((p) => p.userID !== participantID),
      { userID: participantID, username: participantName, stream: null },
    ]);
  }, []);

  const handleParticipantLeft = useCallback((participantID) => {
    setParticipantStreams((prev) => prev.filter((p) => p.userID !== participantID));
  }, []);

  const handleHostAnswer = useCallback(async (msg) => {
    const { sdp } = msg;
    if (pcRef.current && sdp) {
      try {
        await pcRef.current.setRemoteDescription({ type: "answer", sdp });
        setIsParticipant(true);
        console.log("[Viewer] Now participating with video");
      } catch (err) {
        console.error("[Viewer] Failed to set host answer:", err);
      }
    }
  }, []);

  const handleHostICE = useCallback((msg) => {
    const { candidate } = msg;
    if (pcRef.current && candidate) {
      pcRef.current.addIceCandidate(new RTCIceCandidate(candidate)).catch((err) => {
        console.error("[Viewer] Failed to add host ICE candidate:", err);
      });
    }
  }, []);

  // Connect as viewer (receive only)
  const connect = useCallback(async () => {
    if (isConnecting || isConnected) return;
    setIsConnecting(true);
    setError(null);

    try {
      // Get stream info first
      const infoRes = await fetch(`${STREAM_SERVER}/api/streams/${streamID}`);
      if (!infoRes.ok) {
        throw new Error("Stream not found or has ended");
      }
      const info = await infoRes.json();
      setStreamInfo(info);

      // Create peer connection
      const pc = new RTCPeerConnection({ iceServers: ICE_SERVERS });

      // We need to add transceivers to receive video/audio
      pc.addTransceiver("video", { direction: "recvonly" });
      pc.addTransceiver("video", { direction: "recvonly" }); // For presenter track
      pc.addTransceiver("audio", { direction: "recvonly" });

      pc.ontrack = (event) => {
        const track = event.track;
        const stream = event.streams[0];

        console.log("[Viewer] Received track:", track.kind, stream?.id);

        // Determine which stream this is based on stream ID or order
        // First video stream is main, second is presenter
        if (track.kind === "video") {
          setMainStream((prev) => {
            if (!prev) return stream;
            setPresenterStream(stream);
            return prev;
          });
        }
      };

      pc.onicecandidate = (event) => {
        if (event.candidate && wsRef.current?.readyState === WebSocket.OPEN) {
          wsRef.current.send(
            JSON.stringify({
              type: "ice_candidate",
              candidate: event.candidate,
            })
          );
        }
      };

      pc.onconnectionstatechange = () => {
        console.log("[Viewer] Connection state:", pc.connectionState);
        if (pc.connectionState === "failed" || pc.connectionState === "disconnected") {
          setError("Connection lost");
          setIsConnected(false);
        }
      };

      // Create offer
      const offer = await pc.createOffer();
      await pc.setLocalDescription(offer);

      // Send to server
      const res = await fetch(`${STREAM_SERVER}/api/webrtc/watch/${streamID}`, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          ...(token ? { Authorization: `Bearer ${token}` } : {}),
        },
        body: JSON.stringify({
          sdp: offer.sdp,
          userID,
          username,
        }),
      });

      if (!res.ok) {
        const errData = await res.json().catch(() => ({}));
        throw new Error(errData.error || "Failed to connect to stream");
      }

      const { sdp: answerSDP } = await res.json();
      await pc.setRemoteDescription({ type: "answer", sdp: answerSDP });

      pcRef.current = pc;
      setIsConnected(true);

      // Connect event socket
      connectEventSocket();

      console.log("[Viewer] Connected successfully");
    } catch (err) {
      console.error("[Viewer] Failed to connect:", err);
      setError(err.message || "Failed to connect");
    } finally {
      setIsConnecting(false);
    }
  }, [streamID, userID, username, token, isConnecting, isConnected, connectEventSocket]);

  // Disconnect
  const disconnect = useCallback(() => {
    if (pcRef.current) {
      pcRef.current.close();
      pcRef.current = null;
    }

    if (wsRef.current) {
      wsRef.current.close();
      wsRef.current = null;
    }

    stopLocalStream();

    setIsConnected(false);
    setMainStream(null);
    setPresenterStream(null);
    setParticipantStreams([]);
    setIsParticipant(false);
    setJoinRequestStatus(null);

    console.log("[Viewer] Disconnected");
  }, []);

  // Request to join as participant
  const requestToJoin = useCallback(async () => {
    if (joinRequestStatus === "pending") return;

    try {
      const res = await fetch(
        `${STREAM_SERVER}/api/streams/${streamID}/join-request`,
        {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            ...(token ? { Authorization: `Bearer ${token}` } : {}),
          },
          body: JSON.stringify({ userID, username }),
        }
      );

      if (!res.ok) {
        throw new Error("Failed to send join request");
      }

      setJoinRequestStatus("pending");
    } catch (err) {
      console.error("[Viewer] Failed to request join:", err);
      setError(err.message);
    }
  }, [streamID, userID, username, token, joinRequestStatus]);

  // Cancel join request
  const cancelJoinRequest = useCallback(async () => {
    try {
      await fetch(
        `${STREAM_SERVER}/api/streams/${streamID}/join-request/${userID}`,
        {
          method: "DELETE",
          headers: {
            ...(token ? { Authorization: `Bearer ${token}` } : {}),
          },
        }
      );
      setJoinRequestStatus(null);
    } catch (err) {
      console.error("[Viewer] Failed to cancel join request:", err);
    }
  }, [streamID, userID, token]);

  // Start participating (send our video after approval)
  const startParticipating = useCallback(async () => {
    try {
      // Get our camera
      const localStream = await navigator.mediaDevices.getUserMedia({
        video: { facingMode: "user", width: 640, height: 480 },
        audio: true,
      });
      localStreamRef.current = localStream;

      // Create new peer connection for sending our stream
      // (or reuse existing one with renegotiation)
      if (pcRef.current) {
        // Add our tracks
        localStream.getTracks().forEach((track) => {
          pcRef.current.addTrack(track, localStream);
        });

        // Renegotiate
        const offer = await pcRef.current.createOffer();
        await pcRef.current.setLocalDescription(offer);

        // Send offer to host via WebSocket
        if (wsRef.current?.readyState === WebSocket.OPEN) {
          wsRef.current.send(
            JSON.stringify({
              type: "participant_offer",
              sdp: offer.sdp,
            })
          );
        }
      }
    } catch (err) {
      console.error("[Viewer] Failed to start participating:", err);
      setError("Failed to access camera");
    }
  }, []);

  // Stop local stream
  const stopLocalStream = useCallback(() => {
    if (localStreamRef.current) {
      localStreamRef.current.getTracks().forEach((track) => track.stop());
      localStreamRef.current = null;
    }
  }, []);

  // Leave as participant (become viewer again)
  const leaveAsParticipant = useCallback(async () => {
    stopLocalStream();
    setIsParticipant(false);

    // Notify server
    try {
      await fetch(
        `${STREAM_SERVER}/api/streams/${streamID}/leave`,
        {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            ...(token ? { Authorization: `Bearer ${token}` } : {}),
          },
          body: JSON.stringify({ userID }),
        }
      );
    } catch (err) {
      console.error("[Viewer] Failed to notify leave:", err);
    }
  }, [streamID, userID, token, stopLocalStream]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      disconnect();
    };
  }, []);

  return {
    isConnected,
    isConnecting,
    error,
    mainStream,
    presenterStream,
    participantStreams,
    streamInfo,
    viewerCount,
    joinRequestStatus,
    isParticipant,
    localStream: localStreamRef.current,
    connect,
    disconnect,
    requestToJoin,
    cancelJoinRequest,
    leaveAsParticipant,
  };
}

export default useWebRTCViewer;