package ocelot.desktop.graphics.render

import ocelot.desktop.graphics.ShaderProgram
import ocelot.desktop.graphics.buffer.{IndexBuffer, VertexBuffer}
import ocelot.desktop.graphics.mesh.{Mesh, Vertex, VertexArray, VertexType}
import ocelot.desktop.util.{Logging, Resource}
import org.lwjgl.BufferUtils
import org.lwjgl.opengl._

import java.nio.ByteBuffer
import scala.collection.mutable.ArrayBuffer

class InstanceRenderer[V <: Vertex, I <: Vertex](mesh: Mesh[V], instanceType: VertexType, shader: ShaderProgram)
    extends Resource with Logging {
  private val InitialCapacity: Int = 1

  private val vertexBuffer = new VertexBuffer[V](mesh.vertices)
  private val indexBuffer: Option[IndexBuffer] = mesh.indices.map(indices => new IndexBuffer(indices))
  private val instanceBuffer = new VertexBuffer[I](instanceType, InitialCapacity)

  private val vertexArray = new VertexArray(shader)

  vertexArray.addVertexBuffer(vertexBuffer)
  indexBuffer.foreach(ib => vertexArray.addIndexBuffer(ib))
  vertexArray.addVertexBuffer(instanceBuffer, instanced = true)

  private val instances = new ArrayBuffer[I]

  def schedule(instance: I): Unit = {
    instances += instance
  }

  def sort(ltComparator: (I, I) => Boolean): Unit = {
    instances.sortInPlaceWith(ltComparator)
  }

  def isEmpty: Boolean = instances.isEmpty

  def flush(): Unit = {
    write()
    draw()
  }

  private lazy val (glDrawArraysInstancedARBptr, glDrawElementsInstancedARBptr) = {
    val method = classOf[GLContext].getDeclaredMethod("getFunctionAddress", classOf[String])
    method.setAccessible(true)

    (
      method.invoke(null, "glDrawArraysInstancedARB"),
      method.invoke(null, "glDrawElementsInstancedARB"),
    )
  }

  private lazy val nglDrawElementsInstancedARB = {
    val method = classOf[ARBDrawInstanced].getDeclaredMethod("nglDrawElementsInstancedARB", Integer.TYPE, Integer.TYPE,
      Integer.TYPE, java.lang.Long.TYPE, Integer.TYPE, java.lang.Long.TYPE)
    method.setAccessible(true)
    method
  }

  private lazy val nglDrawArraysInstancedARB = {
    val method = classOf[ARBDrawInstanced].getDeclaredMethod("nglDrawArraysInstancedARB", Integer.TYPE, Integer.TYPE,
      Integer.TYPE, Integer.TYPE, java.lang.Long.TYPE)
    method.setAccessible(true)
    method
  }

  private def draw(): Unit = {
    if (instances.isEmpty) return

    shader.bind()
    vertexArray.bind()

    indexBuffer match {
      case Some(ib) =>
        nglDrawElementsInstancedARB.invoke(null, mesh.primitiveType.toGL, ib.capacity, GL11.GL_UNSIGNED_INT, 0L,
          instances.length, glDrawElementsInstancedARBptr)

      case None =>
        nglDrawArraysInstancedARB.invoke(null, mesh.primitiveType.toGL, 0, vertexBuffer.capacity, instances.length,
          glDrawArraysInstancedARBptr)
    }

    instances.clear()
  }

  private def write(): Unit = {
    if (instances.length > instanceBuffer.capacity) {
      instanceBuffer.resize(instances.length * 2)
      instanceBuffer.write(0, packInstances)
    } else {
      instanceBuffer.write(0, packInstances)
    }
  }

  private def packInstances: ByteBuffer = {
    val instanceStride = if (instances.isEmpty) 0 else instances(0).stride
    val data = BufferUtils.createByteBuffer(instances.length * instanceStride)

    for (element <- instances) {
      element.put(data)
    }

    data.flip
    data
  }

  override def freeResource(): Unit = {
    super.freeResource()
    vertexArray.freeResource()
    instanceBuffer.freeResource()
    indexBuffer.foreach(_.freeResource())
    vertexBuffer.freeResource()
  }
}
