{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedRecordDot #-}

module Render.Pass.Offscreen
  ( Settings(..)
  , allocate
  , Offscreen(..)
  , colorTexture
  , colorCube
  , depthTexture
  , depthCube
  ) where

import RIO

import Control.Monad.Trans.Resource (ResourceT)
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)
import Engine.Types.RefCounted qualified as RefCounted
import Engine.Vulkan.Types (HasVulkan(..), HasRenderPass(..), RenderPass(..), MonadVulkan)
import Resource.Image (AllocatedImage)
import Resource.Image qualified as Image
import Resource.Region qualified as Region
import Resource.Texture (CubeMap, Flat, Texture(..))
import Resource.Vulkan.Named qualified as Named

{- 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
  , sColorLayout :: Maybe Vk.ImageLayout -- ^ Target color format when used for export.
  , sDepthFormat :: Vk.Format
  , sDepthLayout :: Maybe Vk.ImageLayout -- ^ Target depth format when used for export.
  , 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 = RefCounted.resourceTRefCount . oRelease

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

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

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

depthCube :: Offscreen -> Texture CubeMap
depthCube Offscreen{..} = Texture
  { tFormat         = oDepth.aiFormat
  , 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
  (_passKey, renderPass) <- allocateRenderPass settings

  (release, resources) <- RefCounted.wrapped $ Region.run do
    Region.logDebug
      ("Allocating Offscreen resources for " <> display sLabel)
      ("Releasing Offscreen resources for " <> display sLabel)
    allocateFramebuffer settings renderPass
  let (color, depth, framebuffer) = resources

  pure Offscreen
    { oRenderPass  = renderPass
    , oRenderArea  = fullSurface
    , oExtent      = sExtent
    , oMipLevels   = 1
    , oLayers      = sLayers
    , oClear       = clear
    , oColor       = color
    , oDepth       = depth
    , oFrameBuffer = framebuffer
    , oRelease     = release
    }
  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    = fromMaybe Vk.IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL (sColorLayout settings)
      , 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    = fromMaybe Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL (sDepthLayout settings)
      , 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 =
  ( Image.AllocatedImage
  , Image.AllocatedImage
  , Vk.Framebuffer
  )

allocateFramebuffer
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     )
  => Settings
  -> Vk.RenderPass
  -> ResourceT m FramebufferOffscreen
allocateFramebuffer Settings{..} renderPass = do
  Region.logDebug
    ("Allocating Offscreen resources for " <> display sLabel)
    ("Releasing Offscreen resources for " <> display sLabel)

  color <- Image.allocate
    (Just $ sLabel <> ".color")
    Vk.IMAGE_ASPECT_COLOR_BIT
    (Image.inflateExtent sExtent 1)
    (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
    )

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

  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
      }

  device <- asks getDevice
  framebuffer <- Vk.createFramebuffer device fbCI Nothing
  void $! Resource.register $ Vk.destroyFramebuffer device framebuffer Nothing
  Named.object framebuffer $ sLabel <> ".FB"

  pure (color, depth, framebuffer)
  where
    Vk.Extent2D{width, height} = sExtent
    mipLevels = extentMips sExtent

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)