workinf_Blender_Wasm/web/app/src/three-adapter/texture-assets.ts

import {
  EquirectangularReflectionMapping,
  NoColorSpace,
  PMREMGenerator,
  SRGBColorSpace,
  Texture,
  type MeshPhysicalMaterial,
  type Object3D,
  type Scene,
  type WebGLRenderer,
} from "../vendor/three/three.module.js";
import type { MaterialIR, SceneSnapshotIR, WorldIR } from "../../../protocol/scene-ir";
import { RenderAssetValidationError, validateGPUTextureAsset, type GPUTextureAsset, type GPUTextureColorSpace, type GPUTextureUsage } from "../../../protocol/render-assets";
import { planPBRTextureBudget, type PBRDeviceLimits, type PBRRenderBackend, type PBRTextureBudgetReport } from "../../../protocol/render-budget";

export interface TextureUploadStatus {
  loaded: number;
  rejected: number;
  bytes: number;
  errors: string[];
  errorCodes: string[];
  budget: PBRTextureBudgetReport;
}

export interface GPUTextureStoreStatsIR {
  entries: number;
  payloadBytes: number;
  decodedGPUBytes: number;
  keys: string[];
}

export interface GPUTextureStoreReleaseIR {
  releasedTextures: number;
  releasedPayloadBytes: number;
  releasedGPUBytes: number;
  resourcesAfter: 0;
}

function key(imageId: string, usage: GPUTextureUsage): string {
  return `${imageId}:${usage}`;
}

function colorSpaceFor(usage: GPUTextureUsage, declared: GPUTextureColorSpace): string {
  if (declared === "SRGB" || usage === "BASE_COLOR" || usage === "EMISSIVE") return SRGBColorSpace;
  return NoColorSpace;
}

async function decodeTexture(asset: GPUTextureAsset): Promise<Texture> {
  if (typeof createImageBitmap !== "function") throw new Error("createImageBitmap is unavailable");
  const blob = new Blob([asset.data], { type: asset.mimeType });
  const bitmap = await createImageBitmap(blob);
  const texture = new Texture(bitmap);
  // ImageBitmap pixels are already oriented for the WebGL upload path.
  texture.flipY = false;
  texture.colorSpace = colorSpaceFor(asset.usage, asset.colorSpace);
  texture.name = `${asset.imageId}:${asset.usage}${asset.tileNumber ? `:${asset.tileNumber}` : ""}`;
  texture.needsUpdate = true;
  return texture;
}

export class GPUTextureStore {
  private readonly textures = new Map<string, Texture>();
  private readonly assets = new Map<string, GPUTextureAsset>();
  private readonly udimTileCounts = new Map<string, number>();
  private revision = 0;

  constructor(
    private readonly backend: PBRRenderBackend = "THREE_WEBGL2",
    private readonly deviceLimits: PBRDeviceLimits = {},
  ) {}

  getRevision(): number {
    return this.revision;
  }

  get(imageId: string, usage: GPUTextureUsage): Texture | undefined {
    return this.textures.get(key(imageId, usage)) ?? this.textures.get(key(imageId, "BASE_COLOR"));
  }

  getAsset(imageId: string, usage: GPUTextureUsage): GPUTextureAsset | undefined {
    return this.assets.get(key(imageId, usage));
  }

  stats(): GPUTextureStoreStatsIR {
    const assets = [...this.assets.values()];
    return {
      entries: this.textures.size,
      payloadBytes: assets.reduce((total, asset) => total + asset.byteLength, 0),
      decodedGPUBytes: assets.reduce((total, asset) => total + asset.width * asset.height * 4, 0),
      keys: [...this.textures.keys()].sort(),
    };
  }

