module Resource.Image where

import RIO

import Data.Bits (shiftR, (.|.))
import RIO.Vector qualified as Vector
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.NamedType ((:::))
import Vulkan.Utils.Debug qualified as Debug
import Vulkan.Zero (zero)
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(..), HasSwapchain(..), Queues(..))
import Resource.CommandBuffer (oneshot_)

data AllocatedImage = AllocatedImage
  { aiAllocation :: VMA.Allocation
  , aiImage      :: Vk.Image
  , aiImageView  :: Vk.ImageView
  }
  deriving (Show)

createColorResource
  :: ( MonadIO io
     , HasVulkan ctx
     , HasSwapchain ctx
     )
  => ctx
  -> Vk.Extent2D
  -> io AllocatedImage
createColorResource context Vk.Extent2D{width, height} = do
  let
    device    = getDevice context
    allocator = getAllocator context
    format    = getSurfaceFormat context
    msaa      = getMultisample context

  (image, allocation, _info) <- VMA.createImage
    allocator
    (imageCI format msaa)
    imageAllocationCI
  Debug.nameObject device image "ColorResource.image"

  imageView <- Vk.createImageView
    device
    (imageViewCI image format)
    Nothing
  Debug.nameObject device image "ColorResource.view"

  pure AllocatedImage
    { aiAllocation = allocation
    , aiImage      = image
    , aiImageView  = imageView
    }
  where
    imageCI format msaa = zero
      { Vk.imageType     = Vk.IMAGE_TYPE_2D
      , Vk.format        = format
      , Vk.extent        = Vk.Extent3D width height 1
      , Vk.mipLevels     = 1
      , Vk.arrayLayers   = 1
      , Vk.tiling        = Vk.IMAGE_TILING_OPTIMAL
      , Vk.initialLayout = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.usage         = Vk.IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT .|. Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT
      , Vk.sharingMode   = Vk.SHARING_MODE_EXCLUSIVE
      , Vk.samples       = msaa
      }

    imageAllocationCI = zero
      { VMA.usage         = VMA.MEMORY_USAGE_GPU_ONLY
      , VMA.requiredFlags = Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
      }

    imageViewCI image format = zero
      { Vk.image            = image
      , Vk.viewType         = Vk.IMAGE_VIEW_TYPE_2D
      , Vk.format           = format
      , Vk.components       = zero
      , Vk.subresourceRange = subr
      }

    subr = zero
      { Vk.aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
      , Vk.baseMipLevel   = 0
      , Vk.levelCount     = 1
      , Vk.baseArrayLayer = 0
      , Vk.layerCount     = 1
      }

createDepthResource
  :: ( MonadIO io
     , HasVulkan context
     , HasSwapchain context
     )
  => context
  -> Vk.Extent2D
  -> "shadowmap layers" ::: Maybe Word32
  -> io AllocatedImage
createDepthResource context Vk.Extent2D{width, height} depthLayers = do
  let
    device    = getDevice context
    allocator = getAllocator context

    depthFormat = getDepthFormat context
    msaa        = getMultisample context

    (samples, usage, numLayers) =
      case depthLayers of
        Nothing ->
          ( msaa
          , Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT
          , 1
          )
        Just nl ->
          ( Vk.SAMPLE_COUNT_1_BIT
          , Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT .|.
            Vk.IMAGE_USAGE_SAMPLED_BIT
          , nl
          )

  (image, allocation, _info) <- VMA.createImage
    allocator
    (imageCI depthFormat usage samples numLayers)
    imageAllocationCI
  Debug.nameObject device image "DepthResource.image"

  imageView <- Vk.createImageView
    device
    (imageViewCI depthFormat image numLayers)
    Nothing
  Debug.nameObject device image "DepthResource.view"

  pure AllocatedImage
    { aiAllocation = allocation
    , aiImage      = image
    , aiImageView  = imageView
    }
  where
    imageCI format usage samples numLayers = zero
      { Vk.imageType     = Vk.IMAGE_TYPE_2D
      , Vk.format        = format
      , Vk.extent        = Vk.Extent3D width height 1
      , Vk.mipLevels     = 1
      , Vk.arrayLayers   = numLayers
      , Vk.tiling        = Vk.IMAGE_TILING_OPTIMAL
      , Vk.initialLayout = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.usage         = usage
      , Vk.sharingMode   = Vk.SHARING_MODE_EXCLUSIVE
      , Vk.samples       = samples
      }

    imageAllocationCI = zero
      { VMA.usage         = VMA.MEMORY_USAGE_GPU_ONLY
      , VMA.requiredFlags = Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
      }

    imageViewCI format image numLayers = zero
      { Vk.image            = image
      , Vk.viewType         = viewType
      , Vk.format           = format
      , Vk.components       = zero
      , Vk.subresourceRange = subr numLayers
      }
      where
        viewType =
          if numLayers > 1 then
            Vk.IMAGE_VIEW_TYPE_2D_ARRAY
          else
            Vk.IMAGE_VIEW_TYPE_2D

    subr numLayers = zero
      { Vk.aspectMask     = Vk.IMAGE_ASPECT_DEPTH_BIT
      , Vk.baseMipLevel   = 0
      , Vk.levelCount     = 1
      , Vk.baseArrayLayer = 0
      , Vk.layerCount     = numLayers
      }

