ChatApp/apps/desktop/src/lib/screenCapture.ts

// Frontend side of the native screen-capture pipeline. Starts a Rust-side
// capture thread via `start_screen_capture` and streams JPEG frames back
// through a Tauri Channel. Each frame is decoded into an ImageBitmap,
// drawn onto an offscreen canvas, and the canvas' captureStream() is
// returned as a MediaStream that LiveKit can publishTrack() directly —
// no OS/browser screen picker is involved.
//
// Video-only: system audio would require WASAPI / ScreenCaptureKit hooks
// that xcap doesn't provide. Callers that request shared audio must
// either fall back to the browser picker or accept video-without-audio.

import { isTauriRuntime } from './globalShortcut';

export interface NativeCaptureHandle {
  /** Rust-side capture id. Pass to `stopNativeCapture` to tear down. */
  captureId: number;
  /** MediaStream fed by a canvas that's drawing each incoming frame. */
  stream: MediaStream;
  /** Cleanup — stops the Rust thread, closes channels, revokes the canvas
   *  stream. Idempotent. */
  stop: () => Promise<void>;
}

interface FramePayload {
  captureId: number;
  width: number;
  height: number;
  jpegBase64: string;
}

/** Returned when the runtime can't support native capture (no Tauri, no
 *  WebAudio, source vanished between enumeration and start, etc.). The
 *  caller is expected to fall back to the browser's getDisplayMedia path. */
export class NativeCaptureUnavailable extends Error {
  constructor(reason: string) {
    super('native capture unavailable: ' + reason);
    this.name = 'NativeCaptureUnavailable';
  }
}

export async function startNativeCapture(opts: {
  sourceId: string;
  maxWidth: number;
  maxHeight: number;
  fps: number;
}): Promise<NativeCaptureHandle> {
  if (!isTauriRuntime()) {
    throw new NativeCaptureUnavailable('not a tauri runtime');
  }

  const canvas = document.createElement('canvas');
  canvas.width = opts.maxWidth;
  canvas.height = opts.maxHeight;
  const ctx = canvas.getContext('2d');
  if (!ctx) {
    throw new NativeCaptureUnavailable('canvas 2d context unavailable');
  }

  // Track whether we got the first frame so we can fail fast if Rust
  // reports "found" but then produces no output (e.g. screen was locked).
  let firstFrameResolved = false;
  let firstFrameResolve!: () => void;
  let firstFrameReject!: (err: Error) => void;
  const firstFramePromise = new Promise<void>((resolve, reject) => {
    firstFrameResolve = resolve;
    firstFrameReject = reject;
  });

  const { Channel, invoke } = await import('@tauri-apps/api/core');
  const channel = new Channel<FramePayload>();

  // Latest-wins frame queue: if the JS side falls behind the Rust producer,
  // we drop stale frames rather than queue them. Keeps memory flat and
  // latency sensible for live screenshare.
  let pendingFrame: FramePayload | null = null;
  let decoding = false;

  const drainQueue = async () => {
    if (decoding) return;
    decoding = true;
    try {
      while (pendingFrame) {
        const frame = pendingFrame;
        pendingFrame = null;
        const bytes = base64ToBytes(frame.jpegBase64);
        // Uint8Array's buffer type is `ArrayBufferLike` (could be a
        // SharedArrayBuffer in theory); Blob wants plain ArrayBuffer.
        // Pass the underlying buffer explicitly so the type narrows.
        const blob = new Blob([bytes.buffer as ArrayBuffer], { type: 'image/jpeg' });
        let bitmap: ImageBitmap;
        try {
          bitmap = await createImageBitmap(blob);
        } catch (err: unknown) {
          console.warn('createImageBitmap failed', err);
          continue;
        }
        if (canvas.width !== bitmap.width || canvas.height !== bitmap.height) {
          canvas.width = bitmap.width;
          canvas.height = bitmap.height;
        }
        ctx.drawImage(bitmap, 0, 0);
        bitmap.close();
        if (!firstFrameResolved) {
          firstFrameResolved = true;
          firstFrameResolve();
        }
      }
    } finally {
      decoding = false;
    }
  };

  channel.onmessage = (frame: FramePayload) => {
    pendingFrame = frame;
    void drainQueue();
  };

  let captureId: number;
  try {
    captureId = await invoke<number>('start_screen_capture', {
      sourceId: opts.sourceId,
      maxWidth: opts.maxWidth,
      maxHeight: opts.maxHeight,
      fps: opts.fps,
      channel,
    });
  } catch (err: unknown) {
    throw new NativeCaptureUnavailable(
      err instanceof Error ? err.message : String(err),
    );
  }

  // Bound the wait: the capture thread may fail silently on some sources
  // (locked screens, protected windows). Fall back to getDisplayMedia in
  // that case rather than hang the user.
  const firstFrameTimeout = window.setTimeout(() => {
    firstFrameReject(new Error('first frame timed out (3s)'));
  }, 3000);
  try {
    await firstFramePromise;
  } catch (err: unknown) {
    window.clearTimeout(firstFrameTimeout);
    try {
      await invoke('stop_screen_capture', { captureId });
    } catch {
      /* ignore */
    }
    throw new NativeCaptureUnavailable(
      err instanceof Error ? err.message : String(err),
    );
  }
  window.clearTimeout(firstFrameTimeout);

  const stream = canvas.captureStream(opts.fps);

  let stopped = false;
  const stop = async (): Promise<void> => {
    if (stopped) return;
    stopped = true;
    try {
      await invoke('stop_screen_capture', { captureId });
    } catch (err: unknown) {
      console.warn('stop_screen_capture failed', err);
    }
    for (const track of stream.getTracks()) {
      try {
        track.stop();
      } catch {
        /* already stopped */
      }
    }
  };

  return { captureId, stream, stop };
}

function base64ToBytes(b64: string): Uint8Array {
  const bin = atob(b64);
  const bytes = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) {
    bytes[i] = bin.charCodeAt(i);
  }
  return bytes;
}