import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
import { GLTFExporter } from 'three/addons/exporters/GLTFExporter.js';
import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1/+esm';
import * as ort from 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/webgpu/+esm';

const MODEL_ROOT = 'https://huggingface.co/needle-tools/SF3D-webgpu/resolve/main/';
const SAMPLE_URL = 'https://huggingface.co/spaces/stabilityai/stable-fast-3d/resolve/main/demo_files/examples/axe.png';
const INPUT_SIZE = 512;
const GRID_RESOLUTION = 160;
const ISO_THRESHOLD = 10;
const TRIPLANE_CHANNELS = 40;
const TRIPLANE_SIZE = 384;
const DECODE_BATCH = 16384;
const BACKGROUND_COLOR = 0.5;

const el = (id) => document.getElementById(id);
const ui = {
  file: el('image-file'),
  dropzone: el('dropzone'),
  previewWrap: el('preview-wrap'),
  preview: el('input-preview'),
  previewName: el('preview-name'),
  previewSize: el('preview-size'),
  sample: el('sample-button'),
  removeBackground: el('remove-background'),
  generate: el('generate-button'),
  generateLabel: el('generate-label'),
  webgpuDot: el('webgpu-dot'),
  webgpuLabel: el('webgpu-label'),
  stage: el('viewer-stage'),
  canvas: el('viewer-canvas'),
  empty: el('viewer-empty'),
  progress: el('viewer-progress'),
  progressPercent: el('progress-percent'),
  progressTitle: el('progress-title'),
  progressDetail: el('progress-detail'),
  error: el('viewer-error'),
  errorTitle: el('error-title'),
  errorDetail: el('error-detail'),
  errorDismiss: el('error-dismiss'),
  reset: el('reset-view'),
  stats: el('mesh-stats'),
  runtime: el('runtime-metric'),
  vertices: el('vertices-metric'),
  faces: el('faces-metric'),
  download: el('download-button'),
};

let selectedFile = null;
let selectedObjectUrl = null;
let preparedCanvas = null;
let backgroundPipeline = null;
let sessions = null;
let currentMeshObject = null;
let currentMeshData = null;
let webgpuReady = false;

function setWebGPUStatus(kind, label) {
  ui.webgpuDot.className = `status-dot ${kind}`;
  ui.webgpuLabel.textContent = label;
}

function setProgress(title, detail, percent = null) {
  ui.progressTitle.textContent = title;
  ui.progressDetail.textContent = detail;
  if (percent === null) {
    ui.progressPercent.textContent = '…';
  } else {
    ui.progressPercent.textContent = `${Math.round(percent)}%`;
  }
}

function showProgress(visible) {
  ui.progress.classList.toggle('is-hidden', !visible);
  ui.empty.classList.toggle('is-hidden', visible || Boolean(currentMeshObject));
}

function showError(title, detail) {
  ui.errorTitle.textContent = title;
  ui.errorDetail.textContent = detail;
  ui.error.classList.remove('is-hidden');
  showProgress(false);
}

function clearError() {
  ui.error.classList.add('is-hidden');
}

function formatBytes(bytes) {
  if (!Number.isFinite(bytes) || bytes <= 0) return '—';
  const units = ['B', 'KB', 'MB', 'GB'];
  const index = Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1);
  return `${(bytes / 1024 ** index).toFixed(index ? 1 : 0)} ${units[index]}`;
}

function formatDuration(ms) {
  return ms < 1000 ? `${Math.round(ms)} ms` : `${(ms / 1000).toFixed(1)} s`;
}

function isTransparentCanvas(canvas) {
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  const { data } = ctx.getImageData(0, 0, canvas.width, canvas.height);
  for (let i = 3; i < data.length; i += 4) {
    if (data[i] < 250) return true;
  }
  return false;
}

function loadImageFromBlob(blob) {
  return new Promise((resolve, reject) => {
    const url = URL.createObjectURL(blob);
    const image = new Image();
    image.onload = () => {
      URL.revokeObjectURL(url);
      resolve(image);
    };
    image.onerror = () => {
      URL.revokeObjectURL(url);
      reject(new Error('The selected file could not be decoded as an image.'));
    };
    image.src = url;
  });
}