  async upload(assets: readonly GPUTextureAsset[]): Promise<TextureUploadStatus> {
    const candidate = new Map(this.assets);
    for (const asset of assets) {
      const lookupUsage = asset.usage === "UDIM_TILE" ? "BASE_COLOR" : asset.usage;
      candidate.set(key(asset.imageId, lookupUsage), asset);
    }
    const budget = planPBRTextureBudget([...candidate.values()], this.backend, this.deviceLimits);
    const status: TextureUploadStatus = { loaded: 0, rejected: 0, bytes: 0, errors: [], errorCodes: [], budget };
    if (budget.status === "BLOCKED") {
      status.rejected = assets.length;
      status.errors.push(...budget.issues.map((issue) => issue.message));
      status.errorCodes.push(...budget.issues.map((issue) => issue.code));
      this.revision += 1;
      return status;
    }
    const incomingUDIMCounts = new Map<string, number>();
    for (const asset of assets) if (asset.usage === "UDIM_TILE") incomingUDIMCounts.set(asset.imageId, (incomingUDIMCounts.get(asset.imageId) ?? 0) + 1);
    for (const asset of assets) {
      try {
        await validateGPUTextureAsset(asset);
        if (asset.usage === "UDIM_TILE" && (incomingUDIMCounts.get(asset.imageId) ?? 0) > 1) {
          throw new Error(`UDIM multi-tile sampling is blocked for ${asset.imageId}`);
        }
        const texture = await decodeTexture(asset);
        const lookupUsage = asset.usage === "UDIM_TILE" ? "BASE_COLOR" : asset.usage;
        const old = this.textures.get(key(asset.imageId, lookupUsage));
        old?.dispose();
        this.textures.set(key(asset.imageId, lookupUsage), texture);
        this.assets.set(key(asset.imageId, lookupUsage), asset);
        if (asset.usage === "UDIM_TILE") this.udimTileCounts.set(asset.imageId, (this.udimTileCounts.get(asset.imageId) ?? 0) + 1);
        status.loaded += 1;
        status.bytes += asset.byteLength;
      }
      catch (error) {
        status.rejected += 1;
        status.errors.push(error instanceof Error ? error.message : "Texture upload failed");
        status.errorCodes.push(error instanceof RenderAssetValidationError ? error.code : "GPU_TEXTURE_DECODE_FAILED");
      }
    }
    this.revision += 1;
    return status;
  }

  applyMaterial(material: MeshPhysicalMaterial, definition?: MaterialIR): void {
    if (!definition) return;
    const compile = material.userData.shaderCompile as { status?: string; textureBindings?: Array<{ imageId: string; usage: "BASE_COLOR" | "NORMAL" }> } | undefined;
    const baseImageId = compile?.status === "COMPILED"
      ? compile.textureBindings?.find((binding) => binding.usage === "BASE_COLOR")?.imageId
      : definition.imageIds?.find((imageId) => imageId !== definition.normalImageId) ?? definition.nodes?.find((node) => node.type === "IMAGE_TEXTURE" && node.imageId && node.imageId !== definition.normalImageId)?.imageId;
    if (baseImageId) {
      const texture = this.get(baseImageId, "BASE_COLOR");
      if (texture) material.map = texture;
    }
    const normalImageId = compile?.status === "COMPILED"
      ? compile.textureBindings?.find((binding) => binding.usage === "NORMAL")?.imageId
      : definition.normalImageId ?? definition.nodes?.find((node) => node.type === "IMAGE_TEXTURE" && node.imageId)?.imageId;
    if (normalImageId) {
      const texture = this.get(normalImageId, "NORMAL");
      if (texture) material.normalMap = texture;
    }
    if (material.map || material.normalMap) material.needsUpdate = true;
  }

  applySnapshotMaterials(root: Object3D, snapshot: SceneSnapshotIR): void {
    const materialById = new Map(snapshot.materials.map((material) => [material.id, material]));
    root.traverse((object) => {
      const mesh = object as { material?: unknown };
      const materialIds = object.userData.materialSlotIds as string[] | undefined;
      const materials = Array.isArray(mesh.material) ? mesh.material : mesh.material ? [mesh.material] : [];
      materials.forEach((material, index) => {
        if (!material || typeof material !== "object" || !("isMeshPhysicalMaterial" in material)) return;
        this.applyMaterial(material as MeshPhysicalMaterial, materialById.get(materialIds?.[index] ?? ""));
      });
    });
  }

  async applyEnvironment(scene: Scene, renderer: WebGLRenderer, world: WorldIR | undefined, background = true): Promise<boolean> {
    if (!world?.environmentImageId) return false;
    const source = this.get(world.environmentImageId, "ENVIRONMENT");
    if (!source) return false;
    source.mapping = EquirectangularReflectionMapping;
    source.needsUpdate = true;
    const pmrem = new PMREMGenerator(renderer);
    const target = pmrem.fromEquirectangular(source);
    pmrem.dispose();
    const oldTarget = scene.userData.pbrIBLTarget as { dispose?: () => void } | undefined;
    oldTarget?.dispose?.();
    scene.userData.pbrIBLTarget = target;
    scene.environment = target.texture;
    scene.environmentIntensity = Math.max(0, world.environmentStrength ?? 1);
    if (background && (world.backgroundVisible ?? true)) scene.background = source;
    return true;
  }

  dispose(): GPUTextureStoreReleaseIR {
    const before = this.stats();
    for (const texture of this.textures.values()) texture.dispose();
    this.textures.clear();
    this.assets.clear();
    this.udimTileCounts.clear();
    return {
      releasedTextures: before.entries,
      releasedPayloadBytes: before.payloadBytes,
      releasedGPUBytes: before.decodedGPUBytes,
      resourcesAfter: 0,
    };
  }
}