module Render.Unlit.Textured.Pipeline
  ( Pipeline
  , allocate
  , allocateBlend
  ) where

import RIO
import Render.DescSets.Set0.Code

import Control.Monad.Trans.Resource (ResourceT)
import Data.Tagged (Tagged(..))
import Vulkan.Core10 qualified as Vk
import Vulkan.Zero (zero)

import Engine.Vulkan.Pipeline qualified as Pipeline
import Engine.Vulkan.Types (HasVulkan, HasRenderPass(..), DsBindings)
import Render.Code (compileVert, compileFrag, glsl)
import Render.DescSets.Set0 (Scene, vertexPos, instanceTransform)
import Render.Unlit.Textured.Model qualified as Model

type Config = Pipeline.Config '[Scene] Model.VertexAttrs Model.InstanceAttrs
type Pipeline = Pipeline.Pipeline '[Scene] Model.VertexAttrs Model.InstanceAttrs

allocate
  :: ( HasVulkan env
     , HasRenderPass renderpass
     )
  => Vk.SampleCountFlagBits
  -> Tagged Scene DsBindings
  -> renderpass
  -> ResourceT (RIO env) Pipeline
allocate multisample tset0 rp = do
  (_, p) <- Pipeline.allocate
    Nothing
    multisample
    (config tset0)
    rp
  pure p

allocateBlend
  :: ( HasVulkan env
     , HasRenderPass renderpass
     )
  => Vk.SampleCountFlagBits
  -> Tagged Scene DsBindings
  -> renderpass
  -> ResourceT (RIO env) Pipeline
allocateBlend multisample tset0 rp = do
  (_, p) <- Pipeline.allocate
    Nothing
    multisample
    (configBlend tset0)
    rp
  pure p

config :: Tagged Scene DsBindings -> Config
config (Tagged set0) = zero
  { Pipeline.cDescLayouts  = Tagged @'[Scene] [set0]
  , Pipeline.cVertexCode   = Just vertCode
  , Pipeline.cVertexInput  = vertexInput
  , Pipeline.cFragmentCode = Just fragCode
  }
  where
    vertexInput = Pipeline.vertexInput
      [ vertexPos -- vPosition
      , (Vk.VERTEX_INPUT_RATE_VERTEX,   Model.vkVertexAttrs)
      , (Vk.VERTEX_INPUT_RATE_INSTANCE, Model.vkInstanceTexture)
      , instanceTransform
      ]

configBlend :: Tagged Scene DsBindings -> Config
configBlend tset0 = (config tset0)
  { Pipeline.cBlend      = True
  , Pipeline.cDepthWrite = False
  }

vertCode :: ByteString
vertCode =
  $(compileVert [glsl|
    #version 450
    #extension GL_ARB_separate_shader_objects : enable

    ${set0binding0}

    // vertexPos
    layout(location = 0) in vec3 vPosition;
    // vertexAttrs
    layout(location = 1) in vec2 vTexCoord;
    // textureParams
    layout(location = 2) in  vec4 iTextureScaleOffset;
    layout(location = 3) in  vec4 iTextureGamma;
    layout(location = 4) in ivec2 iTextureIds;

    // transformMat
    layout(location = 5) in mat4 iModel;

    layout(location = 0)      out  vec2 fTexCoord;
    layout(location = 1) flat out  vec4 fTextureGamma;
    layout(location = 2) flat out ivec2 fTextureIds;

    void main() {
      gl_Position
        = scene.projection
        * scene.view
        * iModel
        * vec4(vPosition, 1.0);

      fTexCoord     = vTexCoord * iTextureScaleOffset.st + iTextureScaleOffset.pq;
      fTextureGamma = iTextureGamma;
      fTextureIds   = iTextureIds;
    }
  |])

fragCode :: ByteString
fragCode =
  $(compileFrag [glsl|
    #version 450
    #extension GL_ARB_separate_shader_objects : enable
    #extension GL_EXT_nonuniform_qualifier : enable

    ${set0binding1}
    ${set0binding2}

    layout(location = 0)      in  vec2 fTexCoord;
    layout(location = 1) flat in  vec4 fTextureGamma;
    layout(location = 2) flat in ivec2 fTextureIds;

    layout(location = 0) out vec4 oColor;

    void main() {
      vec3 color = vec3(0);
      float combinedAlpha = 0;

      if (fTextureIds.t > -1 && fTextureIds.s > -1) {
        vec4 texel = texture(
          sampler2D(
            textures[nonuniformEXT(fTextureIds.t)],
            samplers[nonuniformEXT(fTextureIds.s)]
          ),
          fTexCoord
        );
        color = pow(texel.rgb, fTextureGamma.rgb);
        combinedAlpha = texel.a * fTextureGamma.a;
      }

      // XXX: premultiply alpha due to additive blending
      oColor = vec4(color * combinedAlpha, combinedAlpha);
    }
  |])