function imageToCanvas(image, maxSize = 1200) {
  const scale = Math.min(1, maxSize / Math.max(image.naturalWidth || image.width, image.naturalHeight || image.height));
  const canvas = document.createElement('canvas');
  canvas.width = Math.max(1, Math.round((image.naturalWidth || image.width) * scale));
  canvas.height = Math.max(1, Math.round((image.naturalHeight || image.height) * scale));
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  ctx.clearRect(0, 0, canvas.width, canvas.height);
  ctx.drawImage(image, 0, 0, canvas.width, canvas.height);
  return canvas;
}

function rawImageToCanvas(rawImage) {
  const canvas = document.createElement('canvas');
  canvas.width = rawImage.width;
  canvas.height = rawImage.height;
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  const rgba = new Uint8ClampedArray(canvas.width * canvas.height * 4);
  const source = rawImage.data;
  const channels = rawImage.channels || 4;
  for (let i = 0, j = 0; i < source.length; i += channels, j += 4) {
    rgba[j] = source[i];
    rgba[j + 1] = source[i + 1] ?? source[i];
    rgba[j + 2] = source[i + 2] ?? source[i];
    rgba[j + 3] = channels >= 4 ? source[i + 3] : 255;
  }
  ctx.putImageData(new ImageData(rgba, canvas.width, canvas.height), 0, 0);
  return canvas;
}

async function getBackgroundPipeline() {
  if (!backgroundPipeline) {
    setProgress('Loading cleanup model', 'Transformers.js is downloading MODNet for this browser.', null);
    backgroundPipeline = await pipeline('background-removal', 'Xenova/modnet', {
      device: 'webgpu',
      dtype: 'fp16',
    });
  }
  return backgroundPipeline;
}

async function prepareSourceCanvas(file) {
  const image = await loadImageFromBlob(file);
  const original = imageToCanvas(image);
  if (!ui.removeBackground.checked || isTransparentCanvas(original)) {
    return original;
  }

  let blobUrl = null;
  try {
    const remover = await getBackgroundPipeline();
    blobUrl = URL.createObjectURL(file);
    const output = await remover(blobUrl);
    const first = Array.isArray(output) ? output[0] : output;
    if (!first?.data || !first.width || !first.height) {
      throw new Error('The cleanup model returned an empty image.');
    }
    return rawImageToCanvas(first);
  } catch (error) {
    console.warn('Background cleanup skipped:', error);
    setProgress('Using original image', 'Background cleanup was skipped; continuing with the uploaded pixels.', 5);
    return original;
  } finally {
    if (blobUrl) URL.revokeObjectURL(blobUrl);
  }
}

function alphaBounds(canvas) {
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  const { data, width, height } = ctx.getImageData(0, 0, canvas.width, canvas.height);
  let minX = width;
  let minY = height;
  let maxX = -1;
  let maxY = -1;
  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      if (data[(y * width + x) * 4 + 3] > 8) {
        minX = Math.min(minX, x);
        minY = Math.min(minY, y);
        maxX = Math.max(maxX, x);
        maxY = Math.max(maxY, y);
      }
    }
  }
  if (maxX < 0) return { x: 0, y: 0, width, height };
  return { x: minX, y: minY, width: maxX - minX + 1, height: maxY - minY + 1 };
}

function prepareModelCanvas(source, ratio = 0.85) {
  const bounds = alphaBounds(source);
  const side = Math.max(bounds.width, bounds.height) / ratio;
  const centerX = bounds.x + bounds.width / 2;
  const centerY = bounds.y + bounds.height / 2;
  const cropX = centerX - side / 2;
  const cropY = centerY - side / 2;
  const canvas = document.createElement('canvas');
  canvas.width = INPUT_SIZE;
  canvas.height = INPUT_SIZE;
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  ctx.clearRect(0, 0, INPUT_SIZE, INPUT_SIZE);
  ctx.imageSmoothingEnabled = true;
  ctx.imageSmoothingQuality = 'high';
  ctx.drawImage(source, cropX, cropY, side, side, 0, 0, INPUT_SIZE, INPUT_SIZE);
  return canvas;
}

