module Render.Unlit.TileMap.Pipeline
  ( Pipeline
  , allocate
  , allocateBlend

  , Config
  , config
  , configBlend
  ) 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 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.TileMap.Model qualified as Model

type Config = Pipeline.Configure Pipeline ()
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) = Pipeline.baseConfig
  { 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)
      , (Vk.VERTEX_INPUT_RATE_INSTANCE, Model.vkInstanceTileMap)
      , 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;
    // tilemapParams
    layout(location = 2) in ivec4 iTextureIds; // combined: tileset, tileset sampler, map, repeat
    layout(location = 3) in vec2  iViewOffset;
    layout(location = 4) in vec2  iViewportSize;
    layout(location = 5) in vec2  iMapTextureSize;
    layout(location = 6) in vec2  iTilesetTextureSize;
    layout(location = 7) in vec2  iTileSize;

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

    layout(location = 0)      out  vec2 fTexCoord;
    layout(location = 1)      out  vec2 fPixCoord;
    layout(location = 2) flat out ivec4 fTextureIds;
    layout(location = 3) flat out  vec2 fTilesetTextureSize;
    layout(location = 4) flat out  vec2 fTileSize;

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

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

      fPixCoord = (vTexCoord * iViewportSize) + iViewOffset;
      fTexCoord = fPixCoord / iMapTextureSize / iTileSize;

      fTextureIds = iTextureIds;
      fTilesetTextureSize = iTilesetTextureSize;
      fTileSize = iTileSize;
    }
  |])

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) in vec2 fPixCoord;

    // combined: tileset, tileset sampler, map, repeat
    layout(location = 2) flat in ivec4 fTextureIds;
    layout(location = 3) flat in vec2 fTilesetTextureSize;
    layout(location = 4) flat in vec2 fTileSize;

    layout(location = 0) out vec4 oColor;

    int tilesetTextureIx = fTextureIds[0];
    int tilesetSamplerIx = fTextureIds[1];
    int mapTextureIx     = fTextureIds[2];
    int repeatTiles      = fTextureIds[3];

    // TODO
    // const vec4 fTextureGamma = vec4(1.0);

    void main() {
      if (repeatTiles == 0 && (fTexCoord.x < 0.0 || fTexCoord.y < 0.0 || fTexCoord.x > 1.0 || fTexCoord.y > 1.0)) {
        discard;
      }

      vec4 map = texture(
        sampler2D(
          textures[nonuniformEXT(mapTextureIx)],
          samplers[7] // nearest
        ),
        fTexCoord
      );

      vec2 spriteOffset = floor(map.xy * 256.0) * fTileSize;
      vec2 spriteCoord = mod(fPixCoord, fTileSize);
      vec2 spriteUV = (spriteOffset + spriteCoord) / fTilesetTextureSize;

      oColor = texture(
        sampler2D(
          textures[nonuniformEXT(tilesetTextureIx)],
          samplers[nonuniformEXT(tilesetSamplerIx)]
        ),
        spriteUV
      );
    }
  |])