module Render.Debug.Pipeline
  ( Pipeline
  , allocate
  , Mode(..)
  ) where

import RIO

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

import Engine.Vulkan.Pipeline qualified as Pipeline
import Engine.Vulkan.Shader qualified as Shader
import Engine.Vulkan.Types (HasVulkan, HasRenderPass(..), DsBindings)
import Render.Code (compileVert, compileFrag, glsl)
import Render.Debug.Model qualified as Model
import Render.DescSets.Set0 (Scene, vertexPos, instanceTransform)
import Render.DescSets.Set0.Code (set0binding0, set0binding1, set0binding2, set0binding3, set0binding5color)

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

data Mode
  = UV
  | Texture
  | Shadow Word32
  deriving (Mode -> Mode -> Bool
(Mode -> Mode -> Bool) -> (Mode -> Mode -> Bool) -> Eq Mode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Mode -> Mode -> Bool
$c/= :: Mode -> Mode -> Bool
== :: Mode -> Mode -> Bool
$c== :: Mode -> Mode -> Bool
Eq, Eq Mode
Eq Mode
-> (Mode -> Mode -> Ordering)
-> (Mode -> Mode -> Bool)
-> (Mode -> Mode -> Bool)
-> (Mode -> Mode -> Bool)
-> (Mode -> Mode -> Bool)
-> (Mode -> Mode -> Mode)
-> (Mode -> Mode -> Mode)
-> Ord Mode
Mode -> Mode -> Bool
Mode -> Mode -> Ordering
Mode -> Mode -> Mode
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Mode -> Mode -> Mode
$cmin :: Mode -> Mode -> Mode
max :: Mode -> Mode -> Mode
$cmax :: Mode -> Mode -> Mode
>= :: Mode -> Mode -> Bool
$c>= :: Mode -> Mode -> Bool
> :: Mode -> Mode -> Bool
$c> :: Mode -> Mode -> Bool
<= :: Mode -> Mode -> Bool
$c<= :: Mode -> Mode -> Bool
< :: Mode -> Mode -> Bool
$c< :: Mode -> Mode -> Bool
compare :: Mode -> Mode -> Ordering
$ccompare :: Mode -> Mode -> Ordering
$cp1Ord :: Eq Mode
Ord, Int -> Mode -> ShowS
[Mode] -> ShowS
Mode -> String
(Int -> Mode -> ShowS)
-> (Mode -> String) -> ([Mode] -> ShowS) -> Show Mode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Mode] -> ShowS
$cshowList :: [Mode] -> ShowS
show :: Mode -> String
$cshow :: Mode -> String
showsPrec :: Int -> Mode -> ShowS
$cshowsPrec :: Int -> Mode -> ShowS
Show)

instance Shader.Specialization Mode where
  specializationData :: Mode -> [Word32]
specializationData = \case
    Mode
UV ->
      [Word32
0]
    Mode
Texture ->
      [Word32
1]
    Shadow Word32
bits ->
      [Word32
2, Word32
bits]

allocate
  :: ( HasVulkan env
     , HasRenderPass renderpass
     )
  => Mode
  -> Vk.SampleCountFlagBits
  -> Tagged Scene DsBindings
  -> renderpass
  -> ResourceT (RIO env) Pipeline
allocate :: Mode
-> SampleCountFlagBits
-> Tagged Scene DsBindings
-> renderpass
-> ResourceT (RIO env) Pipeline
allocate Mode
mode SampleCountFlagBits
multisample (Tagged DsBindings
set0) = do
  ((ReleaseKey, Pipeline) -> Pipeline)
-> ResourceT (RIO env) (ReleaseKey, Pipeline)
-> ResourceT (RIO env) Pipeline
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ReleaseKey, Pipeline) -> Pipeline
forall a b. (a, b) -> b
snd (ResourceT (RIO env) (ReleaseKey, Pipeline)
 -> ResourceT (RIO env) Pipeline)
