import { describe, expect, it, vi } from "vitest";
import {
  BoxGeometry,
  Group,
  InstancedMesh,
  Matrix4,
  Mesh,
  MeshBasicMaterial,
  Vector3,
} from "three";
import { cacheStaticMatrices, instanceRepeatedStaticMeshes } from "./staticSceneInstancing";

function batchIn(root: Group): InstancedMesh {
  let result: InstancedMesh | null = null;
  root.traverse((node) => {
    if ((node as InstancedMesh).isInstancedMesh) result = node as InstancedMesh;
  });
  if (!result) throw new Error("Expected an instanced mesh batch.");
  return result;
}

function batchesIn(root: Group): readonly InstancedMesh[] {
  const result: InstancedMesh[] = [];
  root.traverse((node) => {
    if ((node as InstancedMesh).isInstancedMesh) result.push(node as InstancedMesh);
  });
  return result;
}

describe("static scene instancing", () => {
  it("caches local transforms without freezing dungeon placement or imported matrices", () => {
    const parent = new Group();
    const root = new Group();
    const mesh = new Mesh(new BoxGeometry(), new MeshBasicMaterial());
    mesh.matrixAutoUpdate = false;
    mesh.matrix.makeTranslation(3, 0, 0);
    parent.add(root);
    root.add(mesh);
    const updateChild = vi.spyOn(mesh, "updateMatrixWorld");
    cacheStaticMatrices(root);
    root.position.x = 10;
    parent.position.x = 20;
    parent.updateMatrixWorld(true);
    expect(new Vector3().setFromMatrixPosition(mesh.matrixWorld).x).toBe(33);
    parent.updateMatrixWorld(true);
    expect(updateChild).toHaveBeenCalledTimes(1);
    root.position.x = 15;
    parent.updateMatrixWorld(true);
    expect(new Vector3().setFromMatrixPosition(mesh.matrixWorld).x).toBe(38);
    parent.position.x = 30;
    parent.updateMatrixWorld(true);
    expect(new Vector3().setFromMatrixPosition(mesh.matrixWorld).x).toBe(48);
    expect(updateChild).toHaveBeenCalledTimes(3);
    expect(root.matrixAutoUpdate).toBe(true);
    expect(mesh.matrixWorldAutoUpdate).toBe(true);
  });

  it("batches repeated meshes while retaining their world transforms", () => {
    const root = new Group();
    root.position.set(10, 0, 0);
    const geometry = new BoxGeometry();
    const material = new MeshBasicMaterial();
    for (const x of [1, 3, 7]) {
      const mesh = new Mesh(geometry, material);
      mesh.position.x = x;
      root.add(mesh);
    }

    expect(instanceRepeatedStaticMeshes(root)).toEqual({
      batches: 1,
      instances: 3,
      removedMeshes: 3,
    });
    const batch = batchIn(root);
    const matrix = new Matrix4();
    const position = new Vector3();
    const positions = Array.from({ length: batch.count }, (_, index) => {
      batch.getMatrixAt(index, matrix);
      position.setFromMatrixPosition(matrix);
      return position.x;
    });
    expect(positions).toEqual([1, 3, 7]);
  });

  it("leaves small groups and existing instance batches unchanged", () => {
    const root = new Group();
    const geometry = new BoxGeometry();
    const material = new MeshBasicMaterial();
    root.add(new Mesh(geometry, material), new Mesh(geometry, material));
    const existing = new InstancedMesh(geometry, material, 4);
    root.add(existing);

    expect(instanceRepeatedStaticMeshes(root)).toEqual({
      batches: 0,
      instances: 0,
      removedMeshes: 0,
    });
    expect(root.children).toHaveLength(3);
    expect(root.children).toContain(existing);
  });

  it("removes containers emptied by batching while preserving authored empty nodes", () => {
    const root = new Group();
    const geometry = new BoxGeometry();
    const material = new MeshBasicMaterial();
    const authoredMarker = new Group();
    authoredMarker.name = "authored-marker";
    root.add(authoredMarker);
    for (let index = 0; index < 3; index += 1) {
      const outer = new Group();
      const inner = new Group();
      outer.add(inner);
      inner.add(new Mesh(geometry, material));
      root.add(outer);
    }
    instanceRepeatedStaticMeshes(root);
    expect(root.children).toHaveLength(2);
    expect(root.children).toContain(authoredMarker);
    expect(batchIn(root).count).toBe(3);
  });

  it("does not combine meshes whose render state differs", () => {
    const root = new Group();
    const geometry = new BoxGeometry();
    const material = new MeshBasicMaterial();
    for (let index = 0; index < 3; index += 1) {
      const mesh = new Mesh(geometry, material);
      mesh.receiveShadow = index !== 2;
      root.add(mesh);
    }

    expect(instanceRepeatedStaticMeshes(root)).toEqual({
      batches: 0,
      instances: 0,
      removedMeshes: 0,
    });
  });

  it("keeps distant repetitions in separate cullable spatial batches", () => {
    const root = new Group();
    const geometry = new BoxGeometry();
    const material = new MeshBasicMaterial();
    for (const x of [1, 2, 3, 385, 386, 387]) {
      const mesh = new Mesh(geometry, material);
      mesh.position.x = x;
      root.add(mesh);
    }

    expect(instanceRepeatedStaticMeshes(root)).toEqual({
      batches: 2,
      instances: 6,
      removedMeshes: 6,
    });
    expect(batchesIn(root).map((batch) => batch.count)).toEqual([3, 3]);
  });
});