{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedRecordDot #-}

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

  , allocateWith
  , gpuOnly
  , gpuToCpu

  , getSubresourceLayout
  , getColorLayout0

  , DstImage
  , allocateDst
  , copyBufferToDst
  , updateFromStorable

  , DstImageHost(..)
  , allocateDstHost
  , allocateDstHost_
  , grabImageColor0

  , transitionLayout
  , copyBufferToImage

  , subresource
  , inflateExtent
  ) where

import RIO

import Data.Bits (shiftR, (.|.))
import Foreign (Ptr)
import GHC.Records (HasField(..))
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
  , aiAllocationInfo :: VMA.AllocationInfo
  , aiExtent         :: Vk.Extent3D
  , aiFormat         :: Vk.Format
  , aiImage          :: Vk.Image
  , aiImageView      :: Vk.ImageView
  , aiImageRange     :: Vk.ImageSubresourceRange
  }
  deriving (Show)

instance HasField "mappedData" AllocatedImage (Ptr ()) where
  {-# INLINE getField #-}
  getField = VMA.mappedData . aiAllocationInfo

getSubresourceLayout
  :: MonadVulkan env io
  => AllocatedImage
  -> Vk.ImageSubresource
  -> io Vk.SubresourceLayout
getSubresourceLayout ai subr = do
  device <- asks getDevice
  Vk.getImageSubresourceLayout device ai.aiImage subr

getColorLayout0
  :: MonadVulkan env io
  => AllocatedImage
  -> io Vk.SubresourceLayout
getColorLayout0 ai =
  getSubresourceLayout ai Vk.ImageSubresource
    { aspectMask = Vk.IMAGE_ASPECT_COLOR_BIT
    , mipLevel   = 0
    , arrayLayer = 0
    }

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 = allocateWith gpuOnly

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

  (image, allocation, info) <- VMA.createImage
    allocator
    imageCI
    allocationCI
  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
    , aiAllocationInfo = info
    , 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        = tiling
      , Vk.initialLayout = Vk.IMAGE_LAYOUT_UNDEFINED
      , Vk.usage         = usage
      , Vk.sharingMode   = Vk.SHARING_MODE_EXCLUSIVE
      , Vk.samples       = samples
      }

    tiling =
      if allocationCI.usage == VMA.MEMORY_USAGE_GPU_TO_CPU then
        Vk.IMAGE_TILING_LINEAR
      else
        Vk.IMAGE_TILING_OPTIMAL

    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
    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
    (aiFormat ai)
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

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

newtype DstImageHost = DstImageHost AllocatedImage

instance HasField "mappedData" DstImageHost (Ptr ()) where
  {-# INLINE getField #-}
  getField (DstImageHost ai) = VMA.mappedData ai.aiAllocationInfo

-- | Allocate a flat image to grab GPU data.
allocateDstHost_
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Queues Vk.CommandPool
  -> Maybe Text
  -> "flat image dimensions" ::: Vk.Extent2D
  -> Vk.Format
  -> m DstImageHost
allocateDstHost_ pools name extent2d =
  allocateDstHost pools name (inflateExtent extent2d 1) 1 1

-- | Allocate an image to grab GPU data.
allocateDstHost
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Queues Vk.CommandPool
  -> Maybe Text
  -> ("image dimensions" ::: Vk.Extent3D)
  -> ("mip levels" ::: Word32)
  -> ("stored layers" ::: Word32)
  -> Vk.Format
  -> m DstImageHost
allocateDstHost pool name extent3d mipLevels numLayers format = do
  ai <- allocateWith
    gpuToCpu
    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 $ DstImageHost ai

grabImageColor0
  :: MonadVulkan env m
  => Queues Vk.CommandPool
  -> Vk.CommandBuffer
  -> Vk.Image
  -> DstImageHost
  -> m Vk.SubresourceLayout
grabImageColor0 pools cb src (DstImageHost dst) = do
  Vk.cmdCopyImage
    cb
    src
    Vk.IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
    dst.aiImage
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    [ Vk.ImageCopy
        { srcSubresource =
            subrLayers
        , srcOffset =
            Vk.Offset3D 0 0 0
        , dstSubresource =
            subrLayers
        , dstOffset =
            Vk.Offset3D 0 0 0
        , extent =
            dst.aiExtent
        }
    ]
  transitionLayout
    pools
    dst.aiImage
    1
    1
    dst.aiFormat
    Vk.IMAGE_LAYOUT_UNDEFINED
    Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

  Vk.cmdPipelineBarrier
    cb
    Vk.PIPELINE_STAGE_TRANSFER_BIT
    Vk.PIPELINE_STAGE_HOST_BIT
    zero
    mempty
    mempty
    [ SomeStruct zero
        { Vk.srcAccessMask    = Vk.ACCESS_TRANSFER_WRITE_BIT
        , Vk.dstAccessMask    = Vk.ACCESS_HOST_READ_BIT
        , Vk.oldLayout        = Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
        , Vk.newLayout        = Vk.IMAGE_LAYOUT_GENERAL
        , Vk.image            = dst.aiImage
        , Vk.subresourceRange = subrRange
        }
    ]
  getSubresourceLayout dst subr
  where
    subr = Vk.ImageSubresource
      { aspectMask = Vk.IMAGE_ASPECT_COLOR_BIT
      , mipLevel   = 0
      , arrayLayer = 0
      }
    subrLayers = Vk.ImageSubresourceLayers
      { aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
      , mipLevel       = 0
      , baseArrayLayer = 0
      , layerCount     = 1
      }
    subrRange = subresource Vk.IMAGE_ASPECT_COLOR_BIT 1 1

gpuOnly :: VMA.AllocationCreateInfo
gpuOnly = zero
  { VMA.usage =
      VMA.MEMORY_USAGE_GPU_ONLY
  , VMA.requiredFlags =
      Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
  }

gpuToCpu :: VMA.AllocationCreateInfo
gpuToCpu = zero
  { VMA.flags =
      VMA.ALLOCATION_CREATE_MAPPED_BIT
  , VMA.usage =
      VMA.MEMORY_USAGE_GPU_TO_CPU
  , VMA.requiredFlags =
      Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT .|.
      Vk.MEMORY_PROPERTY_HOST_VISIBLE_BIT
  }

{-# 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