package iron.object;

#if lnx_particles

import kha.FastFloat;
import kha.graphics4.Usage;
import kha.arrays.Float32Array;
import iron.data.Data;
import iron.data.ParticleData;
import iron.data.SceneFormat;
import iron.data.Geometry;
import iron.data.MeshData;
import iron.system.Time;
import iron.math.Mat4;
import iron.math.Quat;
import iron.math.Vec3;
import iron.math.Vec4;

class ParticleSystem {
	public var data: ParticleData;
	public var speed = 1.0;
	public var dynamicEmitter: Bool = true;
	var currentSpeed = 0.0;
	var particles: Array<Particle>;
	var ready: Bool;
	var frameRate = 24;
	var lifetime = 0.0;
	var looptime = 0.0;
	var animtime = 0.0;
	var time = 0.0;
	var spawnRate = 0.0;
	var seed = 0;

	var r: TParticleData;
	var gx: Float;
	var gy: Float;
	var gz: Float;
	var alignx: Float;
	var aligny: Float;
	var alignz: Float;
	var dimx: Float;
	var dimy: Float;
	var tilesx: Int;
	var tilesy: Int;
	var tilesFramerate: Int;

	var count = 0;
	var lap = 0;
	var lapTime = 0.0;
	var m = Mat4.identity();

	var ownerLoc = new Vec4();
	var ownerRot = new Quat();
	var ownerScl = new Vec4();
	
	var random = 0.0;

	var tmpV4 = new Vec4();

	var instancedData: Float32Array = null;
	var lastSpawnedCount: Int = 0;
	var hasUniqueGeom: Bool = false; 

	public function new(sceneName: String, pref: TParticleReference) {
		seed = pref.seed;
		currentSpeed = speed;
		speed = 0;
		particles = [];
		ready = false;
		
		Data.getParticle(sceneName, pref.particle, function(b: ParticleData) {
			data = b;
			r = data.raw;
			if (r.dynamic_emitter != null){
				dynamicEmitter = r.dynamic_emitter; 
			} else {
				dynamicEmitter = true;
			}
			if (Scene.active.raw.gravity != null) {
				gx = Scene.active.raw.gravity[0] * r.weight_gravity;
				gy = Scene.active.raw.gravity[1] * r.weight_gravity;
				gz = Scene.active.raw.gravity[2] * r.weight_gravity;
			}
			else {
				gx = 0;
				gy = 0;
				gz = -9.81 * r.weight_gravity;
			}
			alignx = r.object_align_factor[0];
			aligny = r.object_align_factor[1];
			alignz = r.object_align_factor[2];
			looptime = (r.frame_end - r.frame_start) / frameRate;
			lifetime = r.lifetime / frameRate;
			animtime = r.loop ? looptime : looptime + lifetime;
			spawnRate = ((r.frame_end - r.frame_start) / r.count) / frameRate;

			for (i in 0...r.count) {
				particles.push(new Particle(i));
			}

			ready = true;
			if (r.auto_start){
				start();
			}
		});
	}

	public function start() {
		if (r.is_unique) random = Math.random();
		lifetime = r.lifetime / frameRate;
		time = 0;
		lap = 0;
		lapTime = 0;
		speed = currentSpeed;
		lastSpawnedCount = 0;
		instancedData = null;
	}

	public function pause() {
		speed = 0;
	}

	public function resume() {
		lifetime = r.lifetime / frameRate;
		speed = currentSpeed;
	}



	// TODO: interrupt smoothly
	public function stop() {
		end();
	}

	function end() {
		lifetime = 0;
		speed = 0;
		lap = 0;
	}
	
