DepthScene

"use client";

import { useCallback, useEffect, useRef, useState } from "react";
import { Canvas, useFrame } from "@react-three/fiber";
import * as THREE from "three";

/**
 * DepthScene — on-device WebGPU AI: estimates a depth map from a flat photo with
 * Depth-Anything-V2 (Apache-2.0) running entirely in the browser via
 * Transformers.js, then uses it as a displacement map on a subdivided plane for a
 * real 3D parallax you steer with your cursor. Model loads in a Web Worker only
 * on click. Graceful ladder: WebGPU → WASM → flat image (the visual never fails).
 *
 * Pass your own self-hosted photo via `image`. Honors prefers-reduced-motion
 * (no pointer parallax). Frontier flex: zero-backend, in-browser ML → a 3D primitive.
 */
const DEFAULT_IMAGE = "/demo/scene.jpg";

function ParallaxPlane({
  color,
  depth,
  aspect,
  reduce,
}: {
  color: THREE.Texture;
  depth: THREE.Texture;
  aspect: number;
  reduce: boolean;
}) {
  const mouse = useRef({ x: 0, y: 0 });
  useEffect(() => {
    if (reduce) return;
    const onMove = (e: PointerEvent) => {
      mouse.current.x = (e.clientX / window.innerWidth) * 2 - 1;
      mouse.current.y = -(e.clientY / window.innerHeight) * 2 + 1;
    };
    window.addEventListener("pointermove", onMove);
    return () => window.removeEventListener("pointermove", onMove);
  }, [reduce]);
  useFrame((state) => {
    if (reduce) return;
    state.camera.position.x += (mouse.current.x * 0.4 - state.camera.position.x) * 0.05;
    state.camera.position.y += (mouse.current.y * 0.4 - state.camera.position.y) * 0.05;
    state.camera.lookAt(0, 0, 0);
  });
  return (
    <mesh>
      <planeGeometry args={[3, 3 / aspect, 200, 200]} />
      <meshStandardMaterial
        map={color}
        displacementMap={depth}
        displacementScale={0.6}
        roughness={0.9}
      />
    </mesh>
  );
}

type Status = "idle" | "loading" | "ready" | "flat";

export function DepthScene({ image = DEFAULT_IMAGE }: { image?: string }) {
  const workerRef = useRef<Worker | null>(null);
  const [status, setStatus] = useState<Status>("idle");
  const [progress, setProgress] = useState(0);
  const [reduce, setReduce] = useState(false);
  const [tex, setTex] = useState<{ color: THREE.Texture; depth: THREE.Texture; aspect: number } | null>(
    null
  );

  useEffect(() => {
    setReduce(window.matchMedia("(prefers-reduced-motion: reduce)").matches);
  }, []);

  const run = useCallback(async () => {
    setStatus("loading");
    const nav = navigator as Navigator & { gpu?: { requestAdapter: () => Promise<unknown> } };
    let preferWebGPU = false;
    if (nav.gpu) {
      try {
        preferWebGPU = !!(await nav.gpu.requestAdapter());
      } catch {
        preferWebGPU = false;
      }
    }

    const worker = (workerRef.current ??= new Worker(
      new URL("./depth.worker.ts", import.meta.url),
      { type: "module" }
    ));

    worker.onmessage = (e: MessageEvent) => {
      const m = e.data;
      if (m.type === "progress") {
        // only advance on numeric percents; the worker already filters non-progress events
        if (typeof m.p?.progress === "number") setProgress(m.p.progress);
      } else if (m.type === "error") {
        setStatus("flat");
      } else if (m.type === "done") {
        const { width, height, data } = m as { width: number; height: number; data: Uint8Array };
        const depthTex = new THREE.DataTexture(data, width, height, THREE.RedFormat);
        // DataTexture defaults flipY=false; the color photo (TextureLoader) defaults
        // flipY=true — match them so the relief isn't mirrored vertically.
        depthTex.flipY = true;
        depthTex.needsUpdate = true;
        const colorTex = new THREE.TextureLoader().load(image, () => {
          colorTex.colorSpace = THREE.SRGBColorSpace;
          setTex({ color: colorTex, depth: depthTex, aspect: width / height });
          setStatus("ready"); // gate the Canvas on ready (color image decoded)
        });
      }
    };
    worker.postMessage({ url: image, preferWebGPU });
  }, [image]);

  useEffect(() => () => workerRef.current?.terminate(), []);

  return (
    <div className="relative h-full w-full overflow-hidden rounded-2xl border border-line bg-obsidian-2">
      {status !== "ready" && (
        // eslint-disable-next-line @next/next/no-img-element
        <img src={image} alt="" className="absolute inset-0 h-full w-full object-cover opacity-70" />
      )}
      {status === "idle" && (
        <button
          onClick={run}
          aria-busy={false}
          className="absolute inset-0 z-10 grid place-items-center"
        >
          <span className="rounded-full bg-ink px-5 py-2.5 text-sm font-semibold text-obsidian transition-transform hover:scale-[1.03]">
            Generate depth → enter 3D
          </span>
        </button>
      )}
      {status === "loading" && (
        <div
          role="status"
          aria-live="polite"
          aria-busy="true"
          className="absolute inset-0 z-10 grid place-items-center bg-obsidian/50 text-sm text-mist backdrop-blur-sm"
        >
          <span>
            Running model on-device…{" "}
            <span aria-hidden>{Math.round(progress)}%</span>
          </span>
        </div>
      )}
      {status === "flat" && (
        <div
          role="status"
          aria-live="polite"
          className="absolute bottom-3 left-3 z-10 rounded bg-obsidian/70 px-2 py-1 text-xs text-muted backdrop-blur"
        >
          On-device AI unavailable here — showing the flat source image
        </div>
      )}
      {status === "ready" && tex && (
        <Canvas camera={{ position: [0, 0, 2.2], fov: 45 }} dpr={[1, 2]}>
          <ambientLight intensity={0.85} />
          <directionalLight position={[2, 2, 3]} intensity={1.1} />
          <ParallaxPlane color={tex.color} depth={tex.depth} aspect={tex.aspect} reduce={reduce} />
        </Canvas>
      )}
    </div>
  );
}