function renderInputPreview(canvas, name, bytes) {
  if (selectedObjectUrl) URL.revokeObjectURL(selectedObjectUrl);
  canvas.toBlob((blob) => {
    if (!blob) return;
    selectedObjectUrl = URL.createObjectURL(blob);
    ui.preview.src = selectedObjectUrl;
  }, 'image/png');
  ui.previewName.textContent = name;
  ui.previewSize.textContent = formatBytes(bytes);
  ui.previewWrap.classList.remove('is-hidden');
}

function selectFile(file) {
  if (!file || !file.type.startsWith('image/')) {
    showError('Unsupported file', 'Choose a PNG, JPEG, or WebP image.');
    return;
  }
  selectedFile = file;
  clearError();
  preparedCanvas = null;
  ui.generate.disabled = !webgpuReady;
  ui.generateLabel.textContent = webgpuReady ? 'Generate 3D mesh' : 'WebGPU required';
  ui.previewWrap.classList.add('is-hidden');
  const reader = new FileReader();
  reader.onload = () => {
    ui.preview.src = reader.result;
    ui.previewName.textContent = file.name;
    ui.previewSize.textContent = formatBytes(file.size);
    ui.previewWrap.classList.remove('is-hidden');
  };
  reader.readAsDataURL(file);
}

async function loadSample() {
  ui.sample.disabled = true;
  ui.sample.querySelector('span').textContent = 'Loading sample…';
  try {
    const response = await fetch(SAMPLE_URL);
    if (!response.ok) throw new Error(`Sample request failed (${response.status}).`);
    const blob = await response.blob();
    selectFile(new File([blob], 'axe.png', { type: blob.type || 'image/png' }));
  } catch (error) {
    showError('Sample unavailable', error.message);
  } finally {
    ui.sample.disabled = false;
    ui.sample.querySelector('span').textContent = 'Try a sample';
  }
}

function urlFor(path) {
  return `${MODEL_ROOT}${path}?download=true`;
}

function resolveName(names, preferred, index = 0) {
  return preferred.find((candidate) => names.includes(candidate)) || names[index];
}

function resolveOutput(outputs, preferred, index = 0) {
  for (const name of preferred) {
    if (outputs[name]) return outputs[name];
  }
  const first = Object.keys(outputs)[index];
  if (!first) throw new Error('The ONNX graph returned no output tensor.');
  return outputs[first];
}

async function createSessions() {
  if (sessions) return sessions;
  if (!webgpuReady) throw new Error('WebGPU is not available in this browser.');
  if (ort.env?.wasm) {
    ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/dist/';
    ort.env.wasm.numThreads = 1;
  }
  const options = {
    executionProviders: ['webgpu'],
    graphOptimizationLevel: 'all',
    enableMemPattern: true,
  };
  setProgress('Loading SF3D graphs', 'Downloading the tokenizer, backbone, and decoder. They will be cached after this run.', null);
  const [imageTokenizer, backbone, decoder] = await Promise.all([
    ort.InferenceSession.create(urlFor('onnx/image_tokenizer_single.onnx'), options),
    ort.InferenceSession.create(urlFor('onnx/backbone_fp16.onnx'), options),
    ort.InferenceSession.create(urlFor('onnx/decoder_single.onnx'), options),
  ]);
  sessions = { imageTokenizer, backbone, decoder };
  return sessions;
}

async function fetchBinary(path, label, progressOffset, progressScale) {
  const response = await fetch(urlFor(path));
  if (!response.ok) throw new Error(`Could not load ${label} (${response.status}).`);
  const total = Number(response.headers.get('content-length')) || 0;
  const reader = response.body?.getReader();
  if (!reader) return response.arrayBuffer();
  const chunks = [];
  let loaded = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    loaded += value.byteLength;
    if (total) setProgress('Loading mesh grid', `${label} · ${formatBytes(loaded)} / ${formatBytes(total)}`, progressOffset + (loaded / total) * progressScale);
  }
  const buffer = new ArrayBuffer(loaded);
  const output = new Uint8Array(buffer);
  let offset = 0;
  for (const chunk of chunks) {
    output.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return buffer;
}

