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

module Render.DescSets.Sun
  ( Sun(..)
  , createSet0Ds
  , set0

  , pattern MAX_VIEWS
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
import Control.Monad.Trans.Resource qualified as ResourceT
import Data.Tagged (Tagged(..))
import Data.Vector qualified as Vector
import Data.Vector.Storable qualified as VectorS
import Foreign.Storable.Generic (GStorable)
import Geomancy (Vec4, vec4)
import Geomancy.Transform (Transform)
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (Zero(..))

import Engine.Types (StageRIO)
import Engine.Vulkan.DescSets ()
import Engine.Vulkan.Types (DsBindings, HasVulkan(..))
import Resource.Buffer qualified as Buffer
import Resource.DescriptorSet qualified as DescriptorSet

-- * Set0 data

-- | Maximum "guaranteed" amount for multiview passes
pattern MAX_VIEWS :: Int
pattern MAX_VIEWS = 6

data Sun = Sun
  { sunViewProjection :: Transform
  , sunShadow         :: Vec4 -- offsetx, offsety, index, size -- XXX: only index is used
  , sunPosition       :: Vec4 -- XXX: alpha available for stuff
  , sunDirection      :: Vec4 -- XXX: alpha available for stuff
  , sunColor          :: Vec4 -- XXX: RGB premultiplied, alpha is available for stuff
  }
  deriving (Show, Generic)

instance GStorable Sun

instance Zero Sun where
  zero = Sun
    { sunViewProjection = mempty
    , sunShadow         = 0
    , sunPosition       = 0
    , sunDirection      = vec4 0 1 0 0
    , sunColor          = 0
    }

-- * Shadow casting descriptor set

set0
  :: Tagged Sun DsBindings
set0 = Tagged
  [ (set0bind0, zero)
  ]

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

-- * Setup

createSet0Ds
  :: Tagged '[Sun] Vk.DescriptorSetLayout
  -> ResourceT (StageRIO st)
      ( Tagged '[Sun] (Vector Vk.DescriptorSet)
      , Buffer.Allocated 'Buffer.Coherent Sun
      )
createSet0Ds (Tagged set0layout) = do
  context <- asks id

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

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

  let
    initialSuns = VectorS.replicate MAX_VIEWS zero
  (_, sunData) <- ResourceT.allocate
    (Buffer.createCoherent context Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT MAX_VIEWS initialSuns)
    (Buffer.destroyAll context . Just)

  updateSet0Ds descSets sunData

  let device = getDevice context
  Debug.nameObject device descPool "Sun.Pool"
  for_ (unTagged descSets) \ds ->
    Debug.nameObject device ds "Sun.DS"
  Debug.nameObject device (Buffer.aBuffer sunData) "Sun.Data"

  pure (descSets, sunData)

dpSizes :: DescriptorSet.TypeMap Word32
dpSizes =
  [ ( Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
    , uniformBuffers
    )
  -- XXX: may be required to fetch textures for shadows from texture-masked models
  -- , ( Vk.DESCRIPTOR_TYPE_SAMPLED_IMAGE
  --   , sampledImages
  --   )
  -- , ( Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
  --   , sampledImages + shadowMaps
  --   )
  -- , ( Vk.DESCRIPTOR_TYPE_SAMPLER
  --   , staticSamplers
  --   )
  ]
  where
    uniformBuffers = 2    -- 1 scene + 1 light array
    -- sampledImages  = 128  -- max dynamic textures and cubemaps
    -- staticSamplers = 8    -- immutable samplers
    -- shadowMaps     = 2    -- max shadowmaps

updateSet0Ds
  :: Tagged '[Sun] (Vector Vk.DescriptorSet)
  -> Buffer.Allocated 'Buffer.Coherent Sun
  -> ResourceT (StageRIO st) ()
updateSet0Ds (Tagged ds) Buffer.Allocated{aBuffer} = do
  context <- asks id
  Vk.updateDescriptorSets (getDevice context) writeSets mempty

  where
    destSet0 = case Vector.headM ds of
      Nothing ->
        error "assert: descriptor sets promised to contain [Sun]"
      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 = aBuffer
          , Vk.offset = 0
          , Vk.range  = Vk.WHOLE_SIZE
          }

    writeSets =
      Vector.singleton writeSet0b0