package iron.data;

#if lnx_batch

import kha.arrays.ByteArray;
import kha.graphics4.VertexBuffer;
import kha.graphics4.IndexBuffer;
import kha.graphics4.Usage;
import kha.graphics4.VertexStructure;
import kha.graphics4.Graphics;
import iron.object.MeshObject;
import iron.object.Uniforms;
import iron.data.Geometry;
import iron.data.MaterialData;
import iron.data.ShaderData;
import iron.data.SceneFormat;

@:access(iron.object.MeshObject)
class MeshBatch {

	var buckets: Map<ShaderData, Bucket> = new Map();
	public var nonBatched: Array<MeshObject> = [];

	public function new() {}

	public function remove() {
		for (b in buckets) b.remove();
	}

	public static function isBatchable(m: MeshObject): Bool {
		m.computeDepthRead();
		var batch =
			m.materials != null &&
			m.materials.length == 1 &&
			!m.data.geom.instanced &&
			!m.data.isSkinned &&
			m.data.raw.morph_target == null &&
			!m.depthRead;
		return batch;
	}

	public function addMesh(m: MeshObject, isLod: Bool): Bool {
		// No instancing, multimat, skinning, morph targets or lod batching
		if (!isBatchable(m) || isLod) {
			nonBatched.push(m);
			return false;
		}

		var shader = m.materials[0].shader;
		var b = buckets.get(shader);
		if (b == null) {
			b = new Bucket(shader);
			buckets.set(shader, b);
		}
		b.addMesh(m);
		return true;
	}

	public function removeMesh(m: MeshObject) {
		var shader = m.materials[0].shader;
		var b = buckets.get(shader);
		if (b != null) b.removeMesh(m);
		else nonBatched.remove(m);
	}

	@:access(iron.RenderPath)
	public function render(g: Graphics, context: String, bindParams: Array<String>) {

		for (b in buckets) {
			if (!b.batched) b.batch();
			if (b.meshes.length > 0 && b.meshes[0].cullMaterial(context)) continue;

			var scontext = b.shader.getContext(context);
			if (scontext == null) continue;
			g.setPipeline(scontext.pipeState);
			#if lnx_deinterleaved
			g.setVertexBuffers(b.getVertexBuffers(scontext.raw.vertex_elements));
			#else
			g.setVertexBuffer(b.getVertexBuffer(scontext.raw.vertex_elements));
			#end
			g.setIndexBuffer(b.indexBuffer);

			Uniforms.setContextConstants(g, scontext, bindParams);

			var isShadowCtx = context == "shadowmap" || context == "shadowmap_transparent";
			if (!isShadowCtx) RenderPath.sortMeshesDistance(b.meshes);

			for (m in b.meshes) {
				if (!m.visible) continue; // Skip render if object is hidden
				#if lnx_occlusion_culling
				if (!m.occlusionVisible && (context == "mesh" || context == "depth" || context == "translucent")) continue;
				if (!m.shadowOcclusionVisible && isShadowCtx) continue;
				#end
				if (m.cullMesh(context, Scene.active.camera, RenderPath.active.light)) continue;

				// Get context
				var materialContexts: Array<MaterialContext> = null;
				var shaderContexts: Array<ShaderContext> = null;
				if (m.cachedMaterialContexts != null && MeshObject.matsEqual(m.cachedMats.get(context), m.materials)) {
					materialContexts = m.cachedMaterialContexts.get(context);
					shaderContexts = m.cachedShaderContexts.get(context);
				}
				if (materialContexts == null) {
					materialContexts = [];
					shaderContexts = [];
					m.getContexts(context, m.materials, materialContexts, shaderContexts);
					if (m.cachedMaterialContexts == null) {
						m.cachedMats = new Map();
						m.cachedMaterialContexts = new Map();
						m.cachedShaderContexts = new Map();
					}
					m.cachedMats.set(context, m.materials);
					m.cachedMaterialContexts.set(context, materialContexts);
					m.cachedShaderContexts.set(context, shaderContexts);
				}

				Uniforms.posUnpack = m.data.scalePos;
				Uniforms.texUnpack = m.data.scaleTex;
				if (!isShadowCtx) m.transform.update();

				// Render mesh
				Uniforms.setObjectConstants(g, scontext, m);
				Uniforms.setMaterialConstants(g, scontext, materialContexts[0]);

				g.drawIndexedVertices(m.data.start, m.data.count);

				#if lnx_veloc
				m.prevMatrix.setFrom(m.transform.worldUnpack);
				#end

				#if lnx_debug
				RenderPath.drawCalls++;
				RenderPath.batchCalls++;
				#end
			}

			#if lnx_debug
			RenderPath.batchBuckets++;
			#end
		}

		// Render non-batched meshes
		inline RenderPath.meshRenderLoop(g, context, bindParams, nonBatched);
	}
}