async function loadGrid() {
  setProgress('Loading mesh grid', 'Fetching the tetrahedral surface grid.', 68);
  const [verticesBuffer, indicesBuffer] = await Promise.all([
    fetchBinary('tets_vertices.bin', 'vertices', 68, 7),
    fetchBinary('tets_indices.bin', 'tetrahedra', 75, 7),
  ]);
  return {
    vertices: new Float32Array(verticesBuffer),
    indices: new Uint32Array(indicesBuffer),
  };
}

function createInputTensors(canvas, imageTokenizer) {
  const ctx = canvas.getContext('2d', { willReadFrequently: true });
  const { data } = ctx.getImageData(0, 0, INPUT_SIZE, INPUT_SIZE);
  const rgb = new Float32Array(INPUT_SIZE * INPUT_SIZE * 3);
  for (let i = 0, j = 0; i < data.length; i += 4, j += 3) {
    const alpha = data[i + 3] / 255;
    rgb[j] = BACKGROUND_COLOR * (1 - alpha) + (data[i] / 255) * alpha;
    rgb[j + 1] = BACKGROUND_COLOR * (1 - alpha) + (data[i + 1] / 255) * alpha;
    rgb[j + 2] = BACKGROUND_COLOR * (1 - alpha) + (data[i + 2] / 255) * alpha;
  }
  const inputNames = imageTokenizer.inputNames;
  const rgbName = resolveName(inputNames, ['rgb'], 0);
  const c2wName = resolveName(inputNames, ['c2w'], 1);
  const intrinsicName = resolveName(inputNames, ['intrinsic_normed'], 2);
  const focal = 1 / (2 * Math.tan((40 * Math.PI) / 360));
  return {
    [rgbName]: new ort.Tensor('float32', rgb, [1, INPUT_SIZE, INPUT_SIZE, 3]),
    [c2wName]: new ort.Tensor('float32', new Float32Array([
      0, 0, 1, 1.6,
      1, 0, 0, 0,
      0, 1, 0, 0,
      0, 0, 0, 1,
    ]), [1, 4, 4]),
    [intrinsicName]: new ort.Tensor('float32', new Float32Array([
      focal, 0, 0.5,
      0, focal, 0.5,
      0, 0, 1,
    ]), [1, 3, 3]),
  };
}

async function decodeGrid(grid, triplane, decoder) {
  const vertexCount = grid.vertices.length / 3;
  const density = new Float32Array(vertexCount);
  const deformation = new Float32Array(vertexCount * 3);
  const triplaneTensor = new ort.Tensor('float32', triplane, [1, 3, TRIPLANE_CHANNELS, TRIPLANE_SIZE, TRIPLANE_SIZE]);
  const positions = new Float32Array(DECODE_BATCH * 3);
  const inputNames = decoder.inputNames;
  const triplaneName = resolveName(inputNames, ['triplane'], 0);
  const positionsName = resolveName(inputNames, ['positions'], 1);
  const densityName = 'density';
  const offsetName = 'vertex_offset';
  for (let start = 0; start < vertexCount; start += DECODE_BATCH) {
    const count = Math.min(DECODE_BATCH, vertexCount - start);
    const currentPositions = positions.subarray(0, count * 3);
    for (let i = 0; i < count; i += 1) {
      const base = (start + i) * 3;
      currentPositions[i * 3] = grid.vertices[base] * 2 - 1;
      currentPositions[i * 3 + 1] = grid.vertices[base + 1] * 2 - 1;
      currentPositions[i * 3 + 2] = grid.vertices[base + 2] * 2 - 1;
    }
    const outputs = await decoder.run({
      [triplaneName]: triplaneTensor,
      [positionsName]: new ort.Tensor('float32', currentPositions, [1, count, 3]),
    });
    const densityTensor = resolveOutput(outputs, [densityName]);
    const offsetTensor = resolveOutput(outputs, [offsetName]);
    for (let i = 0; i < count; i += 1) {
      density[start + i] = densityTensor.data[i];
      deformation[(start + i) * 3] = offsetTensor.data[i * 3];
      deformation[(start + i) * 3 + 1] = offsetTensor.data[i * 3 + 1];
      deformation[(start + i) * 3 + 2] = offsetTensor.data[i * 3 + 2];
    }
    const percent = 82 + ((start + count) / vertexCount) * 15;
    setProgress('Decoding surface', `${(start + count).toLocaleString()} / ${vertexCount.toLocaleString()} grid points`, percent);
  }
  return { density, deformation };
}

