{-# LANGUAGE OverloadedLists #-}

module Render.Pass.Offscreen
  ( Settings(..)
  , allocate
  , Offscreen(..)
  , colorTexture
  , colorCube
  , depthTexture
  , depthCube
  ) 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.CStruct.Extends (pattern (:&), pattern (::&))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)

import Engine.Types.RefCounted (RefCounted, newRefCounted, resourceTRefCount)
import Engine.Vulkan.Types (HasVulkan(..), HasRenderPass(..), RenderPass(..), MonadVulkan)
import Resource.Image (AllocatedImage)
import Resource.Image qualified as Image
import Resource.Texture (CubeMap, Flat, Texture(..))

{- XXX: Consider spec wrt. parameters and intended use!

https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/vkCmdBlitImage.html
-}
data Settings = Settings
  { sLabel       :: Text
  , sExtent      :: Vk.Extent2D
  , sFormat      :: Vk.Format
  , sDepthFormat :: Vk.Format
  , sLayers      :: Word32
  , sMultiView   :: Bool -- ^ Makes sense only for multiple layers.
  , sSamples     :: Vk.SampleCountFlagBits -- ^ Multisample prevents mipmapping and cubes.
  , sMipMap      :: Bool
  }
  deriving (Eq, Show)

data Offscreen = Offscreen
  { oRenderPass  :: Vk.RenderPass
  , oExtent      :: Vk.Extent2D
  , oColor       :: AllocatedImage
  , oDepth       :: AllocatedImage
  , oLayers      :: Word32
  , oMipLevels   :: Word32

  , oFrameBuffer :: Vk.Framebuffer
  , oRenderArea  :: Vk.Rect2D
  , oClear       :: Vector Vk.ClearValue
  , oRelease     :: RefCounted
  }

instance HasRenderPass Offscreen where
  getRenderPass   = oRenderPass
  getFramebuffers = Vector.replicate 10 . oFrameBuffer
  getClearValues  = oClear
  getRenderArea   = oRenderArea

instance RenderPass Offscreen where
  refcountRenderpass = resourceTRefCount . oRelease

colorTexture :: Offscreen -> Texture Flat
colorTexture Offscreen{..} = Texture
  { tFormat         = Image.aiFormat oColor
  , tMipLevels      = oMipLevels
  , tLayers         = oLayers
  , tAllocatedImage = oColor
  }

colorCube :: Offscreen -> Texture CubeMap
colorCube Offscreen{..} = Texture
  { tFormat         = Image.aiFormat oColor
  , tMipLevels      = oMipLevels
  , tLayers         = oLayers -- TODO: get from type param
  , tAllocatedImage = oColor
  }

depthTexture :: Offscreen -> Texture Flat
depthTexture Offscreen{..} = Texture
  { tFormat         = Image.aiFormat oDepth
  , tMipLevels      = oMipLevels
  , tLayers         = oLayers
  , tAllocatedImage = oDepth
  }

depthCube :: Offscreen -> Texture CubeMap
depthCube Offscreen{..} = Texture
  { tFormat         = Image.aiFormat oDepth
  , tMipLevels      = oMipLevels
  , tLayers         = oLayers -- TODO: get from type param
  , tAllocatedImage = oDepth
  }

allocate
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     )
  => Settings
  -> m Offscreen
allocate settings@Settings{..} = do
  logDebug $ "Allocating Offscreen resources for " <> display sLabel

  (_rpKey, renderPass) <- allocateRenderPass settings

  (refcounted, color, depth, framebuffer) <- allocateFramebuffer settings renderPass

  pure Offscreen
    { oRenderPass  = renderPass
    , oRenderArea  = fullSurface
    , oExtent      = sExtent
    , oMipLevels   = 1
    , oLayers      = sLayers
    , oClear       = clear
    , oColor       = color
    , oDepth       = depth
    , oFrameBuffer = framebuffer
    , oRelease     = refcounted
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = sExtent
      }

    clear = Vector.fromList
      [ Vk.Color $ Vk.Float32 0 0 0 1
      , Vk.DepthStencil (Vk.ClearDepthStencilValue 1.0 0)
      ]

-- ** Render pass

