import * as THREE from 'three/webgpu';
import { IndirectBatchedMesh, indirectBatch } from '@three-blocks/core';
import { positionLocal, morphReference, skinning, uniform, subBuild, normalLocal, Fn } from 'three/tsl';

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

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

// Disable internal culling for this simple demo
mesh.perObjectFrustumCulled = false;

const geoId = mesh.addGeometry(geometry);

// Setup the indirect batch transform
material.setupPosition = function(builder) {
  const { object, geometry } = builder;

  if (geometry.morphAttributes.position || geometry.morphAttributes.normal || geometry.morphAttributes.color) {
    morphReference(object).toStack();
  }

  if (object.isSkinnedMesh === true) {
    skinning(object).toStack();
  }

  if (object.isIndirectBatchedMesh) {
    // inject here
    indirectBatch(object).toStack();

    const worldMatrix = uniform(new THREE.Matrix4()).onObjectUpdate((_, self) => self.value.copy(object.matrixWorld)).setName('batchWorldMatrix');
    object.setCullingMatrixNode(({ baseMatrix }) => worldMatrix.mul(baseMatrix));
  }

  if (this.positionNode !== null) {
    positionLocal.assign(subBuild(this.positionNode, 'POSITION', 'vec3'));
  }

  return positionLocal;
};

// Use normal-based coloring to verify transforms work correctly
material.colorNode = Fn(() => {
  return normalLocal.mul(0.5).add(0.5);
})();

// Add instances in a grid pattern
const tmpMatrix = new THREE.Matrix4();
const gridSize = Math.ceil(Math.cbrt(instanceCount));
const spacing = 2.5;

let idx = 0;
for (let x = 0; x < gridSize && idx < instanceCount; x++) {
  for (let y = 0; y < gridSize && idx < instanceCount; y++) {
    for (let z = 0; z < gridSize && idx < instanceCount; z++) {
      const id = mesh.addInstance(geoId);
      tmpMatrix.makeTranslation(
        (x - gridSize / 2) * spacing,
        (y - gridSize / 2) * spacing,
        (z - gridSize / 2) * spacing
      );
      mesh.setMatrixAt(id, tmpMatrix);
      idx++;
    }
  }
}

mesh.needsUpdate = true;