package ocelot.desktop.graphics

import ocelot.desktop.graphics.mesh.{Mesh3D, MeshInstance3D, MeshVertex3D}
import ocelot.desktop.graphics.render.InstanceRenderer
import ocelot.desktop.graphics.scene.{Scene3D, SceneMesh3D}
import ocelot.desktop.util.{Logging, Resource, Spritesheet}
import org.lwjgl.opengl.{ARBFramebufferObject, GL11, GL21, GL30}

import scala.collection.mutable

class Viewport3D(width: Int, height: Int) extends Resource with Logging {
  private[graphics] val textureColor = new Texture(width, height, GL21.GL_SRGB8, GL11.GL_UNSIGNED_BYTE, GL11.GL_RGB)
  private val textureDepth = new Texture(width, height, GL11.GL_DEPTH_COMPONENT, GL11.GL_FLOAT)

  private val framebuffer = ARBFramebufferObject.glGenFramebuffers()
  GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, framebuffer)

  GL30.glFramebufferTexture2D(GL30.GL_FRAMEBUFFER, GL30.GL_COLOR_ATTACHMENT0, GL11.GL_TEXTURE_2D, textureColor.texture,
    0)

  GL30.glFramebufferTexture2D(GL30.GL_FRAMEBUFFER, GL30.GL_DEPTH_ATTACHMENT, GL11.GL_TEXTURE_2D, textureDepth.texture,
    0)

  GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, 0)

  private type Instancer = InstanceRenderer[MeshVertex3D, MeshInstance3D]

  private val shaderProgram: ShaderProgram = new ShaderProgram("general_3d")
  shaderProgram.set("uTexture", 0)

  private val instancers = new mutable.HashMap[Mesh3D, Instancer]()
  private val usedMeshes = new mutable.HashSet[Mesh3D]()
  private val transparentNodes = new mutable.ArrayBuffer[SceneMesh3D]()

  def draw(scene: Scene3D): Unit = {
    usedMeshes.clear()

    setupCamera(scene)
    setupLight(scene)
    collectOpaque(scene)
    collectTransparent(scene)
    sortOpaque(scene)
    sortTransparent(scene)

    GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, framebuffer)
    GL11.glEnable(GL30.GL_FRAMEBUFFER_SRGB)

    GL11.glViewport(0, 0, width, height)
    GL11.glDisable(GL11.GL_SCISSOR_TEST)

    GL11.glEnable(GL11.GL_CULL_FACE)
    GL11.glEnable(GL11.GL_DEPTH_TEST)
    GL11.glDepthFunc(GL11.GL_LEQUAL)

    GL11.glEnable(GL11.GL_BLEND)
    GL11.glBlendFunc(GL11.GL_ONE, GL11.GL_ONE_MINUS_SRC_ALPHA)

    GL11.glClearColor(0.0f, 0.0f, 0.0f, 1.0f)
    GL11.glClearDepth(1.0f)
    GL11.glClear(GL11.GL_COLOR_BUFFER_BIT | GL11.GL_DEPTH_BUFFER_BIT)

    Spritesheet.texture.bind()
    drawOpaque()

    // TODO: consider AABB
    GL11.glDisable(GL11.GL_CULL_FACE)

    // depth pre-pass
    GL11.glColorMask(false, false, false, false)
    drawTransparent()

    GL11.glColorMask(true, true, true, true)
    GL11.glDepthMask(false)
    drawTransparent()
    GL11.glDepthMask(true)

    instancers.filterInPlace { case (mesh, instancer) =>
      val retain = usedMeshes.contains(mesh)
      if (!retain) instancer.freeResource()
      retain
    }
  }

  private def setupCamera(scene: Scene3D): Unit = {
    shaderProgram.set("uView", scene.camera.transform.inverse)
    shaderProgram.set("uProj", scene.camera.projectionMatrix(width.toFloat / height.toFloat))
  }

  private def setupLight(scene: Scene3D): Unit = {
    shaderProgram.set("uLightDir", scene.lightSource.basis.up)
    shaderProgram.set("uLightColor", scene.lightSource.color.toRGBANorm.toLinear.rgbVector * scene.lightSource.energy)
    shaderProgram.set("uAmbientLightColor",
      scene.ambientLightColor.toRGBANorm.toLinear.rgbVector * scene.ambientLightEnergy)
  }

  private def collectOpaque(scene: Scene3D): Unit = {
    for (node <- scene.meshes.iterator.filter(_.isOpaque)) {
      // TODO: frustum culling
      val instancer = getInstancer(node.mesh)
      usedMeshes += node.mesh
      instancer.schedule(nodeToInstance(node))
    }
  }

  private def collectTransparent(scene: Scene3D): Unit = {
    transparentNodes.clear()
    for (node <- scene.meshes.iterator.filter(!_.isOpaque)) {
      usedMeshes += node.mesh
      transparentNodes.append(node)
    }
  }

  private def sortOpaque(scene: Scene3D): Unit = {
    for ((mesh, instancer) <- instancers) {
      val meshCenter = mesh.boundingBox.center
      instancer.sort((a, b) => {
        val aDist = scene.camera.distanceTo(a.transform * meshCenter)
        val bDist = scene.camera.distanceTo(b.transform * meshCenter)
        aDist < bDist // closest first
      })
    }
  }

  private def sortTransparent(scene: Scene3D): Unit = {
    transparentNodes.sortWith((a, b) => {
      val aDist = scene.camera.distanceTo(a.transform * a.mesh.boundingBox.center)
      val bDist = scene.camera.distanceTo(b.transform * b.mesh.boundingBox.center)
      a.priority < b.priority || aDist > bDist // farthest first
    })
  }

  private def drawOpaque(): Unit = {
    instancers.values.foreach(_.flush())
  }

  private def drawTransparent(): Unit = {
    var i = 0
    while (i < transparentNodes.length) {
      val node = transparentNodes(i)
      i += 1

      val instancer = getInstancer(node.mesh)
      instancer.schedule(nodeToInstance(node))

      while (i < transparentNodes.length && transparentNodes(i).mesh == node.mesh) {
        val similarNode = transparentNodes(i)
        instancer.schedule(nodeToInstance(similarNode))
        i += 1
      }

      instancer.flush()
    }
  }

  private def getInstancer(mesh: Mesh3D): Instancer = {
    instancers.getOrElseUpdate(mesh, new Instancer(mesh, MeshInstance3D, shaderProgram))
  }

  private def nodeToInstance(node: SceneMesh3D): MeshInstance3D = {
    MeshInstance3D(node.color.toRGBANorm, node.transform, node.uvTransform, node.lightingFactor, node.textureFactor)
  }

  override def freeResource(): Unit = {
    super.freeResource()
    textureColor.freeResource()
    textureDepth.freeResource()
    shaderProgram.freeResource()
    instancers.values.foreach(_.freeResource())

    ARBFramebufferObject.glDeleteFramebuffers(framebuffer)
  }
}