allocateRenderPass
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     )
  => Settings
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPass settings = do
  device <- asks getDevice

  res <-
    if sMultiView settings then
      Vk.withRenderPass device createInfoMulti Nothing Resource.allocate
    else
      Vk.withRenderPass device createInfo Nothing Resource.allocate
  Debug.nameObject device (snd res) $ "Offscreen:" <> encodeUtf8 (sLabel settings)
  pure res
  where
    createInfo =
      zero
        { Vk.attachments  = Vector.fromList [color, depth]
        , Vk.subpasses    = Vector.fromList [subpass]
        , Vk.dependencies = deps
        }

    createInfoMulti =
      createInfo
        ::& Khr.RenderPassMultiviewCreateInfo
          { Khr.viewMasks        = [2 ^ sLayers settings - 1]
          , Khr.viewOffsets      = []
          , Khr.correlationMasks = [0]
          }
        :& ()

    color = zero
      { Vk.format         = sFormat settings
      , Vk.samples        = sSamples settings
      , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.finalLayout    = Vk.IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
      , 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
      }

    depth = zero
      { Vk.format         = sDepthFormat settings
      , Vk.samples        = sSamples settings
      , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.finalLayout    = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
      , Vk.loadOp         = Vk.ATTACHMENT_LOAD_OP_CLEAR
      , Vk.storeOp        = Vk.ATTACHMENT_STORE_OP_DONT_CARE
      , Vk.stencilLoadOp  = Vk.ATTACHMENT_LOAD_OP_DONT_CARE
      , Vk.stencilStoreOp = Vk.ATTACHMENT_STORE_OP_DONT_CARE
      }

    subpass = zero
      { Vk.pipelineBindPoint = Vk.PIPELINE_BIND_POINT_GRAPHICS
      , Vk.colorAttachments = Vector.singleton zero
          { Vk.attachment = 0
          , Vk.layout     = Vk.IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
          }
      , Vk.depthStencilAttachment = Just zero
          { Vk.attachment = 1
          , Vk.layout     = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
          }
      }

    {- BUG: Validation Error: [ SYNC-HAZARD-READ_AFTER_WRITE ]

      vkCmdDraw: Hazard READ_AFTER_WRITE for VkImageView 0x2c0a7b0[],
        in VkCommandBuffer 0x3297680[],
        and VkPipeline 0x3203390[Global.Render.EnvCube.Pipeline],
        VkDescriptorSet 0x3248050[],
        type: VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
        imageLayout: VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        binding #3, index 5.
      Access info (
        usage: SYNC_FRAGMENT_SHADER_SHADER_STORAGE_READ,
        prior_usage: SYNC_COLOR_ATTACHMENT_OUTPUT_COLOR_ATTACHMENT_WRITE,
        write_barriers: 0,
        command: vkCmdBeginRenderPass,
        seq_no: 26,
        reset_no: 1
      ).
    -}
    deps =
      [ zero
          { Vk.dependencyFlags = Vk.DEPENDENCY_BY_REGION_BIT
          , Vk.srcSubpass      = Vk.SUBPASS_EXTERNAL
          , Vk.dstSubpass      = 0
          , Vk.srcStageMask    = Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
          , Vk.dstStageMask    = Vk.PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT .|. Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT .|. Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
          , Vk.srcAccessMask   = Vk.ACCESS_SHADER_READ_BIT .|. Vk.ACCESS_SHADER_WRITE_BIT
          , Vk.dstAccessMask   = Vk.ACCESS_COLOR_ATTACHMENT_WRITE_BIT .|. Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
          }
      , zero
          { Vk.dependencyFlags = Vk.DEPENDENCY_BY_REGION_BIT
          , Vk.srcSubpass      = 0
          , Vk.dstSubpass      = Vk.SUBPASS_EXTERNAL
          , Vk.srcStageMask    = Vk.PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT .|. Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT .|. Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT
          , Vk.dstStageMask    = Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
          , Vk.srcAccessMask   = Vk.ACCESS_COLOR_ATTACHMENT_WRITE_BIT .|. Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
          , Vk.dstAccessMask   = Vk.ACCESS_SHADER_READ_BIT
          }
      ]

-- ** Framebuffer

type FramebufferOffscreen =
  ( RefCounted
  , Image.AllocatedImage
  , Image.AllocatedImage
  , Vk.Framebuffer
  )

allocateFramebuffer
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     )
  => Settings
  -> Vk.RenderPass
  -> m FramebufferOffscreen
allocateFramebuffer Settings{..} renderPass = do
  context <- ask
  let device = getDevice context
  let
    Vk.Extent2D{width, height} = sExtent
    mipLevels = extentMips sExtent

  (colorKey, color) <- Resource.allocate
    ( Image.create
        context
        (Just $ sLabel <> ".color")
        Vk.IMAGE_ASPECT_COLOR_BIT
        sExtent
        (if sMipMap then mipLevels else 1)
        sLayers
        sSamples
        sFormat
        ( Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT .|.
          Vk.IMAGE_USAGE_SAMPLED_BIT .|.
          Vk.IMAGE_USAGE_TRANSFER_SRC_BIT
        )
    )
    (Image.destroy context)

  (depthKey, depth) <- Resource.allocate
    ( Image.create
        context
        (Just $ sLabel <> ".depth")
        Vk.IMAGE_ASPECT_DEPTH_BIT
        sExtent
        (if sMipMap then mipLevels else 1)
        sLayers
        sSamples
        sDepthFormat
        (Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT .|. Vk.IMAGE_USAGE_SAMPLED_BIT)
    )
    (Image.destroy context)

  let
    attachments = Vector.fromList
      [ Image.aiImageView color
      , 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 =
      if sMultiView then
        1
      else
        sLayers

    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 $ encodeUtf8 sLabel <> ".FB"

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

  pure (release, color, depth, framebuffer)

extentMips :: Num t => Vk.Extent2D -> t
extentMips Vk.Extent2D{width, height} = go 1 (min width height)
  where
    go levels = \case
      0 ->
        levels
      1 ->
        levels
      side ->
        go (levels + 1) (side `div` 2)