class Bucket {

	public var batched = false;
	public var shader: ShaderData;
#if lnx_deinterleaved
	public var vertexBuffers: Array<InterleavedVertexBuffer> = [];
#else
	var vertexBuffer: VertexBuffer;
	var vertexBufferMap: Map<String, VertexBuffer> = new Map();
#end
	public var indexBuffer: IndexBuffer;
	public var meshes: Array<MeshObject> = [];

	public function new(shader: ShaderData) {
		this.shader = shader;
	}

	public function remove() {
		indexBuffer.delete();
		#if lnx_deinterleaved
		for (buf in vertexBuffers) buf.buffer.delete();
		#else
		// this.vertexBuffer is in the map, so it's also deleted here
		for (buf in vertexBufferMap) buf.delete();
		#end
		meshes = [];
	}

	public function addMesh(m: MeshObject) {
		meshes.push(m);
	}

	public function removeMesh(m: MeshObject) {
		meshes.remove(m);
	}

#if lnx_deinterleaved
	public function getVertexBuffers(elems: Array<TVertexElement>): Array<VertexBuffer> {
		var vbs: Array<VertexBuffer> = [];
		for (e in elems) {
			for (v in vertexBuffers) {
				if (v.name == e.name) {
					vbs.push(v.buffer);
					break;
				}
			}
		}
		return vbs;
	}
#else
	function copyAttribute(attribSize: Int, count: Int, to: ByteArray, toStride: Int, toOffset: Int, from: ByteArray, fromStride: Int, fromOffset: Int) {
		for (i in 0...count) {
			for (j in 0...attribSize) {
				to.setInt16((i * toStride + toOffset + j) * 2, from.getInt16((i * fromStride + fromOffset + j) * 2));
			}
		}
	}

	function extractVertexBuffer(elems: Array<TVertexElement>): VertexBuffer {
		// Build vertex buffer for specific context
		var vs = new VertexStructure();
		for (e in elems) vs.add(e.name, ShaderContext.parseData(e.data));

		var vb = new VertexBuffer(vertexBuffer.count(), vs, Usage.StaticUsage);
		var to = vb.lock();
		var from = vertexBuffer.lock();

		var toOffset = 0;
		var toStride = Std.int(vb.stride() / 2);
		var fromOffset = 0;
		var fromStride = Std.int(vertexBuffer.stride() / 2);

		for (e in elems) {
			var size = 0;
			if (e.name == "pos") { size = 4; fromOffset = 0; }
			else if (e.name == "nor") { size = 2; fromOffset = 4; }
			else if (e.name == "tex") { size = 2; fromOffset = 6; }
			copyAttribute(size, vertexBuffer.count(), to, toStride, toOffset, from, fromStride, fromOffset);
			toOffset += size;
		}

		vb.unlock();
		return vb;
	}

	public function getVertexBuffer(elems: Array<TVertexElement>): VertexBuffer {
		var s = "";
		for (e in elems) s += e.name;
		var vb = vertexBufferMap.get(s);
		if (vb == null) {
			vb = extractVertexBuffer(elems);
			vertexBufferMap.set(s, vb);
		}
		return vb;
	}
#end