const TRIANGLE_TABLE = [
  [-1, -1, -1, -1, -1, -1], [1, 0, 2, -1, -1, -1], [4, 0, 3, -1, -1, -1], [1, 4, 2, 1, 3, 4],
  [3, 1, 5, -1, -1, -1], [2, 3, 0, 2, 5, 3], [1, 4, 0, 1, 5, 4], [4, 2, 5, -1, -1, -1],
  [4, 5, 2, -1, -1, -1], [4, 1, 0, 4, 5, 1], [3, 2, 0, 3, 5, 2], [1, 3, 5, -1, -1, -1],
  [4, 1, 2, 4, 3, 1], [3, 0, 4, -1, -1, -1], [2, 0, 1, -1, -1, -1], [-1, -1, -1, -1, -1, -1],
];
function buildMesh(grid, decoded) {
  const vertexCount = grid.vertices.length / 3;
  const tetCount = grid.indices.length / 4;
  const occupied = new Uint8Array(vertexCount);
  for (let i = 0; i < vertexCount; i += 1) occupied[i] = decoded.density[i] > ISO_THRESHOLD ? 1 : 0;

  const positions = [];
  const faces = [];
  const edgeMap = new Map();
  const deformed = (index, axis) => {
    const value = grid.vertices[index * 3 + axis] + (2 / GRID_RESOLUTION) * Math.tanh(decoded.deformation[index * 3 + axis]);
    return value * 2 - 1;
  };
  const edgeVertex = (a, b) => {
    if (occupied[a] === occupied[b]) return -1;
    const low = Math.min(a, b);
    const high = Math.max(a, b);
    const key = low * vertexCount + high;
    const existing = edgeMap.get(key);
    if (existing !== undefined) return existing;
    const s0 = decoded.density[a] - ISO_THRESHOLD;
    const s1 = decoded.density[b] - ISO_THRESHOLD;
    const denominator = s0 - s1 || 1e-6;
    const t = Math.max(0, Math.min(1, s0 / denominator));
    const x = deformed(a, 0) * (1 - t) + deformed(b, 0) * t;
    const y = deformed(a, 1) * (1 - t) + deformed(b, 1) * t;
    const z = deformed(a, 2) * (1 - t) + deformed(b, 2) * t;
    const index = positions.length / 3;
    positions.push(x, y, z);
    edgeMap.set(key, index);
    return index;
  };

  for (let tet = 0; tet < tetCount; tet += 1) {
    const offset = tet * 4;
    const c0 = grid.indices[offset];
    const c1 = grid.indices[offset + 1];
    const c2 = grid.indices[offset + 2];
    const c3 = grid.indices[offset + 3];
    const mask = occupied[c0] | (occupied[c1] << 1) | (occupied[c2] << 2) | (occupied[c3] << 3);
    if (mask === 0 || mask === 15) continue;
    const edgeIndices = new Int32Array(6);
    edgeIndices[0] = edgeVertex(c0, c1);
    edgeIndices[1] = edgeVertex(c0, c2);
    edgeIndices[2] = edgeVertex(c0, c3);
    edgeIndices[3] = edgeVertex(c1, c2);
    edgeIndices[4] = edgeVertex(c1, c3);
    edgeIndices[5] = edgeVertex(c2, c3);
    const table = TRIANGLE_TABLE[mask];
    for (let i = 0; i < 6 && table[i] !== -1; i += 3) {
      faces.push(edgeIndices[table[i]], edgeIndices[table[i + 1]], edgeIndices[table[i + 2]]);
    }
  }
  if (!faces.length) throw new Error('No surface was found. Try a clearer image with one centered object.');
  return { positions: new Float32Array(positions), faces: new Uint32Array(faces) };
}