	public function update(object: MeshObject, owner: MeshObject) {
		if (!ready || object == null || speed == 0.0) return;
		if (iron.App.pauseUpdates) return;
		
		var prevLap = lap;

		// Copy owner world transform but discard scale
		owner.transform.world.decompose(ownerLoc, ownerRot, ownerScl);
		if (dynamicEmitter) {
			object.transform.loc.x = 0; object.transform.loc.y = 0; object.transform.loc.z = 0;
			object.transform.rot = new Quat();
		} else {
			object.transform.loc = ownerLoc;
			object.transform.rot = ownerRot;
		}

		// Set particle size per particle system
		object.transform.scale = new Vec4(r.particle_size, r.particle_size, r.particle_size, 1);

		object.transform.buildMatrix();
		owner.transform.buildMatrix();
		object.transform.dim.setFrom(owner.transform.dim);

		dimx = object.transform.dim.x;
		dimy = object.transform.dim.y;

		if (object.activeTilesheet != null) {
			tilesx = object.activeTilesheet.raw.tilesx;
			tilesy = object.activeTilesheet.raw.tilesy;
			tilesFramerate = object.activeTilesheet.raw.framerate;
		}

		// Animate
		time += Time.renderDelta * speed; // realDelta to renderDelta
		lap = Std.int(time / animtime);
		lapTime = time - lap * animtime;
		count = Std.int(lapTime / spawnRate);

		if (lap > prevLap && !r.loop) {
			end();
		}

		if (lap > prevLap && r.loop) {
			lastSpawnedCount = 0;
		}
		
		updateGpu(object, owner);
	}

	public function getData(): Mat4 {
		var hair = r.type == 1;
		// Store loop flag in the sign: positive -> loop, negative -> no loop
		m._00 = r.loop ? animtime : -animtime;
		m._01 = hair ? 1 / particles.length : spawnRate;
		m._02 = hair ? 1 : lifetime;
		m._03 = particles.length;
		m._10 = hair ? 0 : alignx;
		m._11 = hair ? 0 : aligny;
		m._12 = hair ? 0 : alignz;
		m._13 = hair ? 0 : r.factor_random;
		m._20 = hair ? 0 : gx;
		m._21 = hair ? 0 : gy;
		m._22 = hair ? 0 : gz;
		m._23 = hair ? 0 : r.lifetime_random;
		m._30 = tilesx;
		m._31 = tilesy;
		m._32 = 1 / tilesFramerate;
		m._33 = hair ? 1 : lapTime;
		return m;
	}

	public function getSizeRandom(): FastFloat {
		return r.size_random;
	}

	public inline function getRandom(): FastFloat {
		return random;
	}

	public inline function getSize(): FastFloat {
		return r.particle_size;
	}

	function updateGpu(object: MeshObject, owner: MeshObject) {
		if (dynamicEmitter) {
			if (!hasUniqueGeom) ensureUniqueGeom(object);
			var needSetup = instancedData == null || object.data.geom.instancedVB == null;
			if (needSetup) setupGeomGpuDynamic(object, owner);
			updateSpawnedInstances(object, owner);
		}
		else {
			if (!hasUniqueGeom) ensureUniqueGeom(object);
			if (!object.data.geom.instanced) setupGeomGpu(object, owner);
		}
		// GPU particles transform is attached to owner object in static mode
	}

	function setupGeomGpu(object: MeshObject, owner: MeshObject) {
		var instancedData = new Float32Array(particles.length * 3);
		var i = 0;

		var normFactor = 1 / 32767; // pa.values are not normalized
		var scalePosOwner = owner.data.scalePos;
		var scalePosParticle = object.data.scalePos;
		var particleSize = r.particle_size;
		var scaleFactor = new Vec4().setFrom(owner.transform.scale);
		scaleFactor.mult(scalePosOwner / (particleSize * scalePosParticle));

		switch (r.emit_from) {
			case 0: // Vert
				var pa = owner.data.geom.positions;

				for (p in particles) {
					var j = Std.int(fhash(i) * (pa.values.length / pa.size));
					instancedData.set(i, pa.values[j * pa.size    ] * normFactor * scaleFactor.x); i++;
					instancedData.set(i, pa.values[j * pa.size + 1] * normFactor * scaleFactor.y); i++;
					instancedData.set(i, pa.values[j * pa.size + 2] * normFactor * scaleFactor.z); i++;
				}

			case 1: // Face
				var positions = owner.data.geom.positions.values;

				for (p in particles) {
					// Choose random index array (there is one per material) and random face
					var ia = owner.data.geom.indices[Std.random(owner.data.geom.indices.length)];
					var faceIndex = Std.random(Std.int(ia.length / 3));

					var i0 = ia[faceIndex * 3 + 0];
					var i1 = ia[faceIndex * 3 + 1];
					var i2 = ia[faceIndex * 3 + 2];

					var v0 = new Vec3(positions[i0 * 4], positions[i0 * 4 + 1], positions[i0 * 4 + 2]);
					var v1 = new Vec3(positions[i1 * 4], positions[i1 * 4 + 1], positions[i1 * 4 + 2]);
					var v2 = new Vec3(positions[i2 * 4], positions[i2 * 4 + 1], positions[i2 * 4 + 2]);

					var pos = randomPointInTriangle(v0, v1, v2);

					instancedData.set(i, pos.x * normFactor * scaleFactor.x); i++;
					instancedData.set(i, pos.y * normFactor * scaleFactor.y); i++;
					instancedData.set(i, pos.z * normFactor * scaleFactor.z); i++;
				}

			case 2: // Volume
				var scaleFactorVolume = new Vec4().setFrom(object.transform.dim);
				scaleFactorVolume.mult(0.5 / (particleSize * scalePosParticle));

				for (p in particles) {
					instancedData.set(i, (Math.random() * 2.0 - 1.0) * scaleFactorVolume.x); i++;
					instancedData.set(i, (Math.random() * 2.0 - 1.0) * scaleFactorVolume.y); i++;
					instancedData.set(i, (Math.random() * 2.0 - 1.0) * scaleFactorVolume.z); i++;
				}
		}
		object.data.geom.setupInstanced(instancedData, 1, Usage.StaticUsage);
	}

	// allocate instanced VB once for this object 
	function setupGeomGpuDynamic(object: MeshObject, owner: MeshObject) {
		if (instancedData == null) instancedData = new Float32Array(particles.length * 3);
		lastSpawnedCount = 0; 
		// Create instanced VB once if missing (seed with our instancedData)
		if (object.data.geom.instancedVB == null) {
			object.data.geom.setupInstanced(instancedData, 1, Usage.DynamicUsage);
		}
	}

	function ensureUniqueGeom(object: MeshObject) {
		if (hasUniqueGeom) return;
		var newData: MeshData = null;
		new MeshData(object.data.raw, function(dat: MeshData) {
			dat.scalePos = object.data.scalePos;
			dat.scaleTex = object.data.scaleTex;
			dat.format = object.data.format;
			newData = dat;
		});
		if (newData != null) object.setData(newData);
		hasUniqueGeom = true;
	}

	function updateSpawnedInstances(object: MeshObject, owner: MeshObject) {
		if (instancedData == null) return;
		var targetCount = count;
		if (targetCount > particles.length) targetCount = particles.length;
		if (targetCount <= lastSpawnedCount) return;

		var normFactor = 1 / 32767;
		var scalePosOwner = owner.data.scalePos;
		var scalePosParticle = object.data.scalePos;
		var particleSize = r.particle_size;
		var base = 1.0 / (particleSize * scalePosParticle);

		switch (r.emit_from) {
			case 0: // Vert
				var pa = owner.data.geom.positions;
				var osx = owner.transform.scale.x;
				var osy = owner.transform.scale.y;
				var osz = owner.transform.scale.z;
				var pCount = Std.int(pa.values.length / pa.size);
				for (idx in lastSpawnedCount...targetCount) {
					var j = Std.int(fhash(idx) * pCount);
					var lx = pa.values[j * pa.size    ] * normFactor * scalePosOwner * osx;
					var ly = pa.values[j * pa.size + 1] * normFactor * scalePosOwner * osy;
					var lz = pa.values[j * pa.size + 2] * normFactor * scalePosOwner * osz;
					tmpV4.x = lx; tmpV4.y = ly; tmpV4.z = lz; tmpV4.w = 1;
					tmpV4.applyQuat(ownerRot);
					var o = idx * 3;
					instancedData.set(o    , (tmpV4.x + ownerLoc.x) * base);
					instancedData.set(o + 1, (tmpV4.y + ownerLoc.y) * base);
					instancedData.set(o + 2, (tmpV4.z + ownerLoc.z) * base);
				}

			case 1: // Face
				var positions = owner.data.geom.positions.values;
				var osx1 = owner.transform.scale.x;
				var osy1 = owner.transform.scale.y;
				var osz1 = owner.transform.scale.z;
				for (idx in lastSpawnedCount...targetCount) {
					var ia = owner.data.geom.indices[Std.random(owner.data.geom.indices.length)];
					var faceIndex = Std.random(Std.int(ia.length / 3));
					var i0 = ia[faceIndex * 3 + 0];
					var i1 = ia[faceIndex * 3 + 1];
					var i2 = ia[faceIndex * 3 + 2];
					var v0x = positions[i0 * 4    ], v0y = positions[i0 * 4 + 1], v0z = positions[i0 * 4 + 2];
					var v1x = positions[i1 * 4    ], v1y = positions[i1 * 4 + 1], v1z = positions[i1 * 4 + 2];
					var v2x = positions[i2 * 4    ], v2y = positions[i2 * 4 + 1], v2z = positions[i2 * 4 + 2];
					var rx = Math.random(); var ry = Math.random(); if (rx + ry > 1) { rx = 1 - rx; ry = 1 - ry; }
					var pxs = v0x + rx * (v1x - v0x) + ry * (v2x - v0x);
					var pys = v0y + rx * (v1y - v0y) + ry * (v2y - v0y);
					var pzs = v0z + rx * (v1z - v0z) + ry * (v2z - v0z);
					var px = pxs * normFactor * scalePosOwner * osx1;
					var py = pys * normFactor * scalePosOwner * osy1;
					var pz = pzs * normFactor * scalePosOwner * osz1;
					tmpV4.x = px; tmpV4.y = py; tmpV4.z = pz; tmpV4.w = 1;
					tmpV4.applyQuat(ownerRot);
					var o1 = idx * 3;
					instancedData.set(o1    , (tmpV4.x + ownerLoc.x) * base);
					instancedData.set(o1 + 1, (tmpV4.y + ownerLoc.y) * base);
					instancedData.set(o1 + 2, (tmpV4.z + ownerLoc.z) * base);
				}

			case 2: // Volume
				var dim = object.transform.dim;
				for (idx in lastSpawnedCount...targetCount) {
					tmpV4.x = (Math.random() * 2.0 - 1.0) * (dim.x * 0.5);
					tmpV4.y = (Math.random() * 2.0 - 1.0) * (dim.y * 0.5);
					tmpV4.z = (Math.random() * 2.0 - 1.0) * (dim.z * 0.5);
					tmpV4.w = 1;
					tmpV4.applyQuat(ownerRot);
					var o2 = idx * 3;
					instancedData.set(o2    , (tmpV4.x + ownerLoc.x) * base);
					instancedData.set(o2 + 1, (tmpV4.y + ownerLoc.y) * base);
					instancedData.set(o2 + 2, (tmpV4.z + ownerLoc.z) * base);
				}
		}

		// Upload full active range [0..targetCount) to this object's instanced VB
		var geom = object.data.geom;
		if (geom.instancedVB == null) {
			geom.setupInstanced(instancedData, 1, Usage.DynamicUsage);
		}
		var vb = geom.instancedVB.lock();
		var totalFloats = targetCount * 3; // xyz per instance
		var i = 0;
		while (i < totalFloats) {
			vb.setFloat32(i * 4, instancedData[i]);
			i++;
		}
		geom.instancedVB.unlock();
		geom.instanceCount = targetCount;
		lastSpawnedCount = targetCount;
	}

	inline function fhash(n: Int): Float {
    var s = n + 1.0;
    s *= 9301.0 % s;
    s = (s * 9301.0 + 49297.0) % 233280.0;
    return s / 233280.0;
}

	public function remove() {}

	/**
		Generates a random point in the triangle with vertex positions abc.
		
		Please note that the given position vectors are changed in-place by this
		function and can be considered garbage afterwards, so make sure to clone
		them first if needed.
	**/
	public static inline function randomPointInTriangle(a: Vec3, b: Vec3, c: Vec3): Vec3 {
		// Generate a random point in a square where (0, 0) <= (x, y) < (1, 1)
		var x = Math.random();
		var y = Math.random();

		if (x + y > 1) {
			// We're in the upper right triangle in the square, mirror to lower left
			x = 1 - x;
			y = 1 - y;
		}

		// Transform the point to the triangle abc
		var u = b.sub(a);
		var v = c.sub(a);
		return a.add(u.mult(x).add(v.mult(y)));
	}
}

class Particle {
	public var i: Int;
	
	public var x = 0.0;
	public var y = 0.0;
	public var z = 0.0;

	public var cameraDistance: Float;

	public function new(i: Int) {
		this.i = i;
	}
}

#end