app/static/viewer-core.js

// viewer-core.js — three.js scene harness for generated img2threejs factories.
// Bundled into every per-job model bundle by esbuild (three included), and
// called with the generated factory exports. Runs inside the sandboxed
// viewer iframe (and inside standalone.html exports).
//
// Design notes:
//  * Generated factories emit `import * as THREE from 'three'` and construct
//    procedural canvas textures at call time, so this must run in a real
//    browser (the factory falls back to flat materials when document is
//    absent — the node smoke test relies on that).
//  * RoomEnvironment + ACES keeps MeshPhysicalMaterial metals/clearcoat
//    readable; the generated LookDevLights rig (neutral) is the rig the
//    upstream review rubric assumes.

import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
import { RoomEnvironment } from 'three/addons/environments/RoomEnvironment.js';

export class ViewerCaptureError extends Error {
  constructor(code, message, cause = null) {
    super(message, cause ? { cause } : undefined);
    this.name = 'ViewerCaptureError';
    this.code = code;
  }
}

function requireBoolean(value, option) {
  if (typeof value !== 'boolean') {
    throw new TypeError(`${option} must be a boolean`);
  }
  return value;
}

function eachMaterial(root, callback) {
  const seen = new Set();
  root.traverse((node) => {
    if (!node.material) return;
    const materials = Array.isArray(node.material) ? node.material : [node.material];
    for (const material of materials) {
      if (!material || seen.has(material)) continue;
      seen.add(material);
      callback(material);
    }
  });
  return seen.size;
}

// Exported as small, DOM-free helpers so the option behavior can be checked in
// Node without constructing a WebGLRenderer.
export function applyWireframeOption(root, enabled) {
  const next = requireBoolean(enabled, 'wireframe');
  let materials = 0;
  eachMaterial(root, (material) => {
    if (!('wireframe' in material)) return;
    materials += 1;
    if (material.wireframe !== next) {
      material.wireframe = next;
      material.needsUpdate = true;
    }
  });
  return { enabled: next, materials };
}

export function applyShadowsOption(renderer, ground, enabled) {
  const next = requireBoolean(enabled, 'shadows');
  renderer.shadowMap.enabled = next;
  renderer.shadowMap.needsUpdate = true;
  if (ground) ground.visible = next;
  return next;
}

export async function captureRendererPng(renderer, scene, camera) {
  try {
    renderer.render(scene, camera);
  } catch (error) {
    throw new ViewerCaptureError(
      'capture-render-failed',
      `Could not render the screenshot: ${(error && error.message) || String(error)}`,
      error,
    );
  }

  const blob = await new Promise((resolve, reject) => {
    try {
      renderer.domElement.toBlob((value) => {
        if (!value) {
          reject(new ViewerCaptureError(
            'capture-encode-failed',
            'The renderer did not produce PNG screenshot data.',
          ));
          return;
        }
        resolve(value);
      }, 'image/png');
    } catch (error) {
      reject(new ViewerCaptureError(
        'capture-encode-failed',
        `Could not encode the screenshot: ${(error && error.message) || String(error)}`,
        error,
      ));
    }
  });

  return new Promise((resolve, reject) => {
    let reader;
    try {
      reader = new FileReader();
      reader.onload = () => {
        if (typeof reader.result !== 'string') {
          reject(new ViewerCaptureError(
            'capture-read-failed',
            'The encoded screenshot could not be read as a data URL.',
          ));
          return;
        }
        resolve(reader.result);
      };
      reader.onerror = () => reject(new ViewerCaptureError(
        'capture-read-failed',
        'The browser could not read the encoded screenshot.',
        reader.error,
      ));
      reader.onabort = () => reject(new ViewerCaptureError(
        'capture-read-failed',
        'Reading the encoded screenshot was interrupted.',
      ));
      reader.readAsDataURL(blob);
    } catch (error) {
      reject(new ViewerCaptureError(
        'capture-read-failed',
        `Could not read the encoded screenshot: ${(error && error.message) || String(error)}`,
        error,
      ));
    }
  });
}

