{-# LANGUAGE OverloadedRecordDot #-}

module Resource.Image
  ( AllocatedImage(..)
  , allocate
  , allocateView

  , DstImage
  , allocateDst
  , copyBufferToDst
  , updateFromStorable

  , transitionLayout
  , copyBufferToImage

  , subresource
  , inflateExtent
  ) where

import RIO

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

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues(..))
import Resource.Buffer qualified as Buffer
import Resource.CommandBuffer (oneshot_)
import Resource.Vulkan.Named qualified as Named

data AllocatedImage = AllocatedImage
  { aiAllocation  :: VMA.Allocation
  , aiExtent      :: Vk.Extent3D
  , aiFormat      :: Vk.Format
  , aiImage       :: Vk.Image
  , aiImageView   :: Vk.ImageView
  , aiImageRange  :: Vk.ImageSubresourceRange
  }
  deriving (Show)

allocate
  :: ( MonadVulkan env io
     , MonadResource io
     )
  => Maybe Text
  -> Vk.ImageAspectFlags
  -> "image dimensions" ::: Vk.Extent3D
  -> "mip levels" ::: Word32
  -> "stored layers" ::: Word32
  -> Vk.SampleCountFlagBits
  -> Vk.Format
  -> Vk.ImageUsageFlags
  -> io AllocatedImage
allocate mlabel aspect extent mipLevels numLayers samples format usage = do
  allocator <- asks getAllocator

  (image, allocation, _info) <- VMA.createImage
    allocator
    imageCI
    imageAllocationCI
  void $! Resource.register $
    VMA.destroyImage allocator image allocation
  traverse_ (Named.object image) mlabel

  imageView <- allocateView image format subr
  traverse_ (Named.object imageView) $ fmap (<> ":view") mlabel

  pure AllocatedImage
    { aiAllocation = allocation
    , aiExtent     = extent
    , aiFormat     = format
    , aiImage      = image
    , aiImageView  = imageView
    , aiImageRange = subr
    }
  where
    imageType =
      case extent of
        Vk.Extent3D{depth=1} ->
          Vk.IMAGE_TYPE_2D
        _ ->
          Vk.IMAGE_TYPE_3D

    imageCI = zero
      { Vk.imageType     = imageType
      , Vk.flags         = createFlags
      , Vk.format        = format
      , Vk.extent        = extent
      , Vk.mipLevels     = mipLevels
      , 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
      }

    createFlags =
      case numLayers of
        6 ->
          Vk.IMAGE_CREATE_CUBE_COMPATIBLE_BIT
        _ ->
          zero

    subr = subresource aspect mipLevels numLayers

allocateView
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Vk.Image
  -> Vk.Format
  -> Vk.ImageSubresourceRange
  -> m Vk.ImageView
allocateView image format subr = do
  device <- asks getDevice
  imageView <- Vk.createImageView
    device
    imageViewCI
    Nothing
  void $! Resource.register $ Vk.destroyImageView device imageView Nothing
  pure imageView

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

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

newtype DstImage = DstImage AllocatedImage

-- | Allocate an image and transition it into TRANSFER_DST_OPTIOMAL
allocateDst
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Queues Vk.CommandPool
  -> Maybe Text
  -> ("image dimensions" ::: Vk.Extent3D)
  -> ("mip levels" ::: Word32)
  -> ("stored layers" ::: Word32)
  -> Vk.Format
  -> m DstImage
allocateDst pool name extent3d mipLevels numLayers format = do
  ai <- allocate
    name
    Vk.IMAGE_ASPECT_COLOR_BIT
    extent3d
    mipLevels
    numLayers
    Vk.SAMPLE_COUNT_1_BIT
    format
    (Vk.IMAGE_USAGE_SAMPLED_BIT .|. Vk.IMAGE_USAGE_TRANSFER_DST_BIT)

  transitionLayout
    pool
    (aiImage ai)
    mipLevels
    numLayers -- XXX: arrayLayers is always 0 for now
    format
    Vk.IMAGE_LAYOUT_UNDEFINED
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

  pure $ DstImage ai

