{-# LANGUAGE OverloadedLists #-}

module Render.Unlit.Sprite.Pipeline
  ( Pipeline
  , allocate
  -- , allocateOutline
  ) 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.Types (HasVulkan, HasRenderPass(..), DsBindings)
import Render.Code (compileVert, compileFrag, glsl)
import Render.DescSets.Set0 (Scene)
import Render.DescSets.Set0.Code (set0binding0, set0binding1, set0binding2)
import Render.Unlit.Sprite.Model qualified as Model

type Config = Pipeline.Configure Pipeline (Float, Bool)
type Pipeline = Pipeline.Pipeline '[Scene] () Model.InstanceAttrs

allocate
  :: ( HasVulkan env
     , HasRenderPass renderpass
     )
  => Vk.SampleCountFlagBits
  -> Maybe Float
  -> Bool
  -> Tagged Scene DsBindings
  -> renderpass
  -> ResourceT (RIO env) Pipeline
allocate :: forall env renderpass.
(HasVulkan env, HasRenderPass renderpass) =>
SampleCountFlagBits
-> Maybe Float
-> Bool
-> Tagged Scene DsBindings
-> renderpass
-> ResourceT (RIO env) Pipeline
allocate SampleCountFlagBits
multisample Maybe Float
discardAlpha Bool
outline Tagged Scene 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] () InstanceAttrs (Float, Bool)
-> 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
    (Maybe Float -> Bool -> Tagged Scene DsBindings -> Config
config Maybe Float
discardAlpha Bool
outline Tagged Scene DsBindings
set0)

config
  :: Maybe Float
  -> Bool
  -> Tagged Scene DsBindings
  -> Config
config :: Maybe Float -> Bool -> Tagged Scene DsBindings -> Config
config Maybe Float
discardAlpha Bool
outline (Tagged DsBindings
set0) =
  Config '[] Any Any ()
forall vertices instances. Config '[] vertices instances ()
Pipeline.baseConfig
    { $sel:cVertexCode:Config :: Maybe ByteString
Pipeline.cVertexCode     = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
vertCode
    , $sel:cFragmentCode:Config :: Maybe ByteString
Pipeline.cFragmentCode   = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
fragCode
    , $sel:cDescLayouts:Config :: Tagged '[Scene] [DsBindings]
Pipeline.cDescLayouts    = forall {s :: [*]} {b}. b -> Tagged s b
forall {k} (s :: k) b. b -> Tagged s b
Tagged @'[Scene] [DsBindings
Item [DsBindings]
set0]
    , $sel:cVertexInput:Config :: SomeStruct PipelineVertexInputStateCreateInfo
Pipeline.cVertexInput    = SomeStruct PipelineVertexInputStateCreateInfo
vertexInput
    , $sel:cDepthTest:Config :: Bool
Pipeline.cDepthTest      = Bool
False
    , $sel:cDepthWrite:Config :: Bool
Pipeline.cDepthWrite     = Bool
False
    , $sel:cBlend:Config :: Bool
Pipeline.cBlend          = Bool
True
    , $sel:cCull:Config :: CullModeFlagBits
Pipeline.cCull           = CullModeFlagBits
Vk.CULL_MODE_NONE
    , $sel:cSpecialization:Config :: (Float, Bool)
Pipeline.cSpecialization = (Float, Bool)
specs
    }
  where
    vertexInput :: SomeStruct PipelineVertexInputStateCreateInfo
vertexInput = [(VertexInputRate, [Format])]
-> SomeStruct PipelineVertexInputStateCreateInfo
Pipeline.vertexInput
      [ (VertexInputRate
Vk.VERTEX_INPUT_RATE_INSTANCE, [Format]
Model.vkInstanceAttrs)
      ]

    specs :: (Float, Bool)
specs =
      ( Float -> Maybe Float -> Float
forall a. a -> Maybe a -> a
fromMaybe Float
0.0 Maybe Float
discardAlpha
      , Bool
outline
      )

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

    ${set0binding0}

    layout(location = 0) in  vec4 iVert;
    layout(location = 1) in  vec4 iFrag;
    layout(location = 2) in  vec4 iTint;
    layout(location = 3) in  vec4 iOutline;
    layout(location = 4) in ivec2 iTextureIds;
    layout(location = 5) in uvec2 iTextureSize;

    layout(location = 0)      out  vec2 fTexCoord;
    layout(location = 1) flat out  vec4 fTint;
    layout(location = 2) flat out  vec4 fOutline;
    layout(location = 3) flat out ivec2 fTextureIds;
    layout(location = 4) flat out uvec2 fTextureSize;

    vec2 positions[6] = vec2[](
      vec2(-0.5, -0.5),
      vec2(0.5, -0.5),
      vec2(-0.5, 0.5),

      vec2(0.5, 0.5),
      vec2(-0.5, 0.5),
      vec2(0.5, -0.5)
    );

    vec2 texCoords[6] = vec2[](
      vec2(0.0, 0.0),
      vec2(1.0, 0.0),
      vec2(0.0, 1.0),

      vec2(1.0, 1.0),
      vec2(0.0, 1.0),
      vec2(1.0, 0.0)
    );

    void main() {
      vec2 pos = positions[gl_VertexIndex];

      gl_Position
        = scene.projection
        * scene.view
        * vec4(iVert.xy + pos * iVert.zw, 0, 1.0);

      vec2 uv = texCoords[gl_VertexIndex];
      fTexCoord = iFrag.xy + uv * iFrag.zw;

      fTint = iTint;
      fOutline = iOutline;
      fTextureIds = iTextureIds;
      fTextureSize = iTextureSize;
    }
  |])

fragCode :: ByteString
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 fTint;
    layout(location = 2) flat in  vec4 fOutline;
    layout(location = 3) flat in ivec2 fTextureIds;
    layout(location = 4) flat in uvec2 fTextureSize;

    layout(location = 0) out vec4 outColor;

    layout(constant_id=0) const float discardAlpha = 0.0;
    layout(constant_id=1) const bool  doOutline = false;

    vec4 tap(vec2 uv) {
      return texture(
        sampler2D(
          textures[nonuniformEXT(fTextureIds.t)],
          samplers[nonuniformEXT(fTextureIds.s)]
        ),
        fTexCoord + uv
      );
    }

    vec2 step_size = 1.0 / vec2(fTextureSize);

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

      if (discardAlpha != 0.0) {
        if (texel.a < discardAlpha) {
          discard;
        }
      }

      outColor = texel * fTint * texel.a;

      if (doOutline) {
        if (fOutline.a > 0.0 && texel.a < 1.0) {
          float alpha = 0.0;

          alpha += tap(vec2(step_size.x,  0.0)         ).a / 4.0;
          alpha += tap(vec2(-step_size.x, 0.0)         ).a / 4.0;
          alpha += tap(vec2(0.0,          step_size.y) ).a / 4.0;
          alpha += tap(vec2(0.0,          -step_size.y)).a / 4.0;

          outColor += fOutline * alpha;
        }
      }
    }
  |])