import { Mat4 } from '../../core/math/mat4.js';
import { Vec3 } from '../../core/math/vec3.js';
import { GraphNode } from '../graph-node.js';
import { GSplatInfo } from './gsplat-info.js';
import { GSplatUnifiedSorter } from './gsplat-unified-sorter.js';
import { GSplatWorkBuffer } from './gsplat-work-buffer.js';
import { GSplatRenderer } from './gsplat-renderer.js';
import { GSplatOctreeInstance } from './gsplat-octree-instance.js';
import { GSplatOctreeResource } from './gsplat-octree.resource.js';
import { GSplatWorldState } from './gsplat-world-state.js';

const cameraPosition = new Vec3();
const cameraDirection = new Vec3();
const translation = new Vec3();
const invModelMat = new Mat4();
const tempNonOctreePlacements = new Set();
const tempOctreePlacements = new Set();
const _updatedSplats = [];
class GSplatManager {
		constructor(device, director, layer, cameraNode){
				this.node = new GraphNode('GSplatManager');
				this.worldStates = new Map();
				this.lastWorldStateVersion = 0;
				this.sortedVersion = 0;
				this.lastCameraPos = new Vec3(Infinity, Infinity, Infinity);
				this.layerPlacements = [];
				this.layerPlacementsDirty = false;
				this.octreeInstances = new Map();
				this.device = device;
				this.director = director;
				this.cameraNode = cameraNode;
				this.workBuffer = new GSplatWorkBuffer(device);
				this.renderer = new GSplatRenderer(device, this.node, this.cameraNode, layer, this.workBuffer);
				this.sorter = this.createSorter();
		}
		destroy() {
				this.workBuffer.destroy();
				this.renderer.destroy();
				this.sorter.destroy();
		}
		createSorter() {
				const sorter = new GSplatUnifiedSorter();
				sorter.on('sorted', (count, version, orderData)=>{
						this.onSorted(count, version, orderData);
				});
				return sorter;
		}
		reconcile(placements) {
				tempNonOctreePlacements.clear();
				for (const p of placements){
						if (p.resource instanceof GSplatOctreeResource) {
								if (!this.octreeInstances.has(p)) {
										this.octreeInstances.set(p, new GSplatOctreeInstance(p.resource.octree, p, this.director.assetLoader));
								}
								tempOctreePlacements.add(p);
						} else {
								tempNonOctreePlacements.add(p);
						}
				}
				for (const [placement, inst] of this.octreeInstances){
						if (!tempOctreePlacements.has(placement)) {
								this.octreeInstances.delete(placement);
								inst.destroy();
						}
				}
				this.layerPlacementsDirty = this.layerPlacements.length !== tempNonOctreePlacements.size;
				if (!this.layerPlacementsDirty) {
						for(let i = 0; i < this.layerPlacements.length; i++){
								const existing = this.layerPlacements[i];
								if (!tempNonOctreePlacements.has(existing)) {
										this.layerPlacementsDirty = true;
										break;
								}
						}
				}
				this.layerPlacements.length = 0;
				for (const p of tempNonOctreePlacements){
						this.layerPlacements.push(p);
				}
				tempNonOctreePlacements.clear();
				tempOctreePlacements.clear();
		}
		updateWorldState() {
				for (const [, inst] of this.octreeInstances){
						this.layerPlacementsDirty ||= inst.update();
				}
				const worldChanged = this.layerPlacementsDirty || this.worldStates.size === 0;
				if (worldChanged) {
						this.lastWorldStateVersion++;
						const splats = [];
						for (const p of this.layerPlacements){
								const splatInfo = new GSplatInfo(this.device, p.resource, p);
								splats.push(splatInfo);
						}
						for (const [, inst] of this.octreeInstances){
								inst.activePlacements.forEach((p)=>{
										splats.push(new GSplatInfo(this.device, p.resource, p));
								});
						}
						splats.forEach((splat)=>{
								this.sorter.setCenters(splat.resource.id, splat.resource.centers);
						});
						const newState = new GSplatWorldState(this.device, this.lastWorldStateVersion, splats);
						this.worldStates.set(this.lastWorldStateVersion, newState);
						this.layerPlacementsDirty = false;
				}
		}
		onSorted(count, version, orderData) {
				this.sortedVersion = version;
				const oldState = this.worldStates.get(version - 1);
				if (oldState) {
						this.worldStates.delete(version - 1);
						oldState.destroy();
				}
				const worldState = this.worldStates.get(version);
				if (worldState) {
						if (!worldState.sortedBefore) {
								worldState.sortedBefore = true;
								const textureSize = worldState.textureSize;
								if (textureSize !== this.workBuffer.textureSize) {
										this.workBuffer.resize(textureSize);
										this.renderer.setMaxNumSplats(textureSize * textureSize);
								}
								this.workBuffer.render(worldState.splats, this.cameraNode);
						}
						this.workBuffer.setOrderData(orderData);
						this.renderer.setNumSplats(count);
				}
		}
		update(scene) {
				let anyOctreeMoved = false;
				for (const [, inst] of this.octreeInstances){
						anyOctreeMoved ||= inst.testMoved();
				}
				const currentCameraPos = this.cameraNode.getPosition();
				const distance = this.lastCameraPos.distance(currentCameraPos);
				const cameraMoved = distance > 1.0;
				if (cameraMoved || anyOctreeMoved) {
						this.lastCameraPos.copy(currentCameraPos);
						for (const [, inst] of this.octreeInstances){
								inst.updateLod(this.cameraNode);
						}
				}
				this.updateWorldState();
				const lastState = this.worldStates.get(this.lastWorldStateVersion);
				if (lastState) {
						if (!lastState.sortParametersSet) {
								lastState.sortParametersSet = true;
								const payload = this.prepareSortParameters(lastState);
								this.sorter.setSortParameters(payload);
						}
						this.sort(lastState);
				}
				const sortedState = this.worldStates.get(this.sortedVersion);
				if (sortedState) {
						sortedState.splats.forEach((splat)=>{
								if (splat.update()) {
										_updatedSplats.push(splat);
								}
						});
						if (_updatedSplats.length > 0) {
								this.workBuffer.render(_updatedSplats, this.cameraNode);
								_updatedSplats.length = 0;
						}
				}
		}
		sort(lastState) {
				const cameraNode = this.cameraNode;
				const cameraMat = cameraNode.getWorldTransform();
				cameraMat.getTranslation(cameraPosition);
				cameraMat.getZ(cameraDirection).normalize();
				const sorterRequest = [];
				lastState.splats.forEach((splat)=>{
						const modelMat = splat.node.getWorldTransform();
						invModelMat.copy(modelMat).invert();
						const uniformScale = modelMat.getScale().x;
						const transformedDirection = invModelMat.transformVector(cameraDirection).normalize();
						modelMat.getTranslation(translation);
						const offset = translation.sub(cameraPosition).dot(cameraDirection);
						const aabbMin = splat.aabb.getMin();
						const aabbMax = splat.aabb.getMax();
						sorterRequest.push({
								transformedDirection,
								offset,
								scale: uniformScale,
								modelMat: modelMat.data.slice(),
								aabbMin: [
										aabbMin.x,
										aabbMin.y,
										aabbMin.z
								],
								aabbMax: [
										aabbMax.x,
										aabbMax.y,
										aabbMax.z
								]
						});
				});
				this.sorter.setSortParams(sorterRequest);
				this.renderer.updateViewport(cameraNode);
		}
		prepareSortParameters(worldState) {
				return {
						command: 'intervals',
						textureSize: worldState.textureSize,
						totalUsedPixels: worldState.totalUsedPixels,
						version: worldState.version,
						ids: worldState.splats.map((splat)=>splat.resource.id),
						lineStarts: worldState.splats.map((splat)=>splat.lineStart),
						padding: worldState.splats.map((splat)=>splat.padding),
						intervals: worldState.splats.map((splat)=>splat.intervals)
				};
		}
}

export { GSplatManager };