destroyAllocatedImage
  :: ( MonadIO io
     , HasVulkan context
     )
  => context
  -> AllocatedImage
  -> io ()
destroyAllocatedImage context AllocatedImage{..} = do
  -- traceM "destroyAllocatedImage"
  Vk.destroyImageView (getDevice context) aiImageView Nothing
  VMA.destroyImage (getAllocator context) aiImage aiAllocation

--------------------------------------------

transitionImageLayout
  :: (HasVulkan context)
  => context
  -> Queues Vk.CommandPool
  -> Vk.Image
  -> "mip levels" ::: Word32
  -> "layer count" ::: Word32
  -> Vk.Format
  -> ("old" ::: Vk.ImageLayout)
  -> ("new" ::: Vk.ImageLayout)
  -> RIO env ()
transitionImageLayout ctx pool image mipLevels layerCount format old new =
  case (old, new) of
    (Vk.IMAGE_LAYOUT_UNDEFINED, Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) ->
      oneshot_ ctx pool qTransfer \buf ->
        Vk.cmdPipelineBarrier
          buf
          Vk.PIPELINE_STAGE_TOP_OF_PIPE_BIT
          Vk.PIPELINE_STAGE_TRANSFER_BIT
          zero
          mempty
          mempty
          ( Vector.singleton $
              barrier Vk.IMAGE_ASPECT_COLOR_BIT zero Vk.ACCESS_TRANSFER_WRITE_BIT
          )
    (Vk.IMAGE_LAYOUT_UNDEFINED, Vk.IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL) ->
      oneshot_ ctx pool qTransfer \buf ->
        Vk.cmdPipelineBarrier
          buf
          Vk.PIPELINE_STAGE_TOP_OF_PIPE_BIT
          Vk.PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT
          zero
          mempty
          mempty
          ( Vector.singleton $
              barrier aspectMask zero $
                Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT .|.
                Vk.ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT
          )
        where
          aspectMask =
            if hasStencilComponent then
              Vk.IMAGE_ASPECT_DEPTH_BIT .|. Vk.IMAGE_ASPECT_STENCIL_BIT
            else
              Vk.IMAGE_ASPECT_DEPTH_BIT
          hasStencilComponent =
            format == Vk.FORMAT_D32_SFLOAT_S8_UINT ||
            format == Vk.FORMAT_D24_UNORM_S8_UINT

    (Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) ->
      oneshot_ ctx pool qGraphics \buf ->
        Vk.cmdPipelineBarrier
          buf
          Vk.PIPELINE_STAGE_TRANSFER_BIT
          Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
          zero
          mempty
          mempty
          ( Vector.singleton $
              barrier
                Vk.IMAGE_ASPECT_COLOR_BIT
                Vk.ACCESS_TRANSFER_WRITE_BIT
                Vk.ACCESS_SHADER_READ_BIT
          )
    _ ->
      error $ "Unsupported image layout transfer: " <> show (old, new)
  where
    barrier aspectMask srcMask dstMask = SomeStruct zero
      { Vk.srcAccessMask       = srcMask
      , Vk.dstAccessMask       = dstMask
      , Vk.oldLayout           = old
      , Vk.newLayout           = new
      , Vk.srcQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.dstQueueFamilyIndex = Vk.QUEUE_FAMILY_IGNORED
      , Vk.image               = image
      , Vk.subresourceRange    = subresource aspectMask mipLevels layerCount
      }

subresource
  :: Vk.ImageAspectFlags
  -> "mip levels" ::: Word32
  -> "layer count" ::: Word32
  -> Vk.ImageSubresourceRange
subresource aspectMask mipLevels layerCount = Vk.ImageSubresourceRange
  { aspectMask     = aspectMask
  , baseMipLevel   = 0
  , levelCount     = mipLevels -- XXX: including base
  , baseArrayLayer = 0
  , layerCount     = layerCount
  }

copyBufferToImage
  :: (HasVulkan context, Foldable t, Integral deviceSize)
  => context
  -> Queues Vk.CommandPool
  -> Vk.Buffer
  -> Vk.Image
  -> "base extent" ::: Vk.Extent3D
  -> "mip offsets" ::: t deviceSize
  -> "layer count" ::: Word32
  -> RIO env ()
copyBufferToImage ctx pool src dst Vk.Extent3D{..} mipOffsets layerCount =
  oneshot_ ctx pool qTransfer \cmd ->
    Vk.cmdCopyBufferToImage cmd src dst Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL $
      Vector.fromList copyRegions
  where
    copyRegions = do
      (offset, mipLevel) <- zip (toList mipOffsets) [0..]
      pure Vk.BufferImageCopy
        { Vk.bufferOffset      = fromIntegral offset
        , Vk.bufferRowLength   = zero -- XXX: "use extent width"
        , Vk.bufferImageHeight = zero -- XXX: "use extent height"
        , Vk.imageSubresource = Vk.ImageSubresourceLayers
            { aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
            , mipLevel       = fromIntegral mipLevel
            , baseArrayLayer = 0
            , layerCount     = layerCount
            }
        , Vk.imageOffset = zero
        , Vk.imageExtent = Vk.Extent3D
            { width  = max 1 $ width `shiftR` mipLevel
            , height = max 1 $ height `shiftR` mipLevel
            , depth  = depth
            }
        }