{-# LANGUAGE OverloadedLists #-}

module Render.ShadowMap.RenderPass where

import RIO

import Control.Monad.Trans.Resource qualified as Resource
import Data.Vector qualified as Vector
import RIO.Vector.Partial ((!))
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, newRefCounted, resourceTRefCount)
import Engine.Vulkan.Types (HasVulkan(..), HasSwapchain(..), HasRenderPass(..), RenderPass(..), MonadVulkan)
import Resource.Image (AllocatedImage)
import Resource.Image qualified as Image

-- * 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
  allocateRenderpass_ = error "Don't allocate ShadowMap via class"
  updateRenderpass    = error "Don't recreate ShadowMap on resize"
  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
  logDebug "Allocating ShadowMap resources"
  (_rpKey, renderPass) <- allocateRenderPass context (2 ^ layerCount - 1) 0
  (refcounted, depthImage, framebuffer) <- allocateFramebuffer context extent layerCount renderPass

  pure ShadowMap
    { smRenderPass  = renderPass
    , smExtent      = extent
    , smLayerCount  = layerCount
    , smRenderArea  = fullSurface
    , smClear       = clear
    , smDepthImage  = depthImage
    , smFrameBuffer = framebuffer
    , smRelease     = refcounted
    }
  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 context
     )
  => context
  -> Word32
  -> Word32
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPass context viewMask correlationMask = do
  device <- asks getDevice
  let depthFormat = getDepthFormat context

  res@(_key, object) <- Vk.withRenderPass device (createInfo depthFormat) Nothing Resource.allocate
  Debug.nameObject device object "ShadowMap.RP"
  pure res
  where
    createInfo depthFormat =
      zero
        { Vk.attachments  = Vector.fromList [depth depthFormat]
        , Vk.subpasses    = Vector.fromList [subpass]
        , Vk.dependencies = Vector.fromList [pre, post]
        }
      ::& Khr.RenderPassMultiviewCreateInfo
        { Khr.viewMasks        = [viewMask]
        , Khr.viewOffsets      = []
        , Khr.correlationMasks = [correlationMask]
        }
      :& ()

    depth depthFormat = 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 =
  ( RefCounted
  , Image.AllocatedImage
  , Vk.Framebuffer
  )

allocateFramebuffer
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain context
     )
  => context
  -> Vk.Extent2D
  -> Word32
  -> Vk.RenderPass
  -> m Framebuffers
allocateFramebuffer context extent layerCount renderPass = do
  device <- asks getDevice
  context' <- asks (, context)
  let Vk.Extent2D{width, height} = extent

  (depthKey, depth) <- Resource.allocate
    (Image.createDepthResource context' extent $ Just layerCount)
    (Image.destroyAllocatedImage context')

  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

    fbCI = zero
      { Vk.renderPass  = renderPass
      , Vk.width       = width
      , Vk.height      = height
      , Vk.attachments = attachments
      , Vk.layers      = fbNumLayers
      }
  (framebufferKey, framebuffer) <- Vk.withFramebuffer device fbCI Nothing Resource.allocate
  Debug.nameObject device framebuffer "ShadowMap.FB"

  releaseDebug <- toIO $ logDebug "Releasing ShadowMap resources"
  release <- newRefCounted do
    releaseDebug
    Resource.release depthKey
    Resource.release framebufferKey

  pure (release, depth, framebuffer)

usePass :: (MonadIO io, HasRenderPass a) => a -> Word32 -> Vk.CommandBuffer -> io r -> io r
usePass render imageIndex cb =
  Vk.cmdUseRenderPass
    cb
    (beginInfo render imageIndex)
    Vk.SUBPASS_CONTENTS_INLINE

beginInfo :: HasRenderPass a => a -> Word32 -> Vk.RenderPassBeginInfo '[]
beginInfo rp imageIndex = zero
  { Vk.renderPass  = getRenderPass rp
  , Vk.framebuffer = getFramebuffers rp ! fromIntegral imageIndex -- XXX: we only have one
  , Vk.renderArea  = getRenderArea rp
  , Vk.clearValues = getClearValues rp
  }