module Resource.Buffer where

import RIO

import Data.Bits ((.|.))
import Data.Vector.Storable qualified as VectorS
import Foreign (Storable)
import Foreign qualified
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.NamedType ((:::))
import Vulkan.Zero (zero)
import VulkanMemoryAllocator qualified as VMA

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

data Store = Staged | Coherent

data Allocated (s :: Store) a = Allocated
  { aBuffer         :: Vk.Buffer
  , aAllocation     :: VMA.Allocation
  , aAllocationInfo :: VMA.AllocationInfo
  , aCapacity       :: Int
  , aUsed           :: Word32
  , aUsage          :: Vk.BufferUsageFlagBits
  } deriving (Show)

allocateCoherent
  :: (Resource.MonadResource m, Storable a, HasVulkan context)
  => context
  -> Vk.BufferUsageFlagBits
  -> "initial size" ::: Int
  -> VectorS.Vector a
  -> m (Resource.ReleaseKey, Allocated 'Coherent a)
allocateCoherent context usage initialSize xs =
  Resource.allocate
    (createCoherent context usage initialSize xs)
    (destroy context)

createCoherent
  :: forall a context io . (Storable a, HasVulkan context, MonadUnliftIO io)
  => context
  -> Vk.BufferUsageFlagBits
  -> "initial size" ::: Int
  -> VectorS.Vector a
  -> io (Allocated 'Coherent a)
createCoherent context usage initialSize xs = do
  (aBuffer, aAllocation, aAllocationInfo) <- VMA.createBuffer (getAllocator context) bci aci

  when (len /= 0) $
    if VMA.mappedData aAllocationInfo == Foreign.nullPtr then
      error "TODO: recover from unmapped data and flush manually"
    else
      liftIO $ VectorS.unsafeWith xs \src ->
        Foreign.copyBytes (Foreign.castPtr $ VMA.mappedData aAllocationInfo) src lenBytes

  pure Allocated
    { aCapacity = max initialSize len
    , aUsed     = fromIntegral len
    , aUsage    = usage
    , ..
    }
  where

    len = VectorS.length xs
    lenBytes = Foreign.sizeOf (undefined :: a) * len

    sizeBytes = Foreign.sizeOf (undefined :: a) * max initialSize len

    bci = zero
      { Vk.size        = fromIntegral sizeBytes
      , Vk.usage       = usage
      , Vk.sharingMode = Vk.SHARING_MODE_EXCLUSIVE
      }

    flags =
      Vk.MEMORY_PROPERTY_HOST_VISIBLE_BIT .|.
      Vk.MEMORY_PROPERTY_HOST_COHERENT_BIT

    aci = zero
      { VMA.flags          = VMA.ALLOCATION_CREATE_MAPPED_BIT
      , VMA.usage          = VMA.MEMORY_USAGE_GPU_ONLY
      , VMA.requiredFlags  = flags
      , VMA.preferredFlags = Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
      }

createStaged
  :: forall a context io . (Storable a, HasVulkan context, MonadUnliftIO io)
  => context
  -> Queues Vk.CommandPool
  -> Vk.BufferUsageFlagBits
  -> Int
  -> VectorS.Vector a
  -> io (Allocated 'Staged a)
createStaged context commandQueues usage initialSize xs = do
  (aBuffer, aAllocation, aAllocationInfo) <- VMA.createBuffer vma bci aci

  when (len /= 0) . liftIO $
    VMA.withBuffer vma stageCI stageAI bracket \(staging, _stage, stageInfo) ->
      if VMA.mappedData stageInfo == Foreign.nullPtr then
        error "TODO: recover from unmapped data and flush manually"
      else do
        VectorS.unsafeWith xs \src ->
          Foreign.copyBytes (Foreign.castPtr $ VMA.mappedData stageInfo) src lenBytes
        copyBuffer_ context commandQueues aBuffer staging (fromIntegral sizeBytes)

  pure Allocated
    { aCapacity = max initialSize len
    , aUsed     = fromIntegral len
    , aUsage    = usage
    , ..
    }
  where
    vma = getAllocator context

    len = VectorS.length xs
    lenBytes = Foreign.sizeOf (undefined :: a) * len

    sizeBytes = Foreign.sizeOf (undefined :: a) * max initialSize len

    bci = zero
      { Vk.size        = fromIntegral sizeBytes
      , Vk.usage       = Vk.BUFFER_USAGE_TRANSFER_DST_BIT .|. usage
      , Vk.sharingMode = Vk.SHARING_MODE_EXCLUSIVE
      }

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

    stageCI = zero
      { Vk.size        = fromIntegral sizeBytes
      , Vk.usage       = Vk.BUFFER_USAGE_TRANSFER_SRC_BIT
      , Vk.sharingMode = Vk.SHARING_MODE_EXCLUSIVE
      }

    stageAI = zero
      { VMA.flags         = VMA.ALLOCATION_CREATE_MAPPED_BIT
      , VMA.usage         = VMA.MEMORY_USAGE_CPU_TO_GPU
      , VMA.requiredFlags = Vk.MEMORY_PROPERTY_HOST_VISIBLE_BIT
      }

destroy
  :: (MonadUnliftIO io, HasVulkan context)
  => context
  -> Allocated s a
  -> io ()
destroy context a =
  VMA.destroyBuffer (getAllocator context) (aBuffer a) (aAllocation a)

destroyAll
  :: (MonadUnliftIO io, HasVulkan context, Foldable t)
  => context
  -> t (Allocated s a)
  -> io ()
destroyAll context =
  traverse_ \a ->
    VMA.destroyBuffer (getAllocator context) (aBuffer a) (aAllocation a)

peekCoherent :: (MonadIO m, Storable a) => Word32 -> Allocated 'Coherent a -> m (Maybe a)
peekCoherent ix Allocated{..} =
  case VMA.mappedData aAllocationInfo of
    _ | ix + 1 > aUsed ->
      pure Nothing
    ptr | ptr == Foreign.nullPtr ->
      pure Nothing
    ptr ->
      liftIO . fmap Just $
        Foreign.peekElemOff (Foreign.castPtr ptr) (fromIntegral ix)

updateCoherent
  :: (Foreign.Storable a, MonadUnliftIO io)
  => VectorS.Vector a
  -> Allocated 'Coherent a
  -> io (Allocated 'Coherent a)
updateCoherent xs old = do
  liftIO $ VectorS.unsafeWith xs \src ->
    Foreign.copyBytes dst src lenBytes
  pure old
    { aUsed = fromIntegral len
    }
  where
    dst = Foreign.castPtr $ VMA.mappedData (aAllocationInfo old)
    len = VectorS.length xs
    lenBytes = len * Foreign.sizeOf (VectorS.head xs)

updateCoherentResize_
  :: (Storable a, HasVulkan context, MonadUnliftIO io)
  => context
  -> Allocated 'Coherent a
  -> VectorS.Vector a
  -> io (Allocated 'Coherent a)
updateCoherentResize_ ctx old xs =
  if newSize <= oldSize then
    updateCoherent xs old
  else do
    destroyAll ctx [old]
    createCoherent ctx (aUsage old) newSize xs
  where
    oldSize = aCapacity old
    newSize = VectorS.length xs

-- TODO: add a staged buffer to check out the transfer queue
copyBuffer_
  :: (MonadUnliftIO io, HasVulkan context)
  => context
  -> Queues Vk.CommandPool
  -> ("dstBuffer" ::: Vk.Buffer)
  -> ("srcBuffer" ::: Vk.Buffer)
  -> Vk.DeviceSize
  -> io ()
copyBuffer_ context commandQueues dst src sizeBytes =
  oneshot_ context commandQueues qTransfer \cmd ->
    Vk.cmdCopyBuffer cmd src dst
      (pure $ Vk.BufferCopy 0 0 sizeBytes)