module Render.ForwardMsaa where

import RIO

import Control.Monad.Trans.Resource qualified as Resource
import Data.Bits ((.|.))
import Data.Vector qualified as Vector
import RIO.Vector.Partial ((!))
import Vulkan.Core10 qualified as Vk
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)

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

-- * Simple MSAA-enabled pass

data ForwardMsaa = ForwardMsaa
  { fmRenderPass   :: Vk.RenderPass
  , fmColor        :: AllocatedImage
  , fmDepth        :: AllocatedImage
  , fmFrameBuffers :: Vector Vk.Framebuffer
  , fmRenderArea   :: Vk.Rect2D
  , fmClear        :: Vector Vk.ClearValue
  , fmRelease      :: RefCounted
  }

instance HasRenderPass ForwardMsaa where
  getRenderPass          = fmRenderPass
  getFramebuffers        = fmFrameBuffers
  getClearValues         = fmClear
  getRenderArea          = fmRenderArea

instance RenderPass ForwardMsaa where
  allocateRenderpass_ = allocateMsaa -- error "Thou shalt allocateMsaa instead"

  updateRenderpass = updateMsaa

  refcountRenderpass = resourceTRefCount . fmRelease

allocateMsaa
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain context
     )
  => context
  -> m ForwardMsaa
allocateMsaa context = do
  logDebug "Allocating ForwardMsaa resources"
  (_rpKey, renderPass) <- allocateRenderPassMsaa context
  (refcounted, color, depth, framebuffers) <- allocateFramebufferMsaa context renderPass

  pure ForwardMsaa
    { fmRenderPass   = renderPass
    , fmRenderArea   = fullSurface
    , fmClear        = clear
    , fmColor        = color
    , fmDepth        = depth
    , fmFrameBuffers = framebuffers
    , fmRelease      = refcounted
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = getSurfaceExtent context
      }
    clear = Vector.fromList
      [ clearColor
      , Vk.DepthStencil (Vk.ClearDepthStencilValue 1.0 0)
      , clearColor
      ]
    clearColor = Vk.Color zero

updateMsaa
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain context
     )
  => context
  -> ForwardMsaa
  -> m ForwardMsaa
updateMsaa context old@ForwardMsaa{fmRelease, fmRenderPass} = do
  releaseRefCounted fmRelease
  (refcounted, color, depth, framebuffers) <- allocateFramebufferMsaa context fmRenderPass
  pure old
    { fmColor        = color
    , fmDepth        = depth
    , fmFrameBuffers = framebuffers
    , fmRelease      = refcounted
    , fmRenderArea   = fullSurface
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = getSurfaceExtent context
      }

-- ** Render pass

allocateRenderPassMsaa
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     , HasSwapchain context
     )
  => context
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPassMsaa context = do
  device <- asks getDevice
  let
    format = getSurfaceFormat context
    depthFormat = getDepthFormat context
    msaa = getMultisample context

    attachments =
      [ color format msaa
      , depth depthFormat msaa
      , colorResolve format
      ]

  res@(_key, object) <- Vk.withRenderPass device (createInfo attachments) Nothing Resource.allocate
  Debug.nameObject device object "ForwardMSAA"
  pure res
  where
    createInfo attachments = zero
      { Vk.attachments  = Vector.fromList attachments
      , Vk.subpasses    = Vector.fromList [subpass]
      , Vk.dependencies = Vector.fromList [colorDeps, depthDeps]
      }

    color format msaa = zero
      { Vk.format         = format
      , Vk.samples        = msaa
      , Vk.finalLayout    = Vk.IMAGE_LAYOUT_COLOR_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
      , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
      }

    depth format msaa = zero
      { Vk.format         = format
      , Vk.samples        = msaa
      , 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
      , Vk.initialLayout  = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.finalLayout    = Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
      }

    colorResolve format = zero
      { Vk.format      = format
      , Vk.samples     = Vk.SAMPLE_COUNT_1_BIT
      , Vk.finalLayout = Vk.IMAGE_LAYOUT_PRESENT_SRC_KHR
      , Vk.loadOp      = Vk.ATTACHMENT_LOAD_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
          }
      , Vk.resolveAttachments = Vector.singleton zero
          { Vk.attachment = 2
          , 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

    depthDeps = zero
      { Vk.srcSubpass    = Vk.SUBPASS_EXTERNAL
      , Vk.dstSubpass    = 0
      , Vk.srcStageMask  = fragTests
      , Vk.srcAccessMask = depthW
      , Vk.dstStageMask  = fragTests
      , Vk.dstAccessMask = depthRW
      }
      where
        fragTests =
          Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT .|.
          Vk.PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT

        depthW =
          Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT

        depthRW =
          Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT .|.
          depthW

-- ** Framebuffer

type FramebuffersMsaa =
  ( RefCounted
  , Image.AllocatedImage
  , Image.AllocatedImage
  , Vector Vk.Framebuffer
  )

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

  (colorKey, color) <- Resource.allocate
    (Image.createColorResource context' extent)
    (Image.destroyAllocatedImage context')

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

  perView <- for (getSwapchainViews context) \colorResolve -> do
    let
      attachments = Vector.fromList
        [ Image.aiImageView color
        , Image.aiImageView depth
        , colorResolve
        ]

      fbCI = zero
        { Vk.renderPass  = renderPass
        , Vk.width       = width
        , Vk.height      = height
        , Vk.attachments = attachments
        , Vk.layers      = 1
        }

    Vk.withFramebuffer device fbCI Nothing Resource.allocate

  let (framebufferKeys, framebuffers) = Vector.unzip perView
  releaseDebug <- toIO $ logDebug "Releasing ForwardMsaa resources"
  release <- newRefCounted do
    releaseDebug
    Resource.release colorKey
    Resource.release depthKey
    traverse_ Resource.release framebufferKeys

  pure (release, color, depth, framebuffers)

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
  , Vk.renderArea  = getRenderArea rp
  , Vk.clearValues = getClearValues rp
  }