QueueCube/web/frontend/src/hooks/useScreenShare.ts

import { useState, useEffect, useRef, useCallback } from 'react';
import { API } from '../api/player';

const TARGET_FRAME_RATE = 60;
const MAX_CAPTURE_WIDTH = 1920;
const MAX_CAPTURE_HEIGHT = 1080;
const RECORDER_TIMESLICE_MS = 16;
const MIN_VIDEO_BIT_RATE = 2_500_000;
const MAX_VIDEO_BIT_RATE = 8_000_000;
const VIDEO_BITS_PER_PIXEL = 0.065;
const SOCKET_BUFFER_TARGET_MS = 75;
const MIN_SOCKET_BUFFER_BYTES = 32 * 1024;
const MIN_RESUME_BUFFER_BYTES = 8 * 1024;

interface UseScreenShareResult {
  isScreenSharing: boolean;
  isScreenSharingSupported: boolean;
  toggleScreenShare: () => Promise<void>;
  stopScreenShare: () => void;
}

function getBestSupportedMimeType(hasAudio: boolean): string | undefined {
  // VP8 is substantially cheaper to encode than VP9 on most systems. Prefer it
  // for real-time capture, then use hardware-friendly H.264 where available.
  const mimeTypes = hasAudio ? [
    'video/webm;codecs=vp8,opus',
    'video/mp4;codecs=avc1.42E01E,mp4a.40.2',
    'video/webm;codecs=vp9,opus',
    'video/webm',
    'video/mp4',
  ] : [
    'video/webm;codecs=vp8',
    'video/mp4;codecs=avc1.42E01E',
    'video/webm;codecs=vp9',
    'video/webm',
    'video/mp4',
  ];

  // Find the first supported mimetype
  for (const type of mimeTypes) {
    if (MediaRecorder.isTypeSupported(type)) {
      console.log(`Using mime type: ${type}`);
      return type;
    }
  }

  console.warn('No preferred mime types supported by this browser');
  return undefined;
}

function getTargetVideoBitRate(track: MediaStreamTrack): number {
  const settings = track.getSettings();
  const width = settings.width ?? MAX_CAPTURE_WIDTH;
  const height = settings.height ?? MAX_CAPTURE_HEIGHT;
  const frameRate = Math.min(settings.frameRate ?? TARGET_FRAME_RATE, TARGET_FRAME_RATE);
  const target = width * height * frameRate * VIDEO_BITS_PER_PIXEL;

  return Math.round(Math.min(MAX_VIDEO_BIT_RATE, Math.max(MIN_VIDEO_BIT_RATE, target)));
}

function getSocketBufferLimits(recorder: MediaRecorder) {
  const bytesPerSecond = (recorder.videoBitsPerSecond + recorder.audioBitsPerSecond) / 8;
  const maximum = Math.max(
    MIN_SOCKET_BUFFER_BYTES,
    Math.ceil(bytesPerSecond * SOCKET_BUFFER_TARGET_MS / 1000),
  );

  return {
    maximum,
    resumeAt: Math.max(MIN_RESUME_BUFFER_BYTES, Math.floor(maximum / 4)),
  };
}

