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

// Frontend wrapper for the Rust `enumerate_screen_sources` command. Falls
// back to an empty list outside the Tauri runtime so a browser-only dev
// build (pnpm vite:dev in Chrome without Tauri) degrades gracefully to
// "nothing to show" rather than throwing.

import { isTauriRuntime } from './globalShortcut';

export type ScreenSourceKind = 'screen' | 'window';

export interface ScreenSource {
  /** Chromium-format source id ("screen:<id>:0" / "window:<hwnd>:0"). */
  id: string;
  name: string;
  kind: ScreenSourceKind;
  /** Base64-encoded JPEG without a data-URL prefix. Null when capture failed.
   *  Kept under `thumbnailPng` key for rollout stability — the server-side
   *  format switched from PNG to JPEG for payload size, but the field name
   *  preserves the wire contract during the transition. */
  thumbnailPng: string | null;
  width: number;
  height: number;
}

// Fast, metadata-only list. The picker uses this first so names show up
// immediately; thumbnails stream in via captureScreenSourceThumbnail below.
export async function listScreenSources(): Promise<ScreenSource[]> {
  if (!isTauriRuntime()) return [];
  try {
    const { invoke } = await import('@tauri-apps/api/core');
    const raw = await invoke<ScreenSource[]>('list_screen_sources');
    return raw ?? [];
  } catch (err: unknown) {
    console.warn('list_screen_sources failed', err);
    return [];
  }
}

// Single-source thumbnail capture (legacy base64 variant). Callers should
// prefer `captureScreenSourceThumbnailBytes` below — it ships raw JPEG
// bytes over IPC so the main thread avoids both the base64 decode AND
// the JSON parse overhead of a long string result. Kept for fallback.
export async function captureScreenSourceThumbnail(
  sourceId: string,
): Promise<string | null> {
  if (!isTauriRuntime()) return null;
  try {
    const { invoke } = await import('@tauri-apps/api/core');
    const result = await invoke<string | null>('capture_screen_source_thumbnail', {
      sourceId,
    });
    return result ?? null;
  } catch (err: unknown) {
    console.warn('capture_screen_source_thumbnail failed', { sourceId, err });
    return null;
  }
}

// Thumbnail fetch with a two-tier fallback:
//   1. binary-IPC path (`..._bytes`) — ArrayBuffer over Tauri's raw channel
//   2. base64 path (legacy) — same command minus the ArrayBuffer wrapper
//
// The binary path can come through in several shapes depending on the
// Tauri / WebView2 version combo: a real ArrayBuffer, a Uint8Array, or
// occasionally a plain number[] when the response got re-serialised.
// We normalise all three into an ArrayBuffer before handing it to Blob.
// If the binary path returns nothing usable we retry once on the base64
// command — keeps thumbnails visible while the binary contract settles.
let warnedBinaryShape = false;
export async function captureScreenSourceThumbnailBytes(
  sourceId: string,
): Promise<Blob | null> {
  if (!isTauriRuntime()) return null;
  const { invoke } = await import('@tauri-apps/api/core');

  // ---- Tier 1: binary IPC ---------------------------------------------
  try {
    const result = await invoke<ArrayBuffer | Uint8Array | number[] | null>(
      'capture_screen_source_thumbnail_bytes',
      { sourceId },
    );
    let bytes: Uint8Array | null = null;
    if (result instanceof ArrayBuffer) {
      bytes = new Uint8Array(result);
    } else if (result instanceof Uint8Array) {
      bytes = result;
    } else if (Array.isArray(result) && result.length > 0) {
      bytes = new Uint8Array(result);
    } else if (result && typeof result === 'object') {
      // One-time diagnostic so we can see the unexpected shape in the
      // console if WebView2 de-serialises the Response body into a bag
      // of properties instead of a transferable binary buffer.
      if (!warnedBinaryShape) {
        warnedBinaryShape = true;
        console.warn(
          'capture_screen_source_thumbnail_bytes: unexpected shape, falling back',
          result,
        );
      }
    }
    if (bytes && bytes.byteLength > 0) {
      const copy = new Uint8Array(bytes.byteLength);
      copy.set(bytes);
      return new Blob([copy.buffer], { type: 'image/jpeg' });
    }
  } catch (err: unknown) {
    console.warn('binary thumbnail path threw, trying base64 fallback', {
      sourceId,
      err,
    });
  }

  // ---- Tier 2: base64 fallback ----------------------------------------
  try {
    const b64 = await invoke<string | null>('capture_screen_source_thumbnail', {
      sourceId,
    });
    if (!b64) return null;
    const bin = atob(b64);
    const fallbackBytes = new Uint8Array(bin.length);
    for (let i = 0; i < bin.length; i++) fallbackBytes[i] = bin.charCodeAt(i);
    return new Blob([fallbackBytes.buffer], { type: 'image/jpeg' });
  } catch (err: unknown) {
    console.warn('base64 thumbnail fallback failed', { sourceId, err });
    return null;
  }
}

// Legacy single-shot variant. Captures everything serially on the Rust side
// before returning. Prefer listScreenSources + captureScreenSourceThumbnail
// for user-facing flows — they feel 5–10× more responsive in practice.
export async function enumerateScreenSources(): Promise<ScreenSource[]> {
  if (!isTauriRuntime()) return [];
  try {
    const { invoke } = await import('@tauri-apps/api/core');
    const raw = await invoke<ScreenSource[]>('enumerate_screen_sources');
    return raw ?? [];
  } catch (err: unknown) {
    console.warn('enumerate_screen_sources failed', err);
    return [];
  }
}

// Data-URL helper — picker tiles bind `src={thumbnailDataUrl(src)}` so the
// thumbnail bytes never leave the component's render pass. Rust encodes
// JPEG now (smaller payload, faster decode); the mime type here must
// match or the <img> element silently fails to paint.
export function thumbnailDataUrl(src: ScreenSource): string | null {
  if (!src.thumbnailPng) return null;
  return 'data:image/jpeg;base64,' + src.thumbnailPng;
}