// Web Worker: runs monocular depth estimation fully on-device.
// Device ladder: WebGPU fp16 → WebGPU fp32 → WASM q8. Loaded only when the main
// thread posts a message (i.e. on user click), so zero AI bytes on page load.
import { pipeline, env } from "@huggingface/transformers";

env.allowLocalModels = false; // pull ONNX weights from the HF CDN

const MODEL = "onnx-community/depth-anything-v2-small"; // Apache-2.0
const LADDER: { device: "webgpu" | "wasm"; dtype: "fp16" | "fp32" | "q8" }[] = [
  { device: "webgpu", dtype: "fp16" },
  { device: "webgpu", dtype: "fp32" },
  { device: "wasm", dtype: "q8" },
];

type Estimator = (
  input: string
) => Promise<{ depth: { data: Uint8Array; width: number; height: number } }>;

// Typed message port without pulling in the webworker lib (keeps tsc on dom lib).
const port = globalThis as unknown as {
  postMessage: (message: unknown, transfer?: Transferable[]) => void;
  onmessage: ((e: MessageEvent) => void) | null;
};

let estimator: Estimator | null = null;

async function getEstimator(preferWebGPU: boolean): Promise<Estimator> {
  if (estimator) return estimator;
  let lastErr: unknown;
  for (let i = preferWebGPU ? 0 : 2; i < LADDER.length; i++) {
    const cfg = LADDER[i];
    try {
      port.postMessage({ type: "status", msg: `loading (${cfg.device}/${cfg.dtype})` });
      estimator = (await pipeline("depth-estimation", MODEL, {
        device: cfg.device,
        dtype: cfg.dtype,
        // forward only events that carry a numeric percent, so the bar doesn't
        // snap back to 0 on non-download statuses (initiate/done/ready).
        progress_callback: (p: unknown) => {
          if (typeof (p as { progress?: number }).progress === "number") {
            port.postMessage({ type: "progress", p });
          }
        },
      })) as unknown as Estimator;
      return estimator;
    } catch (e) {
      lastErr = e;
      estimator = null;
    }
  }
  throw lastErr ?? new Error("no backend available");
}

port.onmessage = async (e: MessageEvent) => {
  const { url, preferWebGPU } = e.data as { url: string; preferWebGPU: boolean };
  try {
    const est = await getEstimator(preferWebGPU);
    const out = await est(url);
    const depth = out.depth;
    port.postMessage(
      { type: "done", width: depth.width, height: depth.height, data: depth.data },
      [depth.data.buffer]
    );
  } catch (err) {
    port.postMessage({ type: "error", error: String(err) });
  }
};