export const useScreenShare = (): UseScreenShareResult => {
  const [isScreenSharing, setIsScreenSharing] = useState(false);
  const [isScreenSharingSupported, setIsScreenSharingSupported] = useState(false);
  const screenShareSocketRef = useRef<WebSocket | null>(null);
  const mediaRecorderRef = useRef<MediaRecorder | null>(null);
  const mediaStreamRef = useRef<MediaStream | null>(null);
  const backpressureTimerRef = useRef<number | null>(null);
  const isStartingRef = useRef(false);
  const sessionGenerationRef = useRef(0);

  // Check if screen sharing is supported
  useEffect(() => {
    setIsScreenSharingSupported(
      typeof navigator !== 'undefined' &&
      navigator.mediaDevices !== undefined &&
      typeof navigator.mediaDevices.getDisplayMedia === 'function'
    );
  }, []);

  const stopScreenShare = useCallback(() => {
    sessionGenerationRef.current += 1;
    isStartingRef.current = false;

    if (backpressureTimerRef.current !== null) {
      window.clearInterval(backpressureTimerRef.current);
      backpressureTimerRef.current = null;
    }

    const recorder = mediaRecorderRef.current;
    mediaRecorderRef.current = null;
    if (recorder) {
      recorder.ondataavailable = null;
      recorder.onerror = null;
      if (recorder.state !== 'inactive') {
        recorder.stop();
      }
    }

    const mediaStream = mediaStreamRef.current;
    mediaStreamRef.current = null;
    mediaStream?.getTracks().forEach((track) => {
      track.onended = null;
      track.stop();
    });

    const socket = screenShareSocketRef.current;
    screenShareSocketRef.current = null;
    if (socket) {
      socket.onopen = null;
      socket.onclose = null;
      socket.onerror = null;
      if (socket.readyState === WebSocket.CONNECTING || socket.readyState === WebSocket.OPEN) {
        socket.close();
      }
    }

    setIsScreenSharing(false);
  }, []);

  const startScreenShare = useCallback(async () => {
    if (isStartingRef.current || screenShareSocketRef.current) {
      return;
    }

    isStartingRef.current = true;
    const sessionGeneration = ++sessionGenerationRef.current;

    try {
      const mediaStream = await navigator.mediaDevices.getDisplayMedia({
        video: {
          frameRate: { ideal: TARGET_FRAME_RATE, max: TARGET_FRAME_RATE },
          width: { max: MAX_CAPTURE_WIDTH },
          height: { max: MAX_CAPTURE_HEIGHT },
        },
        audio: true,
      });

      // getDisplayMedia cannot be programmatically dismissed. If the user
      // stopped or restarted while its picker was open, discard this result.
      if (sessionGenerationRef.current !== sessionGeneration) {
        mediaStream.getTracks().forEach((track) => track.stop());
        return;
      }

      mediaStreamRef.current = mediaStream;

      const videoTrack = mediaStream.getVideoTracks()[0];
      if (!videoTrack) {
        throw new Error('The selected source did not provide a video track');
      }

      // Favor smooth motion over static-text quality when the browser needs to
      // trade capture resolution for frame rate.
      videoTrack.contentHint = 'motion';
      mediaStream.getAudioTracks().forEach((track) => {
        track.contentHint = 'music';
      });

      const hasAudio = mediaStream.getAudioTracks().length > 0;
      const preferredMimeType = getBestSupportedMimeType(hasAudio);
      const recorderOptions: MediaRecorderOptions = {
        videoBitsPerSecond: getTargetVideoBitRate(videoTrack),
      };

      if (preferredMimeType) {
        recorderOptions.mimeType = preferredMimeType;
      }
      if (hasAudio) {
        recorderOptions.audioBitsPerSecond = 128_000;
      }

      const mediaRecorder = new MediaRecorder(mediaStream, recorderOptions);
      mediaRecorderRef.current = mediaRecorder;
      const mimeType = mediaRecorder.mimeType || preferredMimeType || 'application/octet-stream';
      const socketBufferLimits = getSocketBufferLimits(mediaRecorder);
      console.log(
        `Screen sharing at up to ${TARGET_FRAME_RATE} fps, ${Math.round(mediaRecorder.videoBitsPerSecond / 1000)} kbps, using ${mimeType}`,
      );

      // Connect to WebSocket
      const socket = API.startScreenShare(mimeType);
      screenShareSocketRef.current = socket;
      isStartingRef.current = false;

      const resumeRecorderWhenDrained = () => {
        if (screenShareSocketRef.current !== socket || mediaRecorderRef.current !== mediaRecorder) {
          if (backpressureTimerRef.current !== null) {
            window.clearInterval(backpressureTimerRef.current);
            backpressureTimerRef.current = null;
          }
          return;
        }

        if (socket.bufferedAmount <= socketBufferLimits.resumeAt && mediaRecorder.state === 'paused') {
          mediaRecorder.resume();
          if (backpressureTimerRef.current !== null) {
            window.clearInterval(backpressureTimerRef.current);
            backpressureTimerRef.current = null;
          }
        }
      };

      const pauseRecorderForBackpressure = () => {
        if (mediaRecorder.state === 'recording') {
          mediaRecorder.pause();
        }

        if (backpressureTimerRef.current === null) {
          backpressureTimerRef.current = window.setInterval(
            resumeRecorderWhenDrained,
            RECORDER_TIMESLICE_MS,
          );
        }
      };

      // Set up WebSocket event handlers
      socket.onopen = () => {
        if (screenShareSocketRef.current !== socket) {
          socket.close();
          return;
        }

        console.log('Screen sharing WebSocket connected');
        setIsScreenSharing(true);
        mediaRecorder.start(RECORDER_TIMESLICE_MS);
      };

      socket.onclose = () => {
        console.log('Screen sharing WebSocket closed');
        if (screenShareSocketRef.current === socket) {
          stopScreenShare();
        }
      };

      socket.onerror = (error) => {
        console.error('Screen sharing WebSocket error:', error);
        if (screenShareSocketRef.current === socket) {
          stopScreenShare();
        }
      };

      mediaRecorder.onerror = (error) => {
        console.error('Screen sharing recorder error:', error);
        if (mediaRecorderRef.current === mediaRecorder) {
          stopScreenShare();
        }
      };

      // Send data over WebSocket when available
      mediaRecorder.ondataavailable = (event) => {
        if (event.data.size > 0 && screenShareSocketRef.current === socket && socket.readyState === WebSocket.OPEN) {
          socket.send(event.data);

          // Browser WebSockets otherwise queue without a limit, turning a brief
          // network slowdown into permanently stale playback. Pausing capture
          // bounds that queue and resumes at the live edge once it drains.
          if (socket.bufferedAmount >= socketBufferLimits.maximum) {
            pauseRecorderForBackpressure();
          }
        }
      };

      // Handle stream ending (user clicks "Stop sharing")
      videoTrack.onended = stopScreenShare;

    } catch (error) {
      if (sessionGenerationRef.current === sessionGeneration) {
        console.error('Error starting screen share:', error);
        stopScreenShare();
      }
    }
  }, [stopScreenShare]);

  useEffect(() => {
    return () => stopScreenShare();
  }, [stopScreenShare]);

  const toggleScreenShare = useCallback(async () => {
    if (screenShareSocketRef.current || isStartingRef.current) {
      stopScreenShare();
    } else {
      await startScreenShare();
    }
  }, [startScreenShare, stopScreenShare]);

  return {
    isScreenSharing,
    isScreenSharingSupported,
    toggleScreenShare,
    stopScreenShare
  };
};