{-# LANGUAGE OverloadedLists #-}

module Render.ShadowMap.RenderPass
  ( ShadowMap(..)
  , allocate
  , usePass
  ) where

import RIO

import Control.Monad.Trans.Resource qualified as Resource
import Data.Bits ((.|.))
import Data.Vector qualified as Vector
import Vulkan.Core10 qualified as Vk
import Vulkan.Core11.Promoted_From_VK_KHR_multiview qualified as Khr
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)
import Vulkan.CStruct.Extends (pattern (:&), pattern (::&))

import Engine.Types.RefCounted (RefCounted, resourceTRefCount)
import Engine.Types.RefCounted qualified as RefCounted
import Engine.Vulkan.Types (HasVulkan(..), HasSwapchain(..), HasRenderPass(..), RenderPass(..), MonadVulkan)
import Render.Pass (usePass)
import Resource.Image (AllocatedImage)
import Resource.Image qualified as Image
import Resource.Region qualified as Region
import Resource.Vulkan.Named qualified as Named

-- * Depth-only pass for shadowmapping pipelines

data ShadowMap = ShadowMap
  { smRenderPass  :: Vk.RenderPass
  , smDepthImage  :: AllocatedImage
  , smFrameBuffer :: Vk.Framebuffer
  , smRenderArea  :: Vk.Rect2D
  , smExtent      :: Vk.Extent2D
  , smLayerCount  :: Word32
  , smClear       :: Vector Vk.ClearValue
  , smRelease     :: RefCounted
  }

instance HasRenderPass ShadowMap where
  getRenderPass   = smRenderPass
  getFramebuffers = Vector.replicate 10 . smFrameBuffer
  getClearValues  = smClear
  getRenderArea   = smRenderArea

instance RenderPass ShadowMap where
  refcountRenderpass  = resourceTRefCount . smRelease

allocate
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain context
     )
  => context
  -> Word32
  -> "light count" ::: Word32
  -> m ShadowMap
allocate context mapSize layerCount = do
  (_rpKey, renderPass) <- allocateRenderPass context (2 ^ layerCount - 1) 0
  (release, resources) <- RefCounted.wrapped $ Region.run do
    Region.logDebug
      "Allocating ShadowMap resources"
      "Releasing ShadowMap resources"
    allocateFramebuffer context extent layerCount renderPass
  let (depthImage, framebuffer) = resources

  pure ShadowMap
    { smRenderPass  = renderPass
    , smExtent      = extent
    , smLayerCount  = layerCount
    , smRenderArea  = fullSurface
    , smClear       = clear
    , smDepthImage  = depthImage
    , smFrameBuffer = framebuffer
    , smRelease     = release
    }
  where
    extent = Vk.Extent2D{width=mapSize, height=mapSize}

    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = extent
      }

    clear = Vector.fromList
      [ Vk.DepthStencil (Vk.ClearDepthStencilValue 1.0 0)
      ]

-- ** Render pass

allocateRenderPass
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Word32
  -> Word32
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPass swapchain viewMask correlationMask = do
  device <- asks getDevice

  res@(_key, object) <- Vk.withRenderPass device createInfo Nothing Resource.allocate
  Debug.nameObject device object "ShadowMap.RP"
  pure res
  where
    depthFormat = getDepthFormat swapchain

    createInfo =
      zero
        { Vk.attachments  = Vector.fromList [depth]
        , Vk.subpasses    = Vector.fromList [subpass]
        , Vk.dependencies = Vector.fromList [pre, post]
        }
      ::& Khr.RenderPassMultiviewCreateInfo
        { Khr.viewMasks        = [viewMask]
        , Khr.viewOffsets      = []
        , Khr.correlationMasks = [correlationMask]
        }
      :& ()

    depth = zero
      { Vk.format         = depthFormat
      , Vk.samples        = Vk.SAMPLE_COUNT_1_BIT
      , Vk.loadOp         = Vk.ATTACHMENT_LOAD_OP_CLEAR
      , Vk.storeOp        = Vk.ATTACHMENT_STORE_OP_STORE
      , Vk.stencilLoadOp  = Vk.ATTACHMENT_LOAD_OP_DONT_CARE
      , Vk.stencilStoreOp = Vk.ATTACHMENT_STORE_OP_DONT_CARE
      , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.finalLayout    = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL
      }

    subpass = zero
      { Vk.pipelineBindPoint = Vk.PIPELINE_BIND_POINT_GRAPHICS
      , Vk.depthStencilAttachment = Just zero
          { Vk.attachment = 0
          , Vk.layout     = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
          }
      }

    pre = zero
      { Vk.dependencyFlags = Vk.DEPENDENCY_BY_REGION_BIT

      , Vk.srcSubpass      = Vk.SUBPASS_EXTERNAL
      , Vk.srcStageMask    = Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
      , Vk.srcAccessMask   = Vk.ACCESS_SHADER_READ_BIT

      , Vk.dstSubpass      = 0
      , Vk.dstStageMask    = Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT
      , Vk.dstAccessMask   = Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
      }

    post = zero
      { Vk.dependencyFlags = Vk.DEPENDENCY_BY_REGION_BIT

      , Vk.srcSubpass      = 0
      , Vk.srcStageMask    = Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
      , Vk.srcAccessMask   = Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT

      , Vk.dstSubpass      = Vk.SUBPASS_EXTERNAL
      , Vk.dstStageMask    = Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
      , Vk.dstAccessMask   = Vk.ACCESS_SHADER_READ_BIT
      }

-- ** Framebuffer

type Framebuffers =
  ( Image.AllocatedImage
  , Vk.Framebuffer
  )

allocateFramebuffer
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Vk.Extent2D
  -> Word32
  -> Vk.RenderPass
  -> m Framebuffers
allocateFramebuffer swapchain extent layerCount renderPass = do
  depth <- Image.allocate
    (Just "ShadowMap.depth")
    Vk.IMAGE_ASPECT_DEPTH_BIT
    (Image.inflateExtent extent 1)
    1
    layerCount
    Vk.SAMPLE_COUNT_1_BIT
    (getDepthFormat swapchain)
    (Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT .|. Vk.IMAGE_USAGE_SAMPLED_BIT)

  let
    attachments = Vector.fromList
      [ Image.aiImageView depth
      ]

    {- XXX:
      If the render pass uses multiview, then layers must be one and each attachment
      requires a number of layers that is greater than the maximum bit index set in
      the view mask in the subpasses in which it is used.
    -}
    fbNumLayers = 1

    Vk.Extent2D{width, height} = extent
    fbCI = zero
      { Vk.renderPass  = renderPass
      , Vk.width       = width
      , Vk.height      = height
      , Vk.attachments = attachments
      , Vk.layers      = fbNumLayers
      }

  device <- asks getDevice
  (_framebufferKey, framebuffer) <- Vk.withFramebuffer device fbCI Nothing Resource.allocate
  Named.object framebuffer "ShadowMap.FB"

  pure (depth, framebuffer)