function initViewer() {
  const scene = new THREE.Scene();
  const camera = new THREE.PerspectiveCamera(35, 1, 0.01, 100);
  camera.position.set(0, 0.15, 3.4);
  const renderer = new THREE.WebGLRenderer({ canvas: ui.canvas, antialias: true, alpha: true });
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
  renderer.outputColorSpace = THREE.SRGBColorSpace;
  renderer.toneMapping = THREE.ACESFilmicToneMapping;
  renderer.toneMappingExposure = 1.15;
  const controls = new OrbitControls(camera, renderer.domElement);
  controls.enableDamping = true;
  controls.dampingFactor = 0.07;
  controls.minDistance = 1.2;
  controls.maxDistance = 7;
  controls.target.set(0, 0, 0);
  scene.add(new THREE.HemisphereLight(0xdcd7ff, 0x1f2330, 2.1));
  const keyLight = new THREE.DirectionalLight(0xffffff, 3.5);
  keyLight.position.set(3, 4, 4);
  scene.add(keyLight);
  const rimLight = new THREE.DirectionalLight(0xc7baff, 2.2);
  rimLight.position.set(-4, 2, -3);
  scene.add(rimLight);

  const resize = () => {
    const width = ui.stage.clientWidth;
    const height = ui.stage.clientHeight;
    renderer.setSize(width, height, false);
    camera.aspect = width / Math.max(height, 1);
    camera.updateProjectionMatrix();
  };
  const reset = () => {
    camera.position.set(0, 0.15, 3.4);
    controls.target.set(0, 0, 0);
    controls.update();
  };
  window.addEventListener('resize', resize);
  ui.reset.addEventListener('click', reset);
  resize();
  const animate = () => {
    requestAnimationFrame(animate);
    controls.update();
    renderer.render(scene, camera);
  };
  animate();

  return {
    scene,
    camera,
    reset,
    setMesh(meshData) {
      if (currentMeshObject) {
        scene.remove(currentMeshObject);
        currentMeshObject.geometry.dispose();
        currentMeshObject.material.dispose();
      }
      const geometry = new THREE.BufferGeometry();
      geometry.setAttribute('position', new THREE.BufferAttribute(meshData.positions, 3));
      geometry.setIndex(new THREE.BufferAttribute(meshData.faces, 1));
      geometry.computeVertexNormals();
      geometry.computeBoundingBox();
      const center = geometry.boundingBox.getCenter(new THREE.Vector3());
      const size = geometry.boundingBox.getSize(new THREE.Vector3());
      const maxDimension = Math.max(size.x, size.y, size.z) || 1;
      geometry.translate(-center.x, -center.y, -center.z);
      geometry.scale(2.15 / maxDimension, 2.15 / maxDimension, 2.15 / maxDimension);
      const material = new THREE.MeshStandardMaterial({ color: 0xc9b9ff, roughness: 0.44, metalness: 0.14, side: THREE.DoubleSide });
      currentMeshObject = new THREE.Mesh(geometry, material);
      scene.add(currentMeshObject);
      reset();
      ui.empty.classList.add('is-hidden');
      ui.stats.textContent = `${meshData.faces.length / 3} triangles · drag to orbit`;
      ui.download.disabled = false;
    },
  };
}

const viewer = initViewer();

