behind_glass/src/depth/AiDepthProvider.ts

import * as THREE from 'three';
import { DepthProcessor } from './DepthProcessor';
import type { DepthProvider, AiDepthRun } from '../types/depth';

type WorkerReply =
  | { type: 'status'; id: number; message: string }
  | { type: 'result'; id: number; width: number; height: number; values: ArrayBuffer; normalizedValues: ArrayBuffer; run: AiDepthRun }
  | { type: 'fallback'; id: number; preparationMs: number; inferenceMs: number }
  | { type: 'error'; id: number; message: string };

type PendingRequest = {
  id: number;
  started: number;
  image: HTMLImageElement;
  fallbackAttempted: boolean;
  preparationMs: number;
  inferenceMs: number;
  resolve: (texture: THREE.Texture) => void;
  reject: (error: Error) => void;
};

export class AiDepthProvider implements DepthProvider {
  private worker: Worker | null = null;
  private pending: PendingRequest | null = null;
  private requestId = 0;

  constructor(private readonly onStatus: (message: string) => void) {}

  async getDepthMap(image: HTMLImageElement): Promise<THREE.Texture> {
    if (this.pending) throw new Error('AI depth 추론이 이미 진행 중입니다.');
    const id = ++this.requestId;
    return new Promise<THREE.Texture>((resolve, reject) => {
      this.pending = { id, started: performance.now(), image, fallbackAttempted: false, preparationMs: 0, inferenceMs: 0, resolve, reject };
      void this.captureAndSend(image, id).catch((error: unknown) => {
        this.fail(id, error instanceof Error ? error : new Error('AI depth 생성에 실패했습니다.'));
      });
    });
  }

  cancel(): void {
    const pending = this.pending;
    this.pending = null;
    this.requestId += 1;
    this.worker?.terminate();
    this.worker = null;
    pending?.reject(new DOMException('AI depth 요청이 취소되었습니다.', 'AbortError'));
  }

  private async captureAndSend(image: HTMLImageElement, id: number, forceWasm = false): Promise<void> {
    const width = image.naturalWidth || image.width;
    const height = image.naturalHeight || image.height;
    if (!width || !height) throw new Error('사진 크기를 확인할 수 없습니다.');
    const scale = Math.min(1, 1024 / Math.max(width, height));
    const bitmap = await createImageBitmap(image, {
      resizeWidth: Math.max(1, Math.round(width * scale)),
      resizeHeight: Math.max(1, Math.round(height * scale)),
      resizeQuality: 'high',
    });
    if (this.pending?.id !== id) {
      bitmap.close();
      return;
    }
    try {
      const worker = this.ensureWorker();
      worker.postMessage({ type: 'infer', id, image: bitmap, ...(forceWasm ? { forceWasm: true } : {}) }, [bitmap]);
    } catch (error) {
      bitmap.close();
      throw error;
    }
  }

  private ensureWorker(): Worker {
    if (this.worker) return this.worker;
    const worker = new Worker(new URL('./aiDepth.worker.ts', import.meta.url), { type: 'module' });
    worker.onmessage = (event: MessageEvent<WorkerReply>) => {
      if (this.worker !== worker) return;
      const message = event.data;
      if (this.pending?.id !== message.id) return;
      if (message.type === 'status') {
        this.onStatus(message.message);
      } else if (message.type === 'result') {
        const pending = this.pending;
        this.pending = null;
        // Both result buffers are transferred into this realm. Release the
        // inference runtime before the caller starts the full-screen GL scene.
        // A manual new inference uses cached files in a fresh worker.
        this.resetWorker(worker);
        try {
          const texture = DepthProcessor.fromNormalizedDepth(new Uint8Array(message.values), message.width, message.height);
          // Keep unquantized normalized values for inspection; geometry still uses the existing byte texture.
          texture.userData.normalizedDepth = new Float32Array(message.normalizedValues);
          // Main-thread wall time includes capture, worker messaging, preparation and inference.
          texture.userData.aiDepthRun = { ...message.run, preparationMs: pending.preparationMs + message.run.preparationMs, inferenceMs: pending.inferenceMs + message.run.inferenceMs, totalMs: performance.now() - pending.started } satisfies AiDepthRun;
          pending.resolve(texture);
        } catch (error) {
          pending.reject(error instanceof Error ? error : new Error('AI depth 결과를 읽지 못했습니다.'));
        }
      } else if (message.type === 'fallback') {
        const pending = this.pending;
        this.resetWorker(worker);
        if (pending.fallbackAttempted) {
          this.fail(message.id, new Error('이 기기에서 AI 깊이를 계산하지 못했습니다. 다시 시도하거나 평면 사진으로 볼 수 있어요.'));
          return;
        }
        pending.fallbackAttempted = true;
        pending.preparationMs += message.preparationMs;
        pending.inferenceMs += message.inferenceMs;
        this.onStatus('다른 길을 찾고 있어요. 조금만 더 기다려 주세요…');
        void this.captureAndSend(pending.image, pending.id, true).catch((error: unknown) => {
          this.fail(pending.id, error instanceof Error ? error : new Error('AI depth 생성에 실패했습니다.'));
        });
      } else {
        // A failed runtime can keep a rejected session chain; manual retry must be fresh too.
        this.resetWorker(worker);
        this.fail(message.id, new Error(message.message));
      }
    };
    worker.onerror = (event) => {
      if (this.worker !== worker) return;
      this.resetWorker(worker);
      this.fail(this.pending?.id ?? -1, new Error(event.message || 'AI depth worker 오류가 발생했습니다.'));
    };
    worker.onmessageerror = () => {
      if (this.worker !== worker) return;
      this.resetWorker(worker);
      this.fail(this.pending?.id ?? -1, new Error('AI depth worker 결과를 읽지 못했습니다.'));
    };
    this.worker = worker;
    return worker;
  }

  private fail(id: number, error: Error): void {
    if (this.pending?.id !== id) return;
    const pending = this.pending;
    this.pending = null;
    pending.reject(error);
  }

  private resetWorker(worker: Worker): void {
    worker.terminate();
    if (this.worker === worker) this.worker = null;
  }
}