upload-gltf/components/SceneViewer.tsx

'use client'

import { Suspense, useEffect } from 'react'
import { Canvas } from '@react-three/fiber'
import { Stage, OrbitControls } from '@react-three/drei'
import { useLoader } from '@react-three/fiber'
import type { Material, Mesh, Object3D, Texture } from 'three'
import { TextureLoader } from 'three'
import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader.js'

export interface ModelStats {
  drawCalls: number
  materials: number
  meshes: number
  textures: number
  triangles: number
}

interface OpacityMapEntry {
  target: string
  url: string
}

function resolveAssetUrl(requestedUrl: string, assetUrls: Record<string, string>) {
  if (requestedUrl.startsWith('blob:') || requestedUrl.startsWith('data:')) {
    return requestedUrl
  }

  const cleanUrl = decodeURIComponent(requestedUrl.split(/[?#]/)[0] || '')
  const filename = cleanUrl.split(/[\\/]/).pop()?.toLowerCase()

  return filename ? assetUrls[filename] || requestedUrl : requestedUrl
}

function getOpacityMapEntries(assetUrls: Record<string, string>) {
  return Object.entries(assetUrls).reduce<OpacityMapEntry[]>((entries, [filename, url]) => {
    const match = filename.toLowerCase().match(/^opacity(?:[_-](.+))?\.(png|jpe?g|webp)$/)

    if (!match) return entries

    entries.push({ target: match[1] || '', url })
    return entries
  }, [])
}

function getMaterialName(material: Material) {
  return material.name.toLowerCase()
}

function getObjectName(mesh: Mesh) {
  return mesh.name.toLowerCase()
}

function applyAlphaMap(material: Material, texture: Texture) {
  if (!('alphaMap' in material)) return

  texture.flipY = false
  material.alphaMap = texture
  material.transparent = true
  material.alphaTest = 0.01
  material.needsUpdate = true
}

function countGeometryTriangles(mesh: Mesh) {
  const geometry = mesh.geometry

  if (geometry.index) return Math.floor(geometry.index.count / 3)

  const position = geometry.getAttribute('position')
  return position ? Math.floor(position.count / 3) : 0
}

function getMaterialCount(material: Material | Material[]) {
  return Array.isArray(material) ? material.length : 1
}

function estimateMeshDrawCalls(mesh: Mesh) {
  if (mesh.geometry.groups.length > 0) return mesh.geometry.groups.length

  return getMaterialCount(mesh.material)
}

function getTextureCount(assetUrls: Record<string, string>) {
  return Object.keys(assetUrls).filter((filename) => /\.(png|jpe?g|webp)$/i.test(filename)).length
}

function getModelStats(scene: Object3D, assetUrls: Record<string, string>): ModelStats {
  const materials = new Set<Material>()
  let drawCalls = 0
  let meshes = 0
  let triangles = 0

  scene.traverse((object) => {
    const mesh = object as Mesh
    if (!mesh.isMesh || !mesh.geometry || !mesh.material) return

    meshes += 1
    triangles += countGeometryTriangles(mesh)
    drawCalls += estimateMeshDrawCalls(mesh)

    const meshMaterials = Array.isArray(mesh.material) ? mesh.material : [mesh.material]
    meshMaterials.forEach((material) => materials.add(material))
  })

  return {
    drawCalls,
    materials: materials.size,
    meshes,
    textures: getTextureCount(assetUrls),
    triangles,
  }
}

function pickOpacityMap(
  mesh: Mesh,
  material: Material,
  entries: OpacityMapEntry[],
  textures: Texture[],
) {
  const objectName = getObjectName(mesh)
  const materialName = getMaterialName(material)
  const targetedIndex = entries.findIndex((entry) => (
    entry.target && (objectName.includes(entry.target) || materialName.includes(entry.target))
  ))

  if (targetedIndex >= 0) return textures[targetedIndex]

  const genericIndex = entries.findIndex((entry) => entry.target === '')
  if (genericIndex >= 0) return textures[genericIndex]

  return entries.length === 1 ? textures[0] : undefined
}

function Model({
  url,
  assetUrls,
  onStatsReady,
}: {
  url: string
  assetUrls: Record<string, string>
  onStatsReady: (stats: ModelStats) => void
}) {
  const { scene } = useLoader(GLTFLoader, url, (loader) => {
    loader.manager.setURLModifier((requestedUrl) => resolveAssetUrl(requestedUrl, assetUrls))
  })
  const opacityMapEntries = getOpacityMapEntries(assetUrls)
  const opacityMaps = useLoader(TextureLoader, opacityMapEntries.map((entry) => entry.url)) as Texture[]

  useEffect(() => {
    onStatsReady(getModelStats(scene, assetUrls))
  }, [assetUrls, onStatsReady, scene])

  useEffect(() => {
    if (opacityMapEntries.length === 0) return

    scene.traverse((object) => {
      const mesh = object as Mesh
      if (!mesh.isMesh || !mesh.material) return

      const materials = Array.isArray(mesh.material) ? mesh.material : [mesh.material]
      materials.forEach((material) => {
        const opacityMap = pickOpacityMap(mesh, material, opacityMapEntries, opacityMaps)
        if (opacityMap) applyAlphaMap(material, opacityMap)
      })
    })
  }, [scene, opacityMapEntries, opacityMaps])

  return <primitive object={scene} />
}

export default function SceneViewer({
  url,
  assetUrls,
  onStatsReady,
}: {
  url: string
  assetUrls: Record<string, string>
  onStatsReady: (stats: ModelStats) => void
}) {
  return (
    <Canvas dpr={[1, 2]} camera={{ fov: 50 }}>
      <Suspense fallback={null}>
        <Stage environment="city" intensity={0.6} adjustCamera={1.2}>
          <Model url={url} assetUrls={assetUrls} onStatsReady={onStatsReady} />
        </Stage>
      </Suspense>
      <OrbitControls makeDefault autoRotate autoRotateSpeed={0.5} />
    </Canvas>
  )
}