copyBufferToDst
  :: ( MonadVulkan env m
     , Integral deviceSize
     , Foldable t
     )
  => Queues Vk.CommandPool
  -> Vk.Buffer
  -> DstImage
  -> "mip offsets" ::: t deviceSize
  -> m AllocatedImage
copyBufferToDst pool source (DstImage ai) offsets = do
  copyBufferToImage
    pool
    source
    (aiImage ai)
    (aiExtent ai)
    offsets
    layerCount

  transitionLayout
    pool
    (aiImage ai)
    levelCount
    layerCount -- XXX: arrayLayers is always 0 for now
    (aiFormat ai)
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

  pure ai
  where
    Vk.ImageSubresourceRange{layerCount, levelCount} = aiImageRange ai

{-# INLINE updateFromStorable #-}
updateFromStorable
  :: ( Storable a
     , MonadVulkan env m
     , MonadResource m
     )
  => Queues Vk.CommandPool
  -> AllocatedImage
  -> Storable.Vector a
  -> m AllocatedImage
updateFromStorable pools ai update = do
  (_transient, staging) <-
    Buffer.allocateCoherent
      (Just "updateFromStorable:staging")
      Vk.BUFFER_USAGE_TRANSFER_SRC_BIT
      1
      update

  transitionLayout
    pools
    ai.aiImage
    1
    1
    ai.aiFormat
    Vk.IMAGE_LAYOUT_UNDEFINED
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

  copyBufferToImage
    pools
    staging.aBuffer
    ai.aiImage
    (Vk.Extent3D ai.aiExtent.width ai.aiExtent.height 1)
    (Vector.singleton 0 :: Vector Word32) -- offsets
    1

  transitionLayout
    pools
    ai.aiImage
    1
    1
    ai.aiFormat
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

  -- logDebug "Updating map texture... done"
  pure ai

{-# INLINE transitionLayout #-}
transitionLayout
  :: ( MonadVulkan env m
    --  , MonadUnliftIO m
     )
  => Queues Vk.CommandPool
  -> Vk.Image
  -> "mip levels" ::: Word32
  -> "layer count" ::: Word32
  -> Vk.Format
  -> "old" ::: Vk.ImageLayout
  -> "new" ::: Vk.ImageLayout
  -> m ()
transitionLayout pool image mipLevels layerCount format old new = do
  ctx <- ask
  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
      }

{-# INLINE copyBufferToImage #-}
copyBufferToImage
  :: ( Foldable t
     , Integral deviceSize
     , MonadVulkan env m
     )
  => Queues Vk.CommandPool
  -> Vk.Buffer
  -> Vk.Image
  -> "base extent" ::: Vk.Extent3D
  -> "mip offsets" ::: t deviceSize
  -> "layer count" ::: Word32
  -> m ()
copyBufferToImage pools src dst Vk.Extent3D{..} mipOffsets layerCount = do
  context <- ask
  oneshot_ context pools 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  = max 1 $ depth `shiftR` mipLevel
            }
        }

-- * Helpers

{-# INLINEABLE inflateExtent #-}
inflateExtent :: Vk.Extent2D -> Word32 -> Vk.Extent3D
inflateExtent Vk.Extent2D{..} depth = Vk.Extent3D{..}

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
  }

guessViewType :: Vk.ImageSubresourceRange -> Vk.ImageViewType
guessViewType Vk.ImageSubresourceRange{layerCount} =
  case layerCount of
    1 ->
      Vk.IMAGE_VIEW_TYPE_2D
    6 ->
      Vk.IMAGE_VIEW_TYPE_CUBE
    _ ->
      Vk.IMAGE_VIEW_TYPE_2D_ARRAY