async function runGeneration() {
  if (!selectedFile || !webgpuReady) return;
  const started = performance.now();
  ui.generate.disabled = true;
  ui.download.disabled = true;
  clearError();
  showProgress(true);
  try {
    setProgress('Preparing image', 'Composing the subject over SF3D’s neutral background.', 5);
    const sourceCanvas = await prepareSourceCanvas(selectedFile);
    preparedCanvas = prepareModelCanvas(sourceCanvas);
    renderInputPreview(preparedCanvas, `${selectedFile.name} · prepared`, selectedFile.size);
    setProgress('Loading models', 'Connecting to the browser-compatible SF3D export.', 15);
    const [modelSessions, grid] = await Promise.all([createSessions(), loadGrid()]);
    setProgress('Encoding image', 'The image tokenizer is running on WebGPU.', 78);
    const tokenizerInputs = createInputTensors(preparedCanvas, modelSessions.imageTokenizer);
    const tokenizerOutputs = await modelSessions.imageTokenizer.run(tokenizerInputs);
    const imageTokens = resolveOutput(tokenizerOutputs, ['image_tokens']);
    setProgress('Reconstructing shape', 'The SF3D backbone is predicting a 3D triplane.', 80);
    const backboneInputName = resolveName(modelSessions.backbone.inputNames, ['image_tokens'], 0);
    const backboneOutputs = await modelSessions.backbone.run({ [backboneInputName]: imageTokens });
    const triplane = resolveOutput(backboneOutputs, ['triplane']).data;
    setProgress('Decoding surface', 'Sampling the triplane and extracting the tetrahedral surface.', 82);
    const decoded = await decodeGrid(grid, triplane, modelSessions.decoder);
    const meshData = buildMesh(grid, decoded);
    currentMeshData = meshData;
    viewer.setMesh(meshData);
    const elapsed = performance.now() - started;
    ui.runtime.textContent = formatDuration(elapsed);
    ui.vertices.textContent = (meshData.positions.length / 3).toLocaleString();
    ui.faces.textContent = (meshData.faces.length / 3).toLocaleString();
    setProgress('Done', 'Your mesh is ready.', 100);
    showProgress(false);
  } catch (error) {
    console.error(error);
    showError('Generation failed', error?.message || 'The browser could not complete WebGPU inference.');
    ui.stats.textContent = 'No mesh loaded';
  } finally {
    ui.generate.disabled = false;
    ui.generateLabel.textContent = 'Generate 3D mesh';
  }
}

async function downloadGLB() {
  if (!currentMeshObject) return;
  ui.download.disabled = true;
  try {
    const exporter = new GLTFExporter();
    const result = await new Promise((resolve, reject) => exporter.parse(currentMeshObject, resolve, reject, { binary: true }));
    const blob = new Blob([result], { type: 'model/gltf-binary' });
    const url = URL.createObjectURL(blob);
    const anchor = document.createElement('a');
    anchor.href = url;
    anchor.download = 'stable-fast-3d-mesh.glb';
    anchor.click();
    setTimeout(() => URL.revokeObjectURL(url), 1000);
  } catch (error) {
    showError('Download failed', error.message || 'The GLB exporter could not finish.');
  } finally {
    ui.download.disabled = false;
  }
}

async function checkWebGPU() {
  if (!navigator.gpu) {
    setWebGPUStatus('error', 'WebGPU unavailable');
    ui.generateLabel.textContent = 'WebGPU required';
    return;
  }
  try {
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) throw new Error('No compatible GPU adapter found.');
    webgpuReady = true;
    setWebGPUStatus('ready', 'WebGPU ready');
    ui.generate.disabled = !selectedFile;
    ui.generateLabel.textContent = selectedFile ? 'Generate 3D mesh' : 'Select an image';
  } catch (error) {
    setWebGPUStatus('error', 'WebGPU unavailable');
    showError('WebGPU is required', error.message || 'Use a recent Chrome or Edge browser with WebGPU enabled.');
  }
}

ui.file.addEventListener('change', (event) => selectFile(event.target.files?.[0]));
ui.sample.addEventListener('click', loadSample);
ui.generate.addEventListener('click', runGeneration);
ui.download.addEventListener('click', downloadGLB);
ui.errorDismiss.addEventListener('click', clearError);
ui.dropzone.addEventListener('dragover', (event) => { event.preventDefault(); ui.dropzone.classList.add('dragging'); });
ui.dropzone.addEventListener('dragleave', () => ui.dropzone.classList.remove('dragging'));
ui.dropzone.addEventListener('drop', (event) => {
  event.preventDefault();
  ui.dropzone.classList.remove('dragging');
  selectFile(event.dataTransfer.files?.[0]);
});
window.addEventListener('pagehide', () => {
  if (backgroundPipeline?.dispose) backgroundPipeline.dispose();
  if (sessions) Object.values(sessions).forEach((session) => session.release?.());
});

checkWebGPU();