package ocelot.desktop.graphics

import ocelot.desktop.geometry.{ProjectionMatrix3D, Transform2D, Transform3D, Vector3D}
import ocelot.desktop.util.{Logging, Resource}
import org.lwjgl.BufferUtils
import org.lwjgl.opengl.GL20

import scala.io.Source

class ShaderProgram(name: String) extends Logging with Resource {
  logger.info(s"Loading shader program ($name)")

  private val fragmentShader: Int = createShader(
    Source.fromResource(s"ocelot/desktop/shader/$name.frag", getClass.getClassLoader),
    GL20.GL_FRAGMENT_SHADER,
    "fragment",
  )

  private val vertexShader: Int = createShader(
    Source.fromResource(s"ocelot/desktop/shader/$name.vert", getClass.getClassLoader),
    GL20.GL_VERTEX_SHADER,
    "vertex",
  )

  val shaderProgram: Int = GL20.glCreateProgram()

  GL20.glAttachShader(shaderProgram, vertexShader)
  GL20.glAttachShader(shaderProgram, fragmentShader)
  GL20.glLinkProgram(shaderProgram)

  {
    val success = GL20.glGetProgrami(shaderProgram, GL20.GL_LINK_STATUS)

    if (success == 0) {
      val len = GL20.glGetProgrami(shaderProgram, GL20.GL_INFO_LOG_LENGTH)
      val info = GL20.glGetProgramInfoLog(shaderProgram, len)
      logger.error(s"Failed to link shader program ($name)\n$info")
      throw new RuntimeException("Shader linkage failed")
    }
  }

  override def freeResource(): Unit = {
    super.freeResource()
    GL20.glDeleteProgram(shaderProgram)
    GL20.glDeleteShader(vertexShader)
    GL20.glDeleteShader(fragmentShader)
  }

  def bind(): Unit = {
    GL20.glUseProgram(shaderProgram)
  }

  def set(name: String, texture: Texture): Unit = {
    bind()
    GL20.glUniform1i(getLocation(name), texture.texture)
  }

  def set(name: String, vector: Vector3D): Unit = {
    bind()
    GL20.glUniform3f(getLocation(name), vector.x, vector.y, vector.z)
  }

  def set(name: String, transform: Transform2D): Unit = {
    bind()

    val array = transform.array ++ Array(0f, 0f, 1f)
    val buffer = BufferUtils.createFloatBuffer(array.length)
    buffer.put(array)
    buffer.flip()
    GL20.glUniformMatrix3(getLocation(name), true, buffer)
  }

  def set(name: String, transform: Transform3D): Unit = {
    bind()

    val array = transform.array ++ Array(0f, 0f, 0f, 1f)
    val buffer = BufferUtils.createFloatBuffer(array.length)
    buffer.put(array)
    buffer.flip()
    GL20.glUniformMatrix4(getLocation(name), true, buffer)
  }

  def set(name: String, projection: ProjectionMatrix3D): Unit = {
    bind()

    val buffer = BufferUtils.createFloatBuffer(projection.array.length)
    buffer.put(projection.array)
    buffer.flip()
    GL20.glUniformMatrix4(getLocation(name), true, buffer)
  }

  def set(name: String, v: Int): Unit = {
    bind()
    GL20.glUniform1i(getLocation(name), v)
  }

  def getAttributeLocation(name: String): Int = {
    bind()
    GL20.glGetAttribLocation(shaderProgram, name)
  }

  private def getLocation(name: String): Int = {
    GL20.glGetUniformLocation(shaderProgram, name)
  }

  private def createShader(source: Source, ty: Int, s: String): Int = {
    val shader = GL20.glCreateShader(ty)

    GL20.glShaderSource(shader, source.mkString)
    GL20.glCompileShader(shader)

    val success = GL20.glGetShaderi(shader, GL20.GL_COMPILE_STATUS)

    if (success == 0) {
      val len = GL20.glGetShaderi(shader, GL20.GL_INFO_LOG_LENGTH)
      val info = GL20.glGetShaderInfoLog(shader, len)
      logger.error(s"Failed to compile shader ($s)\n$info")
      throw new RuntimeException("Shader compilation failed")
    }

    source.close()

    shader
  }
}
