mrp-qrcode/components/StepViewer.tsx

"use client";

// In-browser STEP viewer. Uses `occt-import-js` (OpenCascade compiled to
// WASM) to triangulate STEP/STP into mesh buffers, then renders them with
// three.js and an OrbitControls camera.
//
// The WASM blob is served from /vendor/occt-import-js.wasm — copied there
// by scripts/copy-viewer-assets.mjs during predev/prebuild.
//
// This component is strictly client-side: both `three` and `occt-import-js`
// reach for `window`/`WebGL`, so we never import them at module top-level in
// server components. All imports inside here are fine because the file is
// marked "use client" and only ever mounted below a client boundary.

import { useEffect, useRef, useState } from "react";

interface Props {
  /** Fully qualified URL that returns the STEP bytes (e.g. /api/v1/files/:id/download). */
  url: string;
  /** Optional class name for the outer wrapper. */
  className?: string;
  /** Show mesh edges for readability. Defaults to true. Disable on very complex assemblies. */
  showEdges?: boolean;
  /**
   * Called once, after the first successful render, with a downsampled JPEG
   * thumbnail of the 3D view (~480x360). Host is responsible for uploading
   * it. Errors in the host callback are swallowed — the viewer still runs.
   */
  onFirstFrame?: (blob: Blob) => void | Promise<void>;
}

type ViewerState =
  | { kind: "idle" }
  | { kind: "loading"; message: string }
  | { kind: "ready"; triangleCount: number; meshCount: number }
  | { kind: "error"; message: string };

