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 { CanvasTexture, Mesh, TextureLoader } from 'three'
import type { Material, Object3D, Texture } from 'three'
import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader.js'
import { normalizeTextureFilename } from '@/lib/asset-naming'
import type { ModelHierarchyNode, ModelStats, SceneViewerProps } from '@/lib/client-types'

interface OpacityMapEntry {
  target: string
  url: string
}

interface AlphaMapMaterial extends Material {
  alphaMap: Texture | null
  transparent: boolean
  alphaTest: number
}

type AlphaImageSource = HTMLImageElement | HTMLCanvasElement | ImageBitmap

const alphaMapTextureCache = new WeakMap<Texture, Texture>()

function getRequestedFilename(requestedUrl: string) {
  const cleanUrl = decodeURIComponent(requestedUrl.split(/[?#]/)[0] || '')
  return cleanUrl.split(/[\\/]/).pop()?.toLowerCase()
}

function resolveSingleBinFallback(filename: string | undefined, assetUrls: Record<string, string>) {
  if (!filename?.endsWith('.bin')) return undefined

  const binEntries = Object.entries(assetUrls).filter(([assetName]) => assetName.endsWith('.bin'))
  return binEntries.length === 1 ? binEntries[0][1] : undefined
}

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

  const filename = getRequestedFilename(requestedUrl)
  const exactAssetUrl = filename ? assetUrls[filename] : undefined
  const fallbackBinUrl = resolveSingleBinFallback(filename, assetUrls)

  return exactAssetUrl || fallbackBinUrl || requestedUrl
}

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

    if (!match) return entries

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

function normalizeMatchName(name: string) {
  return name
    .toLowerCase()
    .normalize('NFD')
    .replace(/[\u0300-\u036f]/g, '')
    .replace(/[^a-z0-9]/g, '')
}

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

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

function isMesh(object: Object3D): object is Mesh {
  return object instanceof Mesh
}

function supportsAlphaMap(material: Material): material is AlphaMapMaterial {
  return 'alphaMap' in material
}

function isAlphaImageSource(image: unknown): image is AlphaImageSource {
  return image instanceof HTMLImageElement
    || image instanceof HTMLCanvasElement
    || (typeof ImageBitmap !== 'undefined' && image instanceof ImageBitmap)
}

function createAlphaMapTexture(texture: Texture) {
  const cachedTexture = alphaMapTextureCache.get(texture)
  if (cachedTexture) return cachedTexture

  const image: unknown = texture.image

  if (!isAlphaImageSource(image)) {
    texture.flipY = false
    alphaMapTextureCache.set(texture, texture)
    return texture
  }

  const canvas = document.createElement('canvas')
  canvas.width = image.width
  canvas.height = image.height

  const context = canvas.getContext('2d')
  if (!context) {
    texture.flipY = false
    alphaMapTextureCache.set(texture, texture)
    return texture
  }

  context.drawImage(image, 0, 0)

  const imageData = context.getImageData(0, 0, canvas.width, canvas.height)
  const pixels = imageData.data

  for (let i = 0; i < pixels.length; i += 4) {
    const alpha = pixels[i + 3]
    const luminance = Math.round((pixels[i] * 0.2126) + (pixels[i + 1] * 0.7152) + (pixels[i + 2] * 0.0722))
    const opacity = alpha < 255 ? alpha : luminance

    pixels[i] = opacity
    pixels[i + 1] = opacity
    pixels[i + 2] = opacity
    pixels[i + 3] = 255
  }

  context.putImageData(imageData, 0, 0)

  const alphaTexture = new CanvasTexture(canvas)
  alphaTexture.flipY = false
  alphaTexture.needsUpdate = true
  alphaMapTextureCache.set(texture, alphaTexture)

  return alphaTexture
}

function applyAlphaMap(material: Material, texture: Texture) {
  if (!supportsAlphaMap(material)) return

  material.alphaMap = createAlphaMapTexture(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) => {
    if (!isMesh(object) || !object.geometry || !object.material) return

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

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

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

function roundTransformValue(value: number) {
  return Number(value.toFixed(4))
}

function getObjectHierarchy(object: Object3D): ModelHierarchyNode {
  return {
    children: object.children.map(getObjectHierarchy),
    id: object.uuid,
    name: object.name || object.type,
    position: [
      roundTransformValue(object.position.x),
      roundTransformValue(object.position.y),
      roundTransformValue(object.position.z),
    ],
    rotation: [
      roundTransformValue(object.rotation.x),
      roundTransformValue(object.rotation.y),
      roundTransformValue(object.rotation.z),
    ],
    type: object.type,
    visible: object.visible,
  }
}

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(normalizeMatchName(entry.target))
      || materialName.includes(normalizeMatchName(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,
  onHierarchyReady,
}: SceneViewerProps) {
  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))

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

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

    scene.traverse((object) => {
      if (!isMesh(object) || !object.material) return

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

  return <primitive object={scene} />
}

export default function SceneViewer({
  url,
  assetUrls,
  onStatsReady,
  onHierarchyReady,
}: SceneViewerProps) {
  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}
            onHierarchyReady={onHierarchyReady}
          />
        </Stage>
      </Suspense>
      <OrbitControls makeDefault autoRotate autoRotateSpeed={0.5} />
    </Canvas>
  )
}