{-# OPTIONS_GHC -fplugin Foreign.Storable.Generic.Plugin #-}

module Render.DescSets.Set0
  ( Scene(..)
  , emptyScene

  , allocate
  , allocateEmpty

  , updateSet0Ds

  , mkBindings

  -- TODO: extract to typeclass magic
  , vertexPos
  , instanceTransform

  , FrameResource(..)
  , extendResourceDS

  , Buffer
  , Process
  , observe

  , withBoundSet0
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Control.Monad.Trans.Resource qualified as ResourceT
import Data.Bits ((.|.))
import Data.ByteString.Char8 qualified as BS8
import Data.Kind (Type)
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import Data.Vector.Storable qualified as VectorS
import Foreign.Storable.Generic (GStorable)
import Geomancy (Vec3, Vec4, vec3)
import Geomancy.Transform (Transform)
import Vulkan.Core10 qualified as Vk
import Vulkan.Core12.Promoted_From_VK_EXT_descriptor_indexing qualified as Vk12
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (Zero(..))

import Engine.Types (StageRIO)
import Engine.Vulkan.DescSets (Bound, Extend, extendDS, withBoundDescriptorSets0)
import Engine.Vulkan.Pipeline (Pipeline)
import Engine.Vulkan.Pipeline qualified as Pipeline
import Engine.Vulkan.Types (DsBindings, HasVulkan(..))
import Engine.Worker qualified as Worker
import Global.Resource.CubeMap.Base qualified as BaseCubeMap
import Global.Resource.Texture.Base qualified as BaseTexture
import Render.DescSets.Sun (Sun)
import Render.Lit.Material (Material)
import Resource.Buffer qualified as Buffer
import Resource.Collection qualified as Collection
import Resource.DescriptorSet qualified as DescriptorSet
import Resource.Image qualified as Image
import Resource.Texture qualified as Texture

-- * Set0 data

data Scene = Scene
  { sceneProjection    :: Transform
  , sceneInvProjection :: Transform

  , sceneView          :: Transform
  , sceneInvView       :: Transform
  , sceneViewPos       :: Vec3 -- XXX: gets extra padding
  , sceneViewDir       :: Vec3 -- XXX: gets extra padding

  , sceneTweaks        :: Vec4 -- ^ 4 debug tweaks bound to Kontrol

  , sceneFog           :: Vec4 -- XXX: RGB color + scatter factor β
  , sceneEnvCube       :: Int32
  , sceneNumLights     :: Word32
  }
  deriving (Show, Generic)

instance GStorable Scene

emptyScene :: Scene
emptyScene = Scene
  { sceneProjection    = mempty
  , sceneInvProjection = mempty

  , sceneView          = mempty
  , sceneInvView       = mempty
  , sceneViewPos       = vec3 0 0 0
  , sceneViewDir       = vec3 0 0 1

  , sceneFog           = 0
  , sceneEnvCube       = minBound
  , sceneNumLights     = 0

  , sceneTweaks        = 0
  }

-- * Common descriptor set

mkBindings
  :: ( Foldable samplers
     , Foldable textures
     , Foldable cubemaps
     )
  => samplers Vk.Sampler
  -> textures a
  -> cubemaps b
  -> Word32
  -> Tagged Scene DsBindings
mkBindings samplers textures cubes shadows = Tagged
  [ (set0bind0,          zero)
  , (set0bind1 samplers, zero)
  , (set0bind2 textures, partialBinding)
  , (set0bind3 cubes,    partialBinding)
  , (set0bind4,          zero)
  , (set0bind5 shadows,  partialBinding)
  , (set0bind6,          zero)
  ]
  where
    partialBinding =
      Vk12.DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT

    _partialVariable =
      partialBinding .|.
      Vk12.DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT
      {-
        VUID-VkDescriptorSetLayoutBindingFlagsCreateInfo-pBindingFlags-03004(ERROR / SPEC):
          msgNum: 222246202 - Validation Error:
          [ VUID-VkDescriptorSetLayoutBindingFlagsCreateInfo-pBindingFlags-03004 ]
          Object 0: handle = 0x157ea80, type = VK_OBJECT_TYPE_DEVICE;
          | MessageID = 0xd3f353a
          | Invalid flags for VkDescriptorSetLayoutBinding entry 2 The Vulkan spec states:
              If an element of pBindingFlags includes VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT,
              then all other elements of VkDescriptorSetLayoutCreateInfo::pBindings must have
              a smaller value of binding
              (https://vulkan.lunarg.com/doc/view/1.2.162.1~rc2/linux/1.2-extensions/vkspec.html#VUID-VkDescriptorSetLayoutBindingFlagsCreateInfo-pBindingFlags-03004)
      -}

set0bind0 :: Vk.DescriptorSetLayoutBinding
set0bind0 = Vk.DescriptorSetLayoutBinding
  { binding           = 0
  , descriptorType    = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , descriptorCount   = 1
  , stageFlags        = Vk.SHADER_STAGE_ALL
  , immutableSamplers = mempty
  }

set0bind1
  :: Foldable samplers
  => samplers Vk.Sampler
  -> Vk.DescriptorSetLayoutBinding
set0bind1 samplers = Vk.DescriptorSetLayoutBinding
  { Vk.binding           = 1
  , Vk.stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , Vk.descriptorType    = Vk.DESCRIPTOR_TYPE_SAMPLER
  , Vk.descriptorCount   = fromIntegral $ Vector.length linearSamplers
  , Vk.immutableSamplers = linearSamplers
  }
  where
    linearSamplers = Collection.toVector samplers

set0bind2
  :: Foldable textures
  => textures a
  -> Vk.DescriptorSetLayoutBinding
set0bind2 textures = Vk.DescriptorSetLayoutBinding
  { binding           = 2
  , descriptorType    = Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
  , descriptorCount   = fromIntegral $ max baseTextures textureCount
  , stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , immutableSamplers = mempty
  }
  where
    baseTextures = length BaseTexture.sources
    textureCount = length textures

set0bind3
  :: Foldable cubes
  => cubes a
  -> Vk.DescriptorSetLayoutBinding
set0bind3 cubes = Vk.DescriptorSetLayoutBinding
  { binding           = 3
  , descriptorType    = Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
  , descriptorCount   = fromIntegral $ max baseCubes cubeCount
  , stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , immutableSamplers = mempty
  }
  where
    baseCubes = length BaseCubeMap.sources
    cubeCount = length cubes

set0bind4 :: Vk.DescriptorSetLayoutBinding
set0bind4 = Vk.DescriptorSetLayoutBinding
  { binding           = 4
  , descriptorType    = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , descriptorCount   = 1
  , immutableSamplers = mempty
  }

set0bind5 :: Word32 -> Vk.DescriptorSetLayoutBinding
set0bind5 shadows = Vk.DescriptorSetLayoutBinding
  { binding           = 5
  , descriptorType    = Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
  , stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , descriptorCount   = max 1 shadows
  , immutableSamplers = mempty
  }

set0bind6 :: Vk.DescriptorSetLayoutBinding
set0bind6 = Vk.DescriptorSetLayoutBinding
  { binding           = 6
  , descriptorType    = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
  , stageFlags        = Vk.SHADER_STAGE_FRAGMENT_BIT
  , descriptorCount   = 1
  , immutableSamplers = mempty
  }

vertexPos :: (Vk.VertexInputRate, [Vk.Format])
vertexPos =
  ( Vk.VERTEX_INPUT_RATE_VERTEX
  , [ Vk.FORMAT_R32G32B32_SFLOAT -- vPosition :: vec3
    ]
  )

instanceTransform :: (Vk.VertexInputRate, [Vk.Format])
instanceTransform =
  ( Vk.VERTEX_INPUT_RATE_INSTANCE
  , [ Vk.FORMAT_R32G32B32A32_SFLOAT -- iModel :: mat4
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    , Vk.FORMAT_R32G32B32A32_SFLOAT
    ]
  )

-- * Setup

allocate
  :: (Traversable textures, Traversable cubes)
  => Tagged '[Scene] Vk.DescriptorSetLayout
  -> textures (Texture.Texture Texture.Flat)
  -> cubes (Texture.Texture Texture.CubeMap)
  -> Maybe (Buffer.Allocated 'Buffer.Coherent Sun)
  -> "shadow maps" ::: Vector Vk.ImageView
  -> Maybe (Buffer.Allocated 'Buffer.Coherent Material)
  -> ResourceT (StageRIO st) (FrameResource '[Scene])
allocate (Tagged set0layout) textures cubes lightsData shadowViews materialsData = do
  context <- asks id

  (_dpKey, descPool) <- DescriptorSet.allocatePool 1 dpSizes

  let
    set0dsCI = zero
      { Vk.descriptorPool = descPool
      , Vk.setLayouts     = Vector.singleton set0layout
      }
  descSets <- fmap (Tagged @'[Scene]) $
    Vk.allocateDescriptorSets (getDevice context) set0dsCI

  (_, sceneData) <- ResourceT.allocate
    (Buffer.createCoherent context Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT 1 $ VectorS.singleton emptyScene)
    (Buffer.destroy context)

  let
    -- TODO: must be checked against depth format and TILING_OPTIMAL
    -- shadowFilter = Vk.FILTER_NEAREST
    shadowFilter = Vk.FILTER_LINEAR

    shadowCI = zero
      { Vk.addressModeU  = Vk.SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE
      , Vk.addressModeV  = Vk.SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE
      , Vk.addressModeW  = Vk.SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE
      , Vk.borderColor   = Vk.BORDER_COLOR_FLOAT_OPAQUE_WHITE
      , Vk.magFilter     = shadowFilter
      , Vk.minFilter     = shadowFilter
      , Vk.compareEnable = True
      , Vk.compareOp     = Vk.COMPARE_OP_LESS
      }

  let ifor = flip Vector.imapM
  shadowMaps <- ifor shadowViews \ix depthView -> do
    (_, shadowSampler) <- Vk.withSampler (getDevice context) shadowCI Nothing ResourceT.allocate
    Debug.nameObject (getDevice context) shadowSampler $
      "ShadowSampler." <> BS8.pack (show ix)
    pure (shadowSampler, depthView)

  updateSet0Ds context descSets sceneData textures cubes lightsData shadowMaps materialsData

  scene <- Worker.newObserverIO emptyScene

  pure $ FrameResource descSets sceneData scene

dpSizes :: DescriptorSet.TypeMap Word32
dpSizes =
  [ ( Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
    , uniformBuffers
    )
  , ( Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
    , sampledImages
    )
  , ( Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
    , sampledImages + shadowMaps
    )
  , ( Vk.DESCRIPTOR_TYPE_SAMPLER
    , staticSamplers
    )
  ]
  where
    uniformBuffers = 3    -- 1 scene + 1 light array + 1 material array
    sampledImages  = 128  -- max dynamic textures and cubemaps
    staticSamplers = 8    -- immutable samplers
    shadowMaps     = 2    -- max shadowmaps

-- | Minimal viable 'Scene' without textures and lighting.
allocateEmpty
  :: Tagged '[Scene] Vk.DescriptorSetLayout
  -> ResourceT (StageRIO st) (FrameResource '[Scene])
allocateEmpty taggedLayout = allocate taggedLayout [] [] Nothing mempty Nothing

updateSet0Ds
  :: (HasVulkan context, Traversable textures, Traversable cubes)
  => context
  -> Tagged '[Scene] (Vector Vk.DescriptorSet)
  -> Buffer.Allocated 'Buffer.Coherent Scene
  -> textures (Texture.Texture Texture.Flat)
  -> cubes (Texture.Texture Texture.CubeMap)
  -> Maybe (Buffer.Allocated 'Buffer.Coherent Sun)
  -> Vector (Vk.Sampler, Vk.ImageView)
  -> Maybe (Buffer.Allocated 'Buffer.Coherent Material)
  -> ResourceT (StageRIO st) ()
updateSet0Ds context (Tagged ds) sceneData textures cubes lightsData shadowMaps materialsData =
  Vk.updateDescriptorSets (getDevice context) writeSets mempty

  where
    linearTextures = Collection.toVector textures
    linearCubes = Collection.toVector cubes

    destSet0 = case Vector.headM ds of
      Nothing ->
        error "assert: descriptor sets promised to contain [Scene]"
      Just one ->
        one

    writeSet0b0 = SomeStruct zero
      { Vk.dstSet          = destSet0
      , Vk.dstBinding      = 0
      , Vk.dstArrayElement = 0
      , Vk.descriptorCount = 1
      , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
      , Vk.bufferInfo      = Vector.singleton set0bind0I
      }
      where
        set0bind0I = Vk.DescriptorBufferInfo
          { Vk.buffer = Buffer.aBuffer sceneData
          , Vk.offset = 0
          , Vk.range  = Vk.WHOLE_SIZE
          }

    writeSet0b2 = SomeStruct zero
      { Vk.dstSet          = destSet0
      , Vk.dstBinding      = 2
      , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
      , Vk.dstArrayElement = 0
      , Vk.descriptorCount = fromIntegral $ Vector.length linearTextures
      , Vk.imageInfo       = textureInfos
      }
      where
        textureInfos = do
          texture <- linearTextures
          pure Vk.DescriptorImageInfo
            { sampler     = zero
            , imageView   = Image.aiImageView $ Texture.tAllocatedImage texture
            , imageLayout = Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            }

    writeSet0b3 = SomeStruct zero
      { Vk.dstSet          = destSet0
      , Vk.dstBinding      = 3
      , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
      , Vk.dstArrayElement = 0
      , Vk.descriptorCount = fromIntegral $ Vector.length linearCubes
      , Vk.imageInfo       = cubeInfos
      }
      where
        cubeInfos = do
          cube <- linearCubes
          pure Vk.DescriptorImageInfo
            { sampler     = zero
            , imageView   = Image.aiImageView $ Texture.tAllocatedImage cube
            , imageLayout = Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            }

    writeSet0b4M =
      case lightsData of
        Nothing ->
          mzero
        Just someLights ->
          pure $ SomeStruct zero
            { Vk.dstSet          = destSet0
            , Vk.dstBinding      = 4
            , Vk.dstArrayElement = 0
            , Vk.descriptorCount = 1
            , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            , Vk.bufferInfo      = Vector.singleton set0bind4I
            }
          where
            set0bind4I = Vk.DescriptorBufferInfo
              { Vk.buffer = Buffer.aBuffer someLights
              , Vk.offset = 0
              , Vk.range  = Vk.WHOLE_SIZE
              }

    writeSet0b5 = SomeStruct zero
      { Vk.dstSet          = destSet0
      , Vk.dstBinding      = 5
      , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
      , Vk.dstArrayElement = 0
      , Vk.descriptorCount = fromIntegral $ Vector.length shadowMaps
      , Vk.imageInfo       = shadowInfos
      }
      where
        shadowInfos = do
          (shadowSampler, shadowImageView) <- shadowMaps
          pure Vk.DescriptorImageInfo
            { sampler     = shadowSampler
            , imageView   = shadowImageView
            , imageLayout = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL
            }

    writeSet0b6M =
      case materialsData of
        Nothing ->
          mzero
        Just someMaterials ->
          pure $ SomeStruct zero
            { Vk.dstSet          = destSet0
            , Vk.dstBinding      = 6
            , Vk.dstArrayElement = 0
            , Vk.descriptorCount = 1
            , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            , Vk.bufferInfo      = Vector.singleton set0bind6I
            }
          where
            set0bind6I = Vk.DescriptorBufferInfo
              { Vk.buffer = Buffer.aBuffer someMaterials
              , Vk.offset = 0
              , Vk.range  = Vk.WHOLE_SIZE
              }

    writeSets = Vector.fromList $ concat
      [ pure writeSet0b0
      -- XXX: binding 1 is immutable samplers, baked into layout.
      , skipEmpty linearTextures writeSet0b2
      , skipEmpty linearCubes writeSet0b3
      , writeSet0b4M
      , skipEmpty shadowMaps writeSet0b5
      , writeSet0b6M
      ]
      where
        skipEmpty items writer
          | Vector.null items = mempty
          | otherwise = pure writer

extendResourceDS
  :: FrameResource ds
  -> Tagged ext Vk.DescriptorSet
  -> FrameResource (Extend ds ext)
extendResourceDS FrameResource{..} ext = FrameResource
  { frDescSets = extendDS frDescSets ext
  , ..
  }

-- * Frame data

data FrameResource (ds :: [Type]) = FrameResource
  { frDescSets :: Tagged ds (Vector Vk.DescriptorSet)
  , frBuffer   :: Buffer
  , frObserver :: Worker.ObserverIO Scene
  }

type Buffer = Buffer.Allocated 'Buffer.Coherent Scene

-- | A process that will assemble 'Scene' values.
type Process = Worker.Merge Scene

observe
  :: (MonadUnliftIO m) -- TODO: compatible '[Scene]
  => Process -> FrameResource ds -> m ()
observe process FrameResource{frBuffer, frObserver}=
  Worker.observeIO_ process frObserver \_old new -> do
    _same <- Buffer.updateCoherent (VectorS.singleton new) frBuffer
    pure new

-- * Rendering

withBoundSet0
  :: MonadIO m
  => FrameResource ds
  -> Pipeline ds vertices instances
  -> Vk.CommandBuffer
  -> Bound ds Void Void m b
  -> m b
withBoundSet0 FrameResource{frDescSets} refPipeline cb action =
  withBoundDescriptorSets0
    cb
    Vk.PIPELINE_BIND_POINT_GRAPHICS
    (Pipeline.pLayout refPipeline)
    frDescSets
    action