import * as THREE from 'three/webgpu';
import { IndirectBatchedMesh, indirectBatch, indirectBatchColor } from '@three-blocks/core';
import { diffuseColor, materialColor, normalLocal, vec4 } from 'three/tsl';

const geometry = new THREE.BoxGeometry(1.2, 1.2, 1.2);
const vertexCount = geometry.getAttribute('position').count;
const indexCount = geometry.getIndex()?.count || 0;
const instanceCount = 300;

const material = new THREE.MeshBasicNodeMaterial();
const mesh = new IndirectBatchedMesh(instanceCount, vertexCount, indexCount, material);

mesh.perObjectFrustumCulled = false;
const geoId = mesh.addGeometry(geometry);

material.colorNode = vec4(materialColor.rgb.add(normalLocal.normalize().mul(0.5)), 1.0);

const originalSetupDiffuseColor = material.setupDiffuseColor ? material.setupDiffuseColor.bind(material) : null;
material.setupDiffuseColor = function(builder) {
  // Run original setupDiffuseColor first (handles user colorNode, vertexColors, etc.)
  if (originalSetupDiffuseColor) {
    originalSetupDiffuseColor(builder);
  }

  // override diffuseColor with batchedTextColor
  diffuseColor.rgb.assign(indirectBatchColor.mul(diffuseColor.rgb));
};

// Add instances with different colors
const tmpMatrix = new THREE.Matrix4();
const tmpPosition = new THREE.Vector3();
const tmpQuat = new THREE.Quaternion();
const tmpEuler = new THREE.Euler();
const tmpScale = new THREE.Vector3();
const color = new THREE.Color();

for (let i = 0; i < instanceCount; i++) {
  const id = mesh.addInstance(geoId);

  tmpEuler.set(Math.random() * Math.PI, Math.random() * Math.PI, Math.random() * Math.PI);
  tmpQuat.setFromEuler(tmpEuler);
  tmpScale.setScalar(0.7 + Math.random() * 0.5);
  tmpPosition.set((Math.random() - 0.5) * 35, (Math.random() - 0.5) * 35, (Math.random() - 0.5) * 35);
  tmpMatrix.compose(tmpPosition, tmpQuat, tmpScale);
  mesh.setMatrixAt(id, tmpMatrix);

  // Set per-instance color using HSL for nice distribution
  mesh.setColorAt(id, color.setHSL(i / instanceCount, 0.7, 0.55));
}

mesh.needsUpdate = true;