module Render.ForwardMsaa
  ( ForwardMsaa(..)
  , allocateMsaa
  , updateMsaa
  , 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.Zero (zero)

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

-- * 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
  updateRenderpass = updateMsaa
  refcountRenderpass = resourceTRefCount . fmRelease

allocateMsaa
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain swapchain
     )
  => swapchain
  -> m ForwardMsaa
allocateMsaa swapchain = do
  (_rpKey, renderPass) <- allocateRenderPassMsaa swapchain

  (release, resources) <- RefCounted.wrapped $ Region.run do
    Region.logDebug
      "Allocating ForwardMsaa resources"
      "Releasing ForwardMsaa resources"
    allocateFramebufferMsaa swapchain renderPass
  let (color, depth, framebuffers) = resources

  pure ForwardMsaa
    { fmRenderPass   = renderPass
    , fmRenderArea   = fullSurface
    , fmClear        = clear
    , fmColor        = color
    , fmDepth        = depth
    , fmFrameBuffers = framebuffers
    , fmRelease      = release
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = getSurfaceExtent swapchain
      }
    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 swapchain
     )
  => swapchain
  -> ForwardMsaa
  -> m ForwardMsaa
updateMsaa swapchain old@ForwardMsaa{fmRelease, fmRenderPass} = do
  releaseRefCounted fmRelease
  (release, resources) <- RefCounted.wrapped $ Region.run do
    Region.logDebug
      "Updating ForwardMsaa resources"
      "Releasing ForwardMsaa resources"
    allocateFramebufferMsaa swapchain fmRenderPass
  let (color, depth, framebuffers) = resources
  pure old
    { fmColor        = color
    , fmDepth        = depth
    , fmFrameBuffers = framebuffers
    , fmRelease      = release
    , fmRenderArea   = fullSurface
    }
  where
    fullSurface = Vk.Rect2D
      { Vk.offset = zero
      , Vk.extent = getSurfaceExtent swapchain
      }

-- ** Render pass

allocateRenderPassMsaa
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     , HasSwapchain swapchain
     )
  => swapchain
  -> m (Resource.ReleaseKey, Vk.RenderPass)
allocateRenderPassMsaa swapchain = do
  device <- asks getDevice

  res <- Vk.withRenderPass device createInfo Nothing Resource.allocate
  Named.object (snd res) "ForwardMSAA"
  pure res
  where
    format = getSurfaceFormat swapchain
    depthFormat = getDepthFormat swapchain
    msaa = getMultisample swapchain

    onMsaa :: a -> a -> a
    onMsaa none more = case msaa of
      Vk.SAMPLE_COUNT_1_BIT ->
        none
      _mucho ->
        more

    attachments =
      color :
      depth :
      onMsaa [] [colorResolve]

    createInfo = zero
      { Vk.attachments  = Vector.fromList attachments
      , Vk.subpasses    = Vector.fromList [subpass]
      , Vk.dependencies = Vector.fromList [colorDeps, depthDeps]
      }

    -- | Main color-rendering attachment. Used as presentation when MSAA is disabled.
    color = zero
      { Vk.format         = format
      , Vk.samples        = msaa
      , Vk.finalLayout    = finalColorLayout
      , 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
      }
    finalColorLayout =
      onMsaa
        Vk.IMAGE_LAYOUT_PRESENT_SRC_KHR
        Vk.IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL

    depth = zero
      { Vk.format         = depthFormat
      , 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
      }

    -- | Extra attachment to collect samples AND present.
    colorResolve = 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 =
          onMsaa mempty $
            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 =
  ( Image.AllocatedImage
  , Image.AllocatedImage
  , Vector Vk.Framebuffer
  )

allocateFramebufferMsaa
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasSwapchain swapchain
     )
  => swapchain
  -> Vk.RenderPass
  -> m FramebuffersMsaa
allocateFramebufferMsaa swapchain renderPass = do
  color <- Image.allocate
    (Just "ForwardMSAA.color")
    Vk.IMAGE_ASPECT_COLOR_BIT
    (Image.inflateExtent extent 1)
    1
    1
    (getMultisample swapchain)
    (getSurfaceFormat swapchain)
    Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT

  depth <- Image.allocate
    (Just "ForwardMSAA.depth")
    Vk.IMAGE_ASPECT_DEPTH_BIT
    (Image.inflateExtent extent 1)
    1
    1
    (getMultisample swapchain)
    (getDepthFormat swapchain)
    Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT

  perView <- Vector.iforM (getSwapchainViews swapchain) \ix colorResolve -> do
    let
      attachments = Vector.fromList
        case getMultisample swapchain of
          Vk.SAMPLE_COUNT_1_BIT ->
            [ colorResolve
            , Image.aiImageView depth
            ]
          _ ->
            [ Image.aiImageView color
            , Image.aiImageView depth
            , colorResolve
            ]

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

    device <- asks getDevice
    (_fbKey, fb) <- Vk.withFramebuffer device fbCI Nothing Resource.allocate
    Named.object fb $
      "ForwardMSAA.FB:" <> fromString (show @Int ix)
    pure fb

  pure (color, depth, perView)
  where
    extent@Vk.Extent2D{width, height} = getSurfaceExtent swapchain