export default function StepViewer({ url, className, showEdges = true, onFirstFrame }: Props) {
  const containerRef = useRef<HTMLDivElement | null>(null);
  const [state, setState] = useState<ViewerState>({ kind: "idle" });

  // Hold the latest callback in a ref so updates don't re-trigger the setup
  // effect (which would reload the STEP file every render).
  const onFirstFrameRef = useRef(onFirstFrame);
  useEffect(() => {
    onFirstFrameRef.current = onFirstFrame;
  }, [onFirstFrame]);

  useEffect(() => {
    const hostNullable = containerRef.current;
    if (!hostNullable) return;
    // Capture in a non-nullable local so the async closure below retains the
    // narrowing (TS flow analysis doesn't carry narrowing across async fns).
    const host: HTMLDivElement = hostNullable;

    let cancelled = false;
    let cleanup: (() => void) | null = null;

    async function run() {
      setState({ kind: "loading", message: "Initialising viewer…" });
      try {
        // Dynamically import the heavy bits — keeps initial JS small and
        // means nothing in here ever runs on the server.
        const [{ default: occtFactory }, THREE, { OrbitControls }] =
          await Promise.all([
            import("occt-import-js"),
            import("three"),
            import("three/examples/jsm/controls/OrbitControls.js"),
          ]);

        if (cancelled) return;
        setState({ kind: "loading", message: "Fetching model…" });

        const response = await fetch(url, { cache: "force-cache" });
        if (!response.ok) {
          throw new Error(`Download failed (${response.status})`);
        }
        const arrayBuffer = await response.arrayBuffer();
        if (cancelled) return;

        setState({ kind: "loading", message: "Parsing STEP file…" });

        const occt = await occtFactory({
          locateFile: (file) => `/vendor/${file}`,
        });
        if (cancelled) return;

        const fileBuffer = new Uint8Array(arrayBuffer);
        const result = occt.ReadStepFile(fileBuffer, null);
        if (!result.success || result.meshes.length === 0) {
          throw new Error(
            result.success
              ? "No geometry found in this STEP file."
              : "OpenCascade could not parse this STEP file.",
          );
        }

        if (cancelled) return;
        setState({ kind: "loading", message: "Building scene…" });

        // --- three.js scene -----------------------------------------------
        const width = host.clientWidth || 600;
        const height = host.clientHeight || 400;

        // `preserveDrawingBuffer` so we can capture a thumbnail via toBlob().
        // Slight perf cost; acceptable for the one-at-a-time viewer use case.
        const renderer = new THREE.WebGLRenderer({
          antialias: true,
          alpha: false,
          preserveDrawingBuffer: true,
        });
        renderer.setPixelRatio(window.devicePixelRatio || 1);
        renderer.setSize(width, height, false);
        renderer.setClearColor(0xf8fafc); // slate-50
        host.appendChild(renderer.domElement);
        renderer.domElement.style.display = "block";
        renderer.domElement.style.width = "100%";
        renderer.domElement.style.height = "100%";
        renderer.domElement.style.touchAction = "none";

        const scene = new THREE.Scene();

        const ambient = new THREE.AmbientLight(0xffffff, 0.55);
        scene.add(ambient);
        const key = new THREE.DirectionalLight(0xffffff, 0.9);
        key.position.set(1, 1.5, 1);
        scene.add(key);
        const fill = new THREE.DirectionalLight(0xffffff, 0.35);
        fill.position.set(-1, -0.5, -1);
        scene.add(fill);

        const group = new THREE.Group();
        let totalTris = 0;
        const disposables: Array<{ dispose: () => void }> = [];

        for (const m of result.meshes) {
          const built = buildMesh(THREE, m, showEdges);
          if (!built) continue;
          group.add(built.mesh);
          if (built.edges) group.add(built.edges);
          totalTris += built.triangleCount;
          disposables.push(...built.disposables);
        }

        scene.add(group);

        // Center and frame.
        const bbox = new THREE.Box3().setFromObject(group);
        const size = new THREE.Vector3();
        const center = new THREE.Vector3();
        bbox.getSize(size);
        bbox.getCenter(center);
        group.position.sub(center); // recentre at origin

        const camera = new THREE.PerspectiveCamera(45, width / height, 0.01, 1e7);
        const radius = Math.max(size.x, size.y, size.z, 1);
        const dist = radius * 2.4;
        camera.position.set(dist, dist * 0.8, dist);
        camera.lookAt(0, 0, 0);
        camera.near = radius / 1000;
        camera.far = radius * 1000;
        camera.updateProjectionMatrix();

        const controls = new OrbitControls(camera, renderer.domElement);
        controls.enableDamping = true;
        controls.dampingFactor = 0.08;
        controls.target.set(0, 0, 0);
        controls.update();

        // Render loop.
        let rafId = 0;
        let capturedFirstFrame = false;
        const render = () => {
          controls.update();
          renderer.render(scene, camera);
          rafId = requestAnimationFrame(render);

          // After the first real paint, grab a thumbnail if the host wants
          // one. We wait one extra frame so antialiasing and OrbitControls
          // have settled. Reads through a ref so callback identity changes
          // don't reload the model.
          const cb = onFirstFrameRef.current;
          if (!capturedFirstFrame && cb) {
            capturedFirstFrame = true;
            // Defer a tick so the paint actually commits before we read back.
            requestAnimationFrame(() => {
              try {
                captureThumbnail(renderer.domElement, 480, 360).then((blob) => {
                  if (!cancelled && blob) {
                    Promise.resolve(cb(blob)).catch(() => {
                      /* host failure must not break the viewer */
                    });
                  }
                });
              } catch {
                /* ignore thumbnail capture errors */
              }
            });
          }
        };
        render();

        // Resize.
        const ro = new ResizeObserver(() => {
          const w = host.clientWidth;
          const h = host.clientHeight;
          if (w === 0 || h === 0) return;
          renderer.setSize(w, h, false);
          camera.aspect = w / h;
          camera.updateProjectionMatrix();
        });
        ro.observe(host);

        cleanup = () => {
          cancelAnimationFrame(rafId);
          ro.disconnect();
          controls.dispose();
          for (const d of disposables) {
            try {
              d.dispose();
            } catch {
              /* ignore */
            }
          }
          renderer.dispose();
          if (renderer.domElement.parentNode === host) {
            host.removeChild(renderer.domElement);
          }
        };

        if (cancelled) {
          cleanup();
          cleanup = null;
          return;
        }
        setState({
          kind: "ready",
          triangleCount: totalTris,
          meshCount: result.meshes.length,
        });
      } catch (err) {
        if (cancelled) return;
        const message = err instanceof Error ? err.message : "Unknown viewer error";
        setState({ kind: "error", message });
      }
    }

    run();

    return () => {
      cancelled = true;
      if (cleanup) cleanup();
    };
    // Reload whenever the target file changes.
  }, [url, showEdges]);

  return (
    <div className={className ?? "relative w-full h-[480px] rounded-lg border border-slate-200 bg-slate-50 overflow-hidden"}>
      <div ref={containerRef} className="absolute inset-0" />
      {state.kind !== "ready" ? (
        <div className="pointer-events-none absolute inset-0 flex items-center justify-center">
          <div className="rounded-md bg-white/80 backdrop-blur px-3 py-2 text-sm text-slate-600 border border-slate-200">
            {state.kind === "error" ? (
              <span className="text-red-600">Viewer error: {state.message}</span>
            ) : state.kind === "loading" ? (
              <span>{state.message}</span>
            ) : (
              <span>Preparing viewer…</span>
            )}
          </div>
        </div>
      ) : (
        <div className="absolute bottom-2 right-2 rounded-md bg-white/80 backdrop-blur px-2 py-1 text-xs text-slate-500 border border-slate-200 pointer-events-none">
          {state.meshCount} mesh{state.meshCount === 1 ? "" : "es"} · {state.triangleCount.toLocaleString()} tris
        </div>
      )}
    </div>
  );
}

// --------------------------------------------------------------------------

import type * as THREENS from "three";
import type { OcctMesh } from "occt-import-js";

/**
 * Downsample the viewer canvas to a fixed size and return a JPEG blob.
 * JPEG instead of PNG because thumbnails don't need transparency and JPEG
 * cuts the byte count roughly 4× on typical 3D renders.
 */
function captureThumbnail(
  source: HTMLCanvasElement,
  width: number,
  height: number,
): Promise<Blob | null> {
  return new Promise((resolve) => {
    try {
      const off = document.createElement("canvas");
      off.width = width;
      off.height = height;
      const ctx = off.getContext("2d");
      if (!ctx) return resolve(null);

      // Letterbox if aspect ratios differ — fill with slate-50 to match the
      // viewer background so thumbnails look continuous with the live view.
      ctx.fillStyle = "#f8fafc";
      ctx.fillRect(0, 0, width, height);

      const srcAspect = source.width / source.height;
      const dstAspect = width / height;
      let dw = width;
      let dh = height;
      if (srcAspect > dstAspect) {
        dh = Math.round(width / srcAspect);
      } else {
        dw = Math.round(height * srcAspect);
      }
      const dx = Math.round((width - dw) / 2);
      const dy = Math.round((height - dh) / 2);
      ctx.drawImage(source, dx, dy, dw, dh);

      off.toBlob(
        (blob) => resolve(blob),
        "image/jpeg",
        0.85,
      );
    } catch {
      resolve(null);
    }
  });
}

interface BuiltMesh {
  mesh: THREENS.Mesh;
  edges: THREENS.Group | null;
  triangleCount: number;
  disposables: Array<{ dispose: () => void }>;
}

/**
 * Adapted from occt-import-js/examples/three_viewer.html. Handles the
 * per-BRep-face colouring scheme OCCT emits: a single indexed buffer
 * geometry with one material per face group.
 */
function buildMesh(
  THREE: typeof THREENS,
  source: OcctMesh,
  showEdges: boolean,
): BuiltMesh | null {
  const positionArr = source.attributes.position?.array;
  const indexArr = source.index?.array;
  if (!positionArr || !indexArr || positionArr.length === 0 || indexArr.length === 0) {
    return null;
  }

  const geometry = new THREE.BufferGeometry();
  geometry.setAttribute(
    "position",
    new THREE.Float32BufferAttribute(
      positionArr instanceof Float32Array ? positionArr : Float32Array.from(positionArr as ArrayLike<number>),
      3,
    ),
  );
  if (source.attributes.normal?.array) {
    const normalArr = source.attributes.normal.array;
    geometry.setAttribute(
      "normal",
      new THREE.Float32BufferAttribute(
        normalArr instanceof Float32Array ? normalArr : Float32Array.from(normalArr as ArrayLike<number>),
        3,
      ),
    );
  } else {
    // Fall back to computed normals if the importer didn't supply any.
  }
  geometry.name = source.name ?? "occt-mesh";

  const indexTyped =
    indexArr instanceof Uint32Array ? indexArr : Uint32Array.from(indexArr as ArrayLike<number>);
  geometry.setIndex(new THREE.BufferAttribute(indexTyped, 1));
  if (!source.attributes.normal?.array) {
    geometry.computeVertexNormals();
  }

  const disposables: Array<{ dispose: () => void }> = [geometry];

  const defaultColor = source.color
    ? new THREE.Color(source.color[0], source.color[1], source.color[2])
    : new THREE.Color(0xb0b8c1); // slate-400-ish
  const defaultMaterial = new THREE.MeshPhongMaterial({
    color: defaultColor,
    specular: 0x101010,
    shininess: 30,
    flatShading: false,
  });
  disposables.push(defaultMaterial);

  const materials: THREENS.Material[] = [defaultMaterial];
  const outlineMaterial = new THREE.LineBasicMaterial({ color: 0x334155 });
  disposables.push(outlineMaterial);

  const edges = showEdges ? new THREE.Group() : null;

  const triangleCount = indexTyped.length / 3;

  if (source.brep_faces && source.brep_faces.length > 0) {
    for (const face of source.brep_faces) {
      const c = face.color
        ? new THREE.Color(face.color[0], face.color[1], face.color[2])
        : defaultMaterial.color;
      const mat = new THREE.MeshPhongMaterial({ color: c, specular: 0x101010, shininess: 30 });
      materials.push(mat);
      disposables.push(mat);
    }

    let triangleIndex = 0;
    let faceGroupIndex = 0;
    while (triangleIndex < triangleCount) {
      const firstIndex = triangleIndex;
      let lastIndex: number;
      let materialIndex: number;
      if (faceGroupIndex >= source.brep_faces.length) {
        lastIndex = triangleCount;
        materialIndex = 0;
      } else if (triangleIndex < source.brep_faces[faceGroupIndex].first) {
        lastIndex = source.brep_faces[faceGroupIndex].first;
        materialIndex = 0;
      } else {
        lastIndex = source.brep_faces[faceGroupIndex].last + 1;
        materialIndex = faceGroupIndex + 1;
        faceGroupIndex++;
      }
      geometry.addGroup(firstIndex * 3, (lastIndex - firstIndex) * 3, materialIndex);
      if (edges) {
        const inner = new THREE.BufferGeometry();
        inner.setAttribute("position", geometry.attributes.position);
        if (geometry.attributes.normal) {
          inner.setAttribute("normal", geometry.attributes.normal);
        }
        inner.setIndex(new THREE.BufferAttribute(indexTyped.slice(firstIndex * 3, lastIndex * 3), 1));
        const edgesGeom = new THREE.EdgesGeometry(inner, 30);
        edges.add(new THREE.LineSegments(edgesGeom, outlineMaterial));
        disposables.push(inner, edgesGeom);
      }
      triangleIndex = lastIndex;
    }
  } else if (edges) {
    const edgesGeom = new THREE.EdgesGeometry(geometry, 30);
    edges.add(new THREE.LineSegments(edgesGeom, outlineMaterial));
    disposables.push(edgesGeom);
  }

  const mesh = new THREE.Mesh(
    geometry,
    materials.length > 1 ? materials : materials[0],
  );
  mesh.name = source.name ?? "occt-mesh";
  if (edges) edges.renderOrder = mesh.renderOrder + 1;

  return { mesh, edges, triangleCount, disposables };
}