{-# LANGUAGE OverloadedLists #-}

module Render.Pass.Compose
  ( Compose(..)
  , allocate
  , update
  , 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.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)

import Engine.Types.RefCounted (RefCounted, releaseRefCounted, resourceTRefCount)
import Engine.Types.RefCounted qualified as RefCounted
import Engine.Vulkan.Types (HasSwapchain(..), HasRenderPass(..), MonadVulkan, RenderPass(..), getDevice)
import Render.Pass (usePass)
import Resource.Region qualified as Region

-- | Composition/postprocessing/presentation pass
--
-- Can be used to transfer images from "Render.Pass.Offscreen.Offscreen" passes and tonemapping.
--
-- Color attachments are derived from swapchain.
-- The pass optmized for image transfer: it has no depth attachment and does not clear.
-- Use image blitting that convers the whole area or a fullscreen shader.
data Compose = Compose
  { renderPass   :: Vk.RenderPass
  , frameBuffers :: Vector Vk.Framebuffer
  , renderArea   :: Vk.Rect2D
  , clear        :: Vector Vk.ClearValue
  , release      :: RefCounted
  }

instance HasRenderPass Compose where
  getRenderPass   = renderPass
  getFramebuffers = frameBuffers
  getClearValues  = clear
  getRenderArea   = renderArea

instance RenderPass Compose where
  updateRenderpass = update
  refcountRenderpass = resourceTRefCount . release

allocate
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain swapchain
     )
  => swapchain
  -> m Compose
allocate swapchain = do
  renderPass <- allocateRenderPass format
  (refcounted, framebuffers) <- RefCounted.wrapped $ Region.run do
    Region.logDebug
      "Allocating Compose resources"
      "Releasing Compose resources"
    allocateFramebufferSwapchain swapchain renderPass

  pure Compose
    { renderPass   = renderPass
    , renderArea   = fullSurface
    , clear        = mempty
    , frameBuffers = framebuffers
    , release      = refcounted
    }
  where
    format = getSurfaceFormat swapchain

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

update
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Compose
  -> m Compose
update swapchain old@Compose{release, renderPass} = do
  releaseRefCounted release
  (refcounted, framebuffers) <- RefCounted.wrapped $ Region.run do
    allocateFramebufferSwapchain swapchain renderPass
  pure old
    { frameBuffers = framebuffers
    , release      = refcounted
    , renderArea   = fullSurface
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = getSurfaceExtent swapchain
      }

-- ** Render pass

allocateRenderPass
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     )
  => Vk.Format
  -> m Vk.RenderPass
allocateRenderPass format = do
  device <- asks getDevice
  (_key, rp) <- Vk.withRenderPass device createInfo Nothing Resource.allocate
  Debug.nameObject device rp "Compose"
  pure rp
  where
    createInfo = zero
      { Vk.attachments  = [color]
      , Vk.subpasses    = [subpass]
      , Vk.dependencies = [colorDeps]
      }

    color = zero
      { Vk.format      = format
      , Vk.samples     = Vk.SAMPLE_COUNT_1_BIT
      , Vk.loadOp      = Vk.ATTACHMENT_LOAD_OP_DONT_CARE
      , Vk.finalLayout = Vk.IMAGE_LAYOUT_PRESENT_SRC_KHR
      }

    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
          }
      }

    colorDeps = zero
      { Vk.srcSubpass    = Vk.SUBPASS_EXTERNAL
      , Vk.dstSubpass    = 0
      , Vk.srcStageMask  = colorOut
      , Vk.srcAccessMask = zero
      , Vk.dstStageMask  = colorOut
      , Vk.dstAccessMask = colorRW
      }
      where
        colorOut =
          Vk.PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT

        colorRW =
          Vk.ACCESS_COLOR_ATTACHMENT_READ_BIT .|.
          Vk.ACCESS_COLOR_ATTACHMENT_WRITE_BIT

-- ** Framebuffer

allocateFramebufferSwapchain
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Vk.RenderPass
  -> m (Vector Vk.Framebuffer)
allocateFramebufferSwapchain swapchain renderPass =
  Vector.iforM (getSwapchainViews swapchain) \ix colorView -> do
    let
      fbCI = zero
        { Vk.renderPass  = renderPass
        , Vk.width       = width
        , Vk.height      = height
        , Vk.attachments = [colorView]
        , Vk.layers      = 1
        }

    device <- asks getDevice
    (_key, fb) <- Vk.withFramebuffer device fbCI Nothing Resource.allocate
    Debug.nameObject device fb $
      "Compose.FB:" <> fromString (show @Int ix)
    pure fb
  where
    Vk.Extent2D{width, height} = getSurfaceExtent swapchain