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.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 Render.Pass (usePass)
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
  updateRenderpass = updateMsaa
  refcountRenderpass = resourceTRefCount . fmRelease

allocateMsaa
  :: ( Resource.MonadResource m
     , MonadVulkan env m
     , HasLogFunc env
     , HasSwapchain swapchain
     )
  => swapchain
  -> m ForwardMsaa
allocateMsaa swapchain = do
  logDebug "Allocating ForwardMsaa resources"
  (_rpKey, renderPass) <- allocateRenderPassMsaa swapchain
  (refcounted, color, depth, framebuffers) <- allocateFramebufferMsaa swapchain 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 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
  (refcounted, color, depth, framebuffers) <- allocateFramebufferMsaa swapchain fmRenderPass
  pure old
    { fmColor        = color
    , fmDepth        = depth
    , fmFrameBuffers = framebuffers
    , fmRelease      = refcounted
    , 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
  let
    format = getSurfaceFormat swapchain
    depthFormat = getDepthFormat swapchain
    msaa = getMultisample swapchain

    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 swapchain
     )
  => swapchain
  -> Vk.RenderPass
  -> m FramebuffersMsaa
allocateFramebufferMsaa swapchain renderPass = do
  context <- ask
  let extent@Vk.Extent2D{width, height} = getSurfaceExtent swapchain

  (colorKey, color) <- Resource.allocate
    ( Image.create
        context
        (Just $ "ForwardMSAA.color")
        Vk.IMAGE_ASPECT_COLOR_BIT
        extent
        1
        1
        (getMultisample swapchain)
        (getSurfaceFormat swapchain)
        Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT
    )
    (Image.destroy context)

  (depthKey, depth) <- Resource.allocate
    ( Image.create
        context
        (Just $ "ForwardMSAA.depth")
        Vk.IMAGE_ASPECT_DEPTH_BIT
        extent
        1
        1
        (getMultisample swapchain)
        (getDepthFormat swapchain)
        Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT
    )
    (Image.destroy context)

  perView <- Vector.iforM (getSwapchainViews swapchain) \ix 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
        }

      device = getDevice context
    res <- Vk.withFramebuffer device fbCI Nothing Resource.allocate
    Debug.nameObject device (snd res) $
      "ForwardMSAA.FB:" <> fromString (show @Int ix)

    pure res

  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)