	function vertexCount(g: Geometry, hasUVs: Bool): Int {
		var vcount = g.getVerticesLength();
		if (hasUVs && g.uvs == null) {
			vcount += Std.int(g.positions.values.length / 4) * 2;
		}
		return vcount;
	}

	public function batch() {
		batched = true;

		// Ensure same vertex structure for batched meshes
		var hasUVs = false;
		for (m in meshes) {
			if (m.data.geom.uvs != null) {
				hasUVs = true;
				break;
			}
		}

		// Unique mesh datas
		var vcount = 0;
		var icount = 0;
		var mdatas: Array<MeshData> = [];
		for (m in meshes) {
			var mdFound = false;
			for (md in mdatas) {
				if (m.data == md) {
					mdFound = true;
					break;
				}
			}
			if (!mdFound) {
				mdatas.push(m.data);
				m.data.start = icount;
				m.data.count = m.data.geom.indices[0].length;
				icount += m.data.count;
				vcount += vertexCount(m.data.geom, hasUVs);
			}
		}

		if (mdatas.length == 0) return;

#if lnx_deinterleaved
		var attrNames: Array<String> = [];
		var attrReps = new Map<String, TVertexArray>();
		for (md in mdatas) {
			for (va in md.geom.vertexArrays) {
				if (attrReps.get(va.attrib) == null) {
					attrReps.set(va.attrib, va);
					attrNames.push(va.attrib);
				}
			}
		}
		for (name in attrNames) {
			var rep = attrReps.get(name);
			var attrStruct = new VertexStructure();
			attrStruct.add(name, ShaderContext.parseData(rep.data));
			var stride = Std.int(attrStruct.byteSize() / 2);
			var vtotal = 0;
			for (md in mdatas) vtotal += md.geom.getVerticesCount();
			var vb = new VertexBuffer(vtotal, attrStruct, Usage.StaticUsage);
			var vertices = vb.lock();
			var ai = -1;
			for (md in mdatas) {
				var va: TVertexArray = null;
				for (v in md.geom.vertexArrays) {
					if (v.attrib == name) {
						va = v;
						break;
					}
				}
				if (va != null && va.values != null) {
					var nverts = md.geom.getVerticesCount();
					for (i in 0...nverts) {
						for (o in 0...va.size) vertices.setInt16(++ai * 2, va.values[i * va.size + o]);
						for (o in va.size...stride) vertices.setInt16(++ai * 2, 0);
					}
				}
				else {
					for (i in 0...md.geom.getVerticesCount() * stride) vertices.setInt16(++ai * 2, 0);
				}
			}
			vb.unlock();
			vertexBuffers.push({name: name, buffer: vb});
		}
#else
		// Pick UVs if present
		var vs = mdatas[0].geom.struct;
		for (md in mdatas) if (md.geom.struct.size() > vs.size()) vs = md.geom.struct;

		// Build shared buffers
		vertexBuffer = new VertexBuffer(vcount, vs, Usage.StaticUsage);
		var vertices = vertexBuffer.lock();
		var offset = 0;
		for (md in mdatas) {
			md.geom.copyVertices(vertices, offset, hasUVs);
			offset += vertexCount(md.geom, hasUVs);
		}
		vertexBuffer.unlock();

		var s = "";
		for (e in vs.elements) s += e.name;
		vertexBufferMap.set(s, vertexBuffer);
#end

		indexBuffer = new IndexBuffer(icount, Usage.StaticUsage);
		var indices = indexBuffer.lock();
		var di = -1;
		var offset = 0;
		for (md in mdatas) {
			for (i in 0...md.geom.indices[0].length) {
				indices[++di] = md.geom.indices[0][i] + offset;
			}
			offset += Std.int(md.geom.getVerticesLength() / md.geom.structLength);
		}
		indexBuffer.unlock();
	}
}

#end