-> (renderpass -> ResourceT (RIO env) (ReleaseKey, Pipeline))
-> renderpass
-> ResourceT (RIO env) Pipeline
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Extent2D
-> SampleCountFlagBits
-> Config '[Scene] VertexAttrs InstanceAttrs Mode
-> renderpass
-> ResourceT (RIO env) (ReleaseKey, Pipeline)
forall env (m :: * -> *) renderpass spec (dsl :: [*]) vertices
       instances.
(MonadVulkan env m, MonadResource m, HasRenderPass renderpass,
 Specialization spec, HasCallStack) =>
Maybe Extent2D
-> SampleCountFlagBits
-> Config dsl vertices instances spec
-> renderpass
-> m (ReleaseKey, Pipeline dsl vertices instances)
Pipeline.allocate Maybe Extent2D
forall a. Maybe a
Nothing SampleCountFlagBits
multisample Config '[Scene] VertexAttrs InstanceAttrs Mode
config
  where
    config :: Config '[Scene] VertexAttrs InstanceAttrs Mode
config = Config '[] Any Any ()
forall vertices instances. Config '[] vertices instances ()
Pipeline.baseConfig
      { $sel:cDescLayouts:Config :: Tagged '[Scene] [DsBindings]
Pipeline.cDescLayouts    = [DsBindings] -> Tagged '[Scene] [DsBindings]
forall k (s :: k) b. b -> Tagged s b
Tagged @'[Scene] [DsBindings
set0]
      , $sel:cVertexCode:Config :: Maybe ByteString
Pipeline.cVertexCode     = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
vertCode
      , $sel:cVertexInput:Config :: SomeStruct PipelineVertexInputStateCreateInfo
Pipeline.cVertexInput    = SomeStruct PipelineVertexInputStateCreateInfo
vertexInput
      , $sel:cFragmentCode:Config :: Maybe ByteString
Pipeline.cFragmentCode   = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
fragCode
      , $sel:cSpecialization:Config :: Mode
Pipeline.cSpecialization = Mode
mode
      }

    vertexInput :: SomeStruct PipelineVertexInputStateCreateInfo
vertexInput = [(VertexInputRate, [Format])]
-> SomeStruct PipelineVertexInputStateCreateInfo
Pipeline.vertexInput
      [ (VertexInputRate, [Format])
vertexPos -- vPosition
      , (VertexInputRate
Vk.VERTEX_INPUT_RATE_VERTEX,   [Format]
Model.vkVertexAttrs)
      , (VertexInputRate
Vk.VERTEX_INPUT_RATE_INSTANCE, [Format]
Model.vkInstanceTexture)
      , (VertexInputRate, [Format])
instanceTransform
      ]

vertCode :: ByteString
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() {
      vec4 fPosition = iModel * vec4(vPosition, 1.0);

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

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

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

    ${set0binding1}
    ${set0binding2}
    ${set0binding3}
    ${set0binding5color}

    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;

    layout (constant_id=0) const uint mode = 0;
    layout (constant_id=1) const uint shadow_mask = 0;

    void main() {
      vec4 texel = vec4(0);

      switch(mode) {
        case 0:
          texel = vec4(fTexCoord, 0, 1);
          break;

        case 1:
          if ((fTextureIds.t < 0) || (fTextureIds.s < 0)) break;

          texel = texture(
            sampler2D(
              textures[nonuniformEXT(fTextureIds.t)],
              samplers[nonuniformEXT(fTextureIds.s)]
            ),
            fTexCoord
          );
          break;

        case 2:
          float d0 = 1.0;
          if ((shadow_mask & 1) == 1)
            d0 = texture(
              shadowmaps,
              vec3(fTexCoord, 0.0)
            ).x;

          float d1 = 1.0;
          if ((shadow_mask & 2) == 2)
            d1 = texture(
              shadowmaps,
              vec3(fTexCoord, 1.0)
            ).x;

          float d2 = 1.0;
          if ((shadow_mask & 4) == 4)
            d2 = texture(
              shadowmaps,
              vec3(fTexCoord, 2.0)
            ).x;

          texel = vec4(1-d0, 1-d1, 1-d2, 1.0);
          break;

        default:
          break;
      }

      vec3 color = pow(texel.rgb, fTextureGamma.rgb);
      float combinedAlpha = texel.a * fTextureGamma.a;

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