ocelot-desktop/src/main/scala/ocelot/desktop/graphics/Graphics.scala

package ocelot.desktop.graphics

import ocelot.desktop.color.{Color, RGBAColorNorm}
import ocelot.desktop.geometry.{Rect2D, Size2D, Transform2D, Vector2D}
import ocelot.desktop.graphics.Texture.MinFilteringMode
import ocelot.desktop.graphics.mesh.{Mesh2D, MeshInstance2D, MeshVertex2D}
import ocelot.desktop.graphics.render.InstanceRenderer
import ocelot.desktop.ui.UiHandler
import ocelot.desktop.util.{Logging, Resource, Spritesheet}
import org.lwjgl.BufferUtils
import org.lwjgl.opengl.{ARBFramebufferObject, GL11, GL21, GL30}

import java.awt.image.BufferedImage
import java.nio.ByteBuffer
import scala.collection.mutable
import scala.util.control.Breaks._

//noinspection ScalaWeakerAccess,ScalaUnusedSymbol
class Graphics(private var width: Int, private var height: Int, private var scalingFactor: Float)
    extends Logging with Resource {
  private var time = 0f

  private var projection = Transform2D.viewport(width, height)

  private val shaderProgram = new ShaderProgram("general")
  private val renderer = new InstanceRenderer[MeshVertex2D, MeshInstance2D](Mesh2D.quad, MeshInstance2D, shaderProgram)

  private var _font: Font = Font.NormalFont
  private var oldFont: Font = _font

  private val stack = mutable.Stack[GraphicsState](GraphicsState())

  private var spriteRect = Spritesheet.sprites("Empty")

  private val emptySpriteTrans =
    Transform2D.translate(spriteRect.x, spriteRect.y) >> Transform2D.scale(spriteRect.w, spriteRect.h)

  private val offscreenTexture = new Texture(width, height, GL21.GL_SRGB8_ALPHA8, GL11.GL_UNSIGNED_BYTE, GL11.GL_RGBA)

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

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

  private var currentFramebuffer = 0
  private val currentScissor: Option[Rect2D] = None

  shaderProgram.set("uTexture", 0)
  shaderProgram.set("uTextTexture", 1)

  scale(scalingFactor)

  private[graphics] val screenShaderProgram = new ShaderProgram("general")
  screenShaderProgram.set("uTexture", 0)
  screenShaderProgram.set("uTextTexture", 1)

  def resize(width: Int, height: Int, scaling: Float): Boolean = {
    var viewportChanged = false

    if (scaling != scalingFactor) {
      scalingFactor = scaling
      stack.last.transform = Transform2D.scale(scalingFactor)
      viewportChanged = true
    }

    if (this.width != width || this.height != height) {
      this.width = width
      this.height = height

      offscreenTexture.bind()
      GL11.glTexImage2D(GL11.GL_TEXTURE_2D, 0, GL21.GL_SRGB8_ALPHA8, width, height, 0, GL11.GL_RGBA,
        GL11.GL_UNSIGNED_BYTE, null.asInstanceOf[ByteBuffer])

      viewportChanged = true
    }

    viewportChanged
  }

  override def freeResource(): Unit = {
    super.freeResource()

    GL30.glDeleteFramebuffers(offscreenFramebuffer)

    offscreenTexture.freeResource()
    Spritesheet.freeResource()
    renderer.freeResource()
    shaderProgram.freeResource()
    screenShaderProgram.freeResource()
  }

  def font: Font = _font

  def setNormalFont(): Unit = {
    if (_font == Font.NormalFont) return
    flush()
    _font = Font.NormalFont
  }

  def setSmallFont(): Unit = {
    if (_font == Font.SmallFont) return
    flush()
    _font = Font.SmallFont
  }

  def save(): Unit = {
    stack.push(stack.head.copy())
  }

  def restore(): Unit = {
    if (stack.length == 1)
      throw new RuntimeException("attempt to pop root graphics state")

    val old = stack.last
    if (old.scissor != stack.head.scissor) {
      flush()
    }

    stack.pop()
  }

  def beginGroupAlpha(): Unit = {
    flush()
    GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, offscreenFramebuffer)
    GL11.glDisable(GL11.GL_SCISSOR_TEST)
    GL11.glViewport(0, 0, width, height)
    GL11.glClearColor(0, 0, 0, 0)
    GL11.glClear(GL11.GL_COLOR_BUFFER_BIT)
    currentFramebuffer = offscreenFramebuffer
  }

  def endGroupAlpha(alpha: Float): Unit = {
    flush()
    currentFramebuffer = 0
    foreground = RGBAColorNorm(1, 1, 1, alpha)
    spriteRect = Rect2D(0, 1f, 1f, -1f)
    _rect(0, 0, width / scalingFactor, height / scalingFactor, fixUV = false)
    flush(mainTexture = offscreenTexture)
  }

  def blitViewport3D(viewport: Viewport3D, bounds: Rect2D, alpha: Float = 1.0f): Unit = {
    flush()
    foreground = RGBAColorNorm(1, 1, 1, alpha)
    spriteRect = Rect2D(0, 1f, 1f, -1f)
    _rect(bounds.x, bounds.y, bounds.w, bounds.h, fixUV = false)
    flush(mainTexture = viewport.textureColor)
  }

  def blitScreenViewport(
    viewport: ScreenViewport,
    bounds: Rect2D,
    filteringMode: MinFilteringMode = MinFilteringMode.Nearest,
    alpha: Float = 1.0f,
  ): Unit = {
    flush()
    foreground = RGBAColorNorm(1, 1, 1, alpha)
    spriteRect = Rect2D(0, 1f, 1f, -1f)
    _rect(bounds.x, bounds.y, bounds.w, bounds.h, fixUV = false)

    viewport.texture.setMinFilter(filteringMode)

    flush(mainTexture = viewport.texture)
  }

  def begin(): Unit = {
    shaderProgram.set("uProj", Transform2D.viewport(width, height))
  }

  def clear(): Unit = {
    flush(clear = true)
  }

  def setScissor(x: Float, y: Float, width: Float, height: Float): Unit = {
    flush()
    stack.head.scissor = Some((x, y, width, height))
  }

  def clearScissor(): Unit = {
    flush()
    stack.head.scissor = None
  }

  def clear(x: Int, y: Int, width: Int, height: Int): Unit = {
    save()
    setScissor(x, y, width, height)
    clear()
    restore()
  }

  def foreground: RGBAColorNorm = stack.head.foreground

  def foreground_=(col: Color): Unit = {
    stack.head.foreground = col.toRGBANorm
  }

  def background: RGBAColorNorm = stack.head.background

  def background_=(col: Color): Unit = {
    stack.head.background = col.toRGBANorm
  }

  def fontSizeMultiplier: Float = stack.head.fontSizeMultiplier

  def fontSizeMultiplier_=(value: Float): Unit = {
    stack.head.fontSizeMultiplier = value
  }

  def alphaMultiplier: Float = stack.head.alphaMultiplier

  def alphaMultiplier_=(value: Float): Unit = {
    stack.head.alphaMultiplier = value
  }

  def sprite: String = stack.head.sprite

  def sprite_=(value: String): Unit = {
    stack.head.sprite = value
    spriteRect = Spritesheet.sprites(value)
  }

  def transform(t: Transform2D): Unit = {
    stack.head.transform = stack.head.transform >> t
  }

  def scale(v: Float): Unit = {
    transform(Transform2D.scale(v))
  }

  def scale(x: Float, y: Float): Unit = {
    transform(Transform2D.scale(x, y))
  }

  def translate(x: Float, y: Float): Unit = {
    transform(Transform2D.translate(x, y))
  }

  def rotate(angle: Float): Unit = {
    transform(Transform2D.rotate(angle))
  }

  def text(x: Float, y: Float, text: String, shrink: Int = 0): Unit = {
    var ox = x

    text.codePoints().forEach { c =>
      char(ox, y, c)
      ox += _font.charWidth(c) - shrink
    }
  }

  def char(x: Float, y: Float, c: Int): Unit = {
    val fontSize = fontSizeMultiplier * _font.fontSize
    val rect = _font.map.getOrElse(c, _font.map('?'))

    val uvTransform = Transform2D.translate(
      rect.x,
      rect.y,
    ) >> Transform2D.scale(
      rect.w - 0.25f / _font.AtlasWidth,
      rect.h - 0.25f / _font.AtlasHeight,
    )

    val transform = {
      stack.head.transform >>
        Transform2D.translate(x.round, y.round) >>
        Transform2D.scale(_font.charWidth(c), fontSize)
    }

    val foreground = stack.head.foreground.toRGBANorm.mapA(_ * alphaMultiplier)
    val background = stack.head.background.toRGBANorm.mapA(_ * alphaMultiplier)

    if (background.a > 0 || foreground.a > 0)
      renderer.schedule(MeshInstance2D(background, foreground, transform, emptySpriteTrans, uvTransform))
  }

  // I hate scala. Overloaded methods with default arguments are not allowed
  def sprite(icon: IconSource, bounds: Rect2D): Unit = {
    sprite(icon, bounds.x, bounds.y, bounds.w, bounds.h, Color.White)
  }

  def sprite(icon: IconSource, bounds: Rect2D, color: Color): Unit = {
    sprite(icon, bounds.x, bounds.y, bounds.w, bounds.h, color)
  }

  def sprite(icon: IconSource, pos: Vector2D, size: Size2D): Unit = {
    sprite(icon, pos.x, pos.y, size.width, size.height, Color.White)
  }

  def sprite(icon: IconSource, pos: Vector2D, color: Color): Unit = {
    sprite(icon, pos.x, pos.y, color)
  }

  def sprite(icon: IconSource, pos: Vector2D, size: Size2D, color: Color): Unit = {
    sprite(icon, pos.x, pos.y, size.width, size.height, color)
  }

  def sprite(icon: IconSource, x: Float, y: Float): Unit = {
    sprite(icon, x, y, Color.White)
  }

  def sprite(icon: IconSource, x: Float, y: Float, width: Float, height: Float): Unit = {
    sprite(icon, x, y, width, height, Color.White)
  }

  def sprite(icon: IconSource, x: Float, y: Float, color: Color): Unit = {
    val size = Spritesheet.spriteSize(icon.path)
    sprite(icon, x, y, size.width, size.height, color)
  }

  def sprite(
    icon: IconSource,
    x: Float,
    y: Float,
    width: Float,
    height: Float,
    color: Color,
  ): Unit = {
    sprite = icon.path
    foreground = color

    val spriteRect = icon.animation.map { animation =>
      val duration = animation.frames.map(_._2).sum
      var timeOffset = 0f
      var curFrame = 0

      breakable {
        for ((idx, dur) <- animation.frames) {
          timeOffset += dur
          curFrame = idx
          if (timeOffset >= time % duration) break
        }
      }

      val size = animation.frameSize match {
        case Some(size) => Size2D(this.spriteRect.w, this.spriteRect.w * size.height / size.width)
        case None => Size2D(this.spriteRect.w, this.spriteRect.w)
      }

      this.spriteRect.copy(y = this.spriteRect.y + curFrame * size.height, h = size.height)
    }

    _rect(x, y, width, height, fixUV = true, spriteRect)
  }

  def sprite(
    name: String,
    x: Float,
    y: Float,
    width: Float,
    height: Float,
    color: Color,
    spriteRect: Option[Rect2D],
    fixUV: Boolean = true,
  ): Unit = {
    sprite = name
    foreground = color

    _rect(x, y, width, height, fixUV, spriteRect)
  }

  def rect(r: Rect2D, color: Color): Unit = {
    rect(r.x, r.y, r.w, r.h, color)
  }

  def rect(x: Float, y: Float, width: Float, height: Float, color: Color = RGBAColorNorm(1f, 1f, 1f)): Unit = {
    sprite(IconSource.Empty, x, y, width, height, color)
  }

  private def checkFont(): Unit = {
    if (_font != oldFont) {
      val newFont = _font
      _font = oldFont
      flush()
      _font = newFont
      oldFont = _font
    }
  }

  private def _rect(x: Float, y: Float, width: Float, height: Float,
                    fixUV: Boolean = true,
                    spriteRectOptional: Option[Rect2D] = None): Unit = {
    val spriteRect = spriteRectOptional.getOrElse(this.spriteRect)

    val uvTransform = Transform2D.translate(spriteRect.x, spriteRect.y) >>
      (if (fixUV)
         Transform2D.scale(spriteRect.w - 0.25f / 1024, spriteRect.h - 0.25f / 1024)
       else
         Transform2D.scale(spriteRect.w, spriteRect.h))

    val transform = {
      stack.head.transform >>
        Transform2D.translate(x, y) >>
        Transform2D.scale(width, height)
    }

    val color = stack.head.foreground.toRGBANorm.mapA(_ * alphaMultiplier)

    if (color.a > 0)
      renderer.schedule(MeshInstance2D(color, Color.Transparent, transform, uvTransform, Transform2D.scale(0)))
  }

  def line(x1: Float, y1: Float, x2: Float, y2: Float, thickness: Float, color: Color): Unit = {
    save()

    val dy = x2 - x1
    val dx = y2 - y1
    val length = math.sqrt(dx * dx + dy * dy).toFloat
    val inclination = math.atan2(dy, dx).toFloat
    translate(x1, y1)
    rotate(-inclination + (math.Pi * 0.5).toFloat)
    rect(0, -thickness * 0.5f, length, thickness, color)

    restore()
  }

  def line(start: Vector2D, end: Vector2D, thickness: Float, color: Color): Unit = {
    line(start.x, start.y, end.x, end.y, thickness, color)
  }

  def screenshot(): BufferedImage = {
    val buffer = BufferUtils.createByteBuffer(width * height * 4)
    GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, 0)
    GL11.glReadPixels(0, 0, width, height, GL11.GL_RGBA, GL11.GL_UNSIGNED_BYTE, buffer)
    val image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB)
    for (y <- 0 until height) {
      for (x <- 0 until width) {
        val i = (x + (height - y - 1) * width) * 4
        val r = buffer.get(i) & 0xff
        val g = buffer.get(i + 1) & 0xff
        val b = buffer.get(i + 2) & 0xff
        val rgb = (r << 16) | (g << 8) | b
        image.setRGB(x, y, rgb)
      }
    }
    image
  }

  def flush(mainTexture: Texture = Spritesheet.texture, clear: Boolean = false): Unit = {
    if (renderer.isEmpty && !clear) return

    GL30.glBindFramebuffer(GL30.GL_FRAMEBUFFER, currentFramebuffer)
    GL11.glEnable(GL30.GL_FRAMEBUFFER_SRGB)
    GL11.glViewport(0, 0, width, height)

    GL11.glDisable(GL11.GL_DEPTH_TEST)
    GL11.glDisable(GL11.GL_CULL_FACE)

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

    stack.head.scissor match {
      case Some((x, y, w, h)) =>
        GL11.glEnable(GL11.GL_SCISSOR_TEST)
        GL11.glScissor(
          Math.round(x * scalingFactor),
          Math.round(height - h * scalingFactor - y * scalingFactor),
          Math.round(w * scalingFactor),
          Math.round(h * scalingFactor),
        )
      case _ =>
        GL11.glDisable(GL11.GL_SCISSOR_TEST)
    }

    mainTexture.bind()
    _font.texture.bind(1)
    renderer.flush()

    if (clear) {
      GL11.glClearColor(1, 1, 1, 1)
      GL11.glClear(GL11.GL_COLOR_BUFFER_BIT)
    }
  }

  def update(): Unit = {
    time += UiHandler.dt * 20f
  }
}