import * as THREE from 'three';
import { MobxReconcilerMap, type ReconcilerMap } from '../utils/reconciler';
import type { ObjectTypeGeometryViewItem } from '../state/viewLayers';

export class ObjectGeometryList {
    private items: ReconcilerMap<ObjectTypeGeometryViewItem>;

    private readonly matrixObj = new THREE.Object3D();
    private readonly _p = new THREE.Vector3();
    private readonly _e = new THREE.Euler();
    private readonly _q = new THREE.Quaternion();

    private readonly _meshToObjectIds = new Map<THREE.InstancedMesh, string[]>();

    constructor(scene: THREE.Scene, typeGeometries: Iterable<ObjectTypeGeometryViewItem>) {
        this.items = new MobxReconcilerMap<ObjectTypeGeometryViewItem>({
            reaction: () => typeGeometries,
            keyOf: (item) => item.objectType.id,
            initializeItem: (item) => {
                for (const mesh of item.meshes ?? [])
                    scene.add(mesh);
            },
            disposeItem: (item) => {
                for (const mesh of item.meshes ?? [])
                    scene.remove(mesh);
            },
        });
    }

    dispose(): void {
        this.items?.dispose();
    }

    getObjectIdByMesh(mesh: THREE.InstancedMesh, instanceId: number): string | undefined {
        return this._meshToObjectIds.get(mesh)?.[instanceId];
    }

    tick(): void {
        this._meshToObjectIds.clear();

        this.items.values.forEach((typeGeometry) => {
            const instances = typeGeometry.instances;
            const meshes = typeGeometry.meshes ?? [];
            const objectIds: string[] = [];

            instances.forEach((item, instanceIdx) => {
                const obj = item.object;

                this._p.set(...obj.position);
                this._e.set(...obj.rotation);
                this._q.setFromEuler(this._e);
                this.matrixObj.position.copy(this._p);
                this.matrixObj.quaternion.copy(this._q);
                this.matrixObj.scale.set(...obj.scale);
                this.matrixObj.updateMatrix();

                objectIds[instanceIdx] = obj.id;

                for (const mesh of meshes)
                    mesh.setMatrixAt(instanceIdx, this.matrixObj.matrix);
            });

            for (const mesh of meshes) {
                this._meshToObjectIds.set(mesh, objectIds);
                mesh.count = instances.length;
                mesh.instanceMatrix.needsUpdate = true;
            }
        });
    }
}