module Render.Lit.Colored.Pipeline
  ( Pipeline
  , allocate
  , allocateBlend
  ) where

import RIO

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.Code.Lit (litMain, structLight, shadowFuns, brdfSpecular)
import Render.DescSets.Set0 (Scene, vertexPos, instanceTransform)
import Render.DescSets.Set0.Code (set0binding0, set0binding1, set0binding2, set0binding3, set0binding4, set0binding5)
import Render.Lit.Colored.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
      , (Vk.VERTEX_INPUT_RATE_VERTEX, Model.vkVertexAttrs)
      , instanceTransform
      ]

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

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

    ${set0binding0}

    layout(location = 0) in vec3 vPosition;
    layout(location = 1) in vec4 vBaseColor;
    layout(location = 2) in vec4 vEmissiveColor;
    layout(location = 3) in vec2 vMetallicRoughness;
    layout(location = 4) in vec3 vNormal;

    layout(location = 5) in mat4 iModel;

    layout(location = 0) out vec4 fPosition;
    layout(location = 1) out vec4 fColor;
    layout(location = 2) out vec4 fEmissiveColor;
    layout(location = 3) out vec2 fMetallicRoughness;
    layout(location = 4) out vec3 fNormal;

    void main() {
      fPosition = iModel * vec4(vPosition, 1.0);

      gl_Position
        = scene.projection
        * scene.view
        * fPosition;

      fNormal = transpose(mat3(inverse(iModel))) * vNormal; // TODO: use modelInv

      fColor = vBaseColor;
      fColor.rgb *= vBaseColor.a;

      fEmissiveColor = vEmissiveColor;
      fMetallicRoughness = vMetallicRoughness;
    }
  |])

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

    // TODO: move to spec constant
    const uint MAX_LIGHTS = 255;
    const float PCF_STEP = 1.5 / 4096;

    // TODO: move to material
    const float reflectivity = 1.0/256.0;

    ${structLight}

    ${set0binding0}
    ${set0binding1}
    ${set0binding2}
    ${set0binding3}
    ${set0binding4}
    ${set0binding5}

    layout(location = 0) in vec4 fPosition;
    layout(location = 1) in vec4 fColor;
    layout(location = 2) in vec4 fEmissiveColor;
    layout(location = 3) in vec2 fMetallicRoughness;
    layout(location = 4) in vec3 fNormal;

    layout(location = 0) out vec4 oColor;

    ${shadowFuns}
    ${brdfSpecular}

    void main() {
      vec4 baseColor = fColor; // XXX: assuming premultiplied alpha
      float metallic = fMetallicRoughness[0];
      float roughness = fMetallicRoughness[1];
      float occlusion = 1.0 - 0.0;

      vec3 normal = normalize(fNormal);

      ${litMain}

      oColor.rgb += pow(fEmissiveColor.rgb, vec3(2.2));
    }
  |])