export function mountViewer(el, makeModel, makeLights, opts = {}) {
  if (typeof makeModel !== 'function') {
    throw new Error('No create*Model export found in the generated factory.');
  }
  const renderer = new THREE.WebGLRenderer({ antialias: true, preserveDrawingBuffer: true });
  renderer.toneMapping = THREE.ACESFilmicToneMapping;
  renderer.setPixelRatio(Math.min(window.devicePixelRatio || 1, 2));
  el.appendChild(renderer.domElement);
  renderer.domElement.style.display = 'block';

  const scene = new THREE.Scene();
  scene.background = new THREE.Color(0x16181d);
  const pmrem = new THREE.PMREMGenerator(renderer);
  const environmentTarget = pmrem.fromScene(new RoomEnvironment(), 0.04);
  scene.environment = environmentTarget.texture;

  const model = makeModel({
    textureSize: opts.textureSize || 512,
    qualityPriority: 'balanced',
    ...(opts.factoryOptions || {}),
  });
  scene.add(model);

  let lights = null;
  try {
    lights = typeof makeLights === 'function' ? makeLights('neutral') : null;
  } catch (_) {
    lights = null;
  }
  if (lights) {
    lights.traverse((node) => {
      if (node.isDirectionalLight && node.shadow && node.shadow.mapSize) {
        node.shadow.mapSize.set(2048, 2048);
      }
    });
    scene.add(lights);
  } else {
    scene.add(new THREE.HemisphereLight(0xffffff, 0x363b42, 1.1));
  }
  // Ambient floor so a spec that zeroes lights still shows silhouettes.
  scene.add(new THREE.AmbientLight(0xffffff, 0.15));

  const box = new THREE.Box3().setFromObject(model);
  const sphere = box.getBoundingSphere(new THREE.Sphere());
  const radius = Math.max(sphere.radius, 1e-3);

  const grid = new THREE.GridHelper(radius * 4, 20, 0x3a3f4a, 0x262a33);
  grid.position.y = box.min.y;
  scene.add(grid);

  const ground = new THREE.Mesh(
    new THREE.CircleGeometry(radius * 2.5, 48),
    new THREE.ShadowMaterial({ opacity: 0.35 }),
  );
  ground.rotation.x = -Math.PI / 2;
  ground.position.y = box.min.y - 0.001;
  ground.receiveShadow = true;
  scene.add(ground);

  const camera = new THREE.PerspectiveCamera(40, 1, radius / 100, radius * 100);
  const dir = new THREE.Vector3(1, 0.6, 1.4).normalize();
  const distance = (radius / Math.sin(THREE.MathUtils.degToRad(20))) * 1.1;
  camera.position.copy(sphere.center).addScaledVector(dir, distance);

  const controls = new OrbitControls(camera, renderer.domElement);
  controls.target.copy(sphere.center);
  controls.enableDamping = true;
  controls.update();
  controls.saveState();

  renderer.shadowMap.enabled = true;
  renderer.shadowMap.type = THREE.PCFSoftShadowMap;

  const fit = () => {
    const w = el.clientWidth || 640;
    const h = el.clientHeight || 420;
    renderer.setSize(w, h, false);
    camera.aspect = w / h;
    camera.updateProjectionMatrix();
  };
  const observer = new ResizeObserver(fit);
  observer.observe(el);
  fit();

  let meshes = 0;
  model.traverse((node) => { if (node.isMesh) meshes += 1; });

  renderer.setAnimationLoop(() => {
    controls.update();
    renderer.render(scene, camera);
  });

  let disposed = false;
  const optionState = { wireframe: false, shadows: true };

  const requireActive = () => {
    if (disposed) throw new Error('viewer session has been disposed');
  };

  const renderNow = () => {
    controls.update();
    renderer.render(scene, camera);
  };

  const setWireframe = (enabled) => {
    requireActive();
    const applied = applyWireframeOption(model, enabled);
    optionState.wireframe = applied.enabled;
    renderNow();
    return optionState.wireframe;
  };

  const setShadows = (enabled) => {
    requireActive();
    optionState.shadows = applyShadowsOption(renderer, ground, enabled);
    renderNow();
    return optionState.shadows;
  };

  const resetCamera = () => {
    requireActive();
    controls.reset();
    controls.update();
    camera.updateProjectionMatrix();
    renderer.render(scene, camera);
    return true;
  };

  const capture = () => {
    requireActive();
    return captureRendererPng(renderer, scene, camera);
  };

  const dispose = () => {
    if (disposed) return;
    disposed = true;
    observer.disconnect();
    renderer.setAnimationLoop(null);
    controls.dispose();
    const geometries = new Set();
    const materials = new Set();
    const textures = new Set();
    scene.traverse((node) => {
      if (node.geometry) geometries.add(node.geometry);
      if (node.material) {
        for (const material of Array.isArray(node.material) ? node.material : [node.material]) {
          if (!material) continue;
          materials.add(material);
          for (const key of Object.keys(material)) {
            const value = material[key];
            if (value && value.isTexture) textures.add(value);
          }
        }
      }
    });
    for (const geometry of geometries) geometry.dispose();
    for (const texture of textures) texture.dispose();
    for (const material of materials) material.dispose();
    environmentTarget.dispose();
    pmrem.dispose();
    renderer.dispose();
    if (typeof renderer.forceContextLoss === 'function') renderer.forceContextLoss();
    if (renderer.domElement.parentNode) renderer.domElement.remove();
  };

  return {
    capture,
    dispose,
    resetCamera,
    setShadows,
    setWireframe,
    get options() {
      return { ...optionState };
    },
    stats: {
      meshes,
      topLevelChildren: model.children.length,
      boundingSphereRadius: Number(radius.toFixed(4)),
      runtimeNodes: Object.keys(
        (model.userData && model.userData.sculptRuntime && model.userData.sculptRuntime.nodes) || {},
      ).length,
    },
  };
}