module Resource.Buffer where

import RIO

import Data.Bits ((.|.))
import Data.Type.Equality (type (~))
import Data.Vector.Storable qualified as VectorS
import Foreign (Storable)
import Foreign qualified
import UnliftIO.Resource (MonadResource)
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), MonadVulkan)
import Engine.Worker qualified as Worker
import Resource.CommandBuffer (oneshot_)
import Resource.Vulkan.Named qualified as Named

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
  , aLabel          :: Maybe Text
  } deriving (Show)

instance Vk.HasObjectType (Allocated s a) where
  objectTypeAndHandle =
    Vk.objectTypeAndHandle . aBuffer

allocateCoherent
  :: ( MonadVulkan env m
     , Resource.MonadResource m
     , Storable a
     )
  => Maybe Text
  -> Vk.BufferUsageFlagBits
  -> "initial size" ::: Int
  -> VectorS.Vector a
  -> m (Resource.ReleaseKey, Allocated 'Coherent a)
allocateCoherent label usage initialSize xs = do
  res <- createCoherent label usage initialSize xs
  destroyIO <- toIO do
    context <- ask
    destroy context res
  key <- Resource.register destroyIO
  pure (key, res)

createCoherent
  :: forall a env m
  .  ( MonadVulkan env m
     , Storable a
     )
  => Maybe Text
  -> Vk.BufferUsageFlagBits
  -> "initial size" ::: Int
  -> VectorS.Vector a
  -> m (Allocated 'Coherent a)
createCoherent aLabel usage initialSize xs = do
  allocator <- asks getAllocator
  (aBuffer, aAllocation, aAllocationInfo) <- VMA.createBuffer allocator bci aci
  traverse_ (Named.object aBuffer) aLabel

  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 env m
  .  ( Storable a
     , MonadVulkan env m
     )
  => Maybe Text
  -> Queues Vk.CommandPool
  -> Vk.BufferUsageFlagBits
  -> Int
  -> VectorS.Vector a
  -> m (Allocated 'Staged a)
createStaged aLabel commandQueues usage initialSize xs = do
  allocator <- asks getAllocator
  (aBuffer, aAllocation, aAllocationInfo) <- VMA.createBuffer allocator bci aci
  traverse_ (Named.object aBuffer) aLabel

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

  pure Allocated
    { aCapacity = max initialSize len
    , aUsed     = fromIntegral len
    , aUsage    = usage
    , aLabel    = Nothing
    , ..
    }
  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       = 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
      }

register
  :: ( MonadVulkan env m
     , MonadResource m
     )
  => Allocated stage a
  -> m Resource.ReleaseKey
register Allocated{aBuffer, aAllocation} = do
  allocator <- asks getAllocator
  Resource.register $ VMA.destroyBuffer allocator aBuffer aAllocation

destroy
  :: (MonadUnliftIO io, HasVulkan context)
  => context
  -> Allocated s a
  -> io ()
destroy context 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)

{-# INLINE pokeCoherent #-}
pokeCoherent
  :: ( MonadVulkan env m
     , Storable a
     )
  => Allocated 'Coherent a
  -> Word32
  -> a
  -> m ()
pokeCoherent Allocated{..} ix new =
  case VMA.mappedData aAllocationInfo of
  _ | ix + 1 > aUsed ->
    pure ()
  ptr | ptr == Foreign.nullPtr ->
    pure ()
  ptr ->
    liftIO  $
      Foreign.pokeElemOff
        (Foreign.castPtr ptr)
        (fromIntegral ix)
        new

updateCoherent
  :: ( MonadUnliftIO io
     , Foreign.Storable a
     )
  => 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)

{-# INLINE updateCoherentResize_ #-}
updateCoherentResize_
  :: ( MonadVulkan env m
     , Storable a
     )
  => Allocated 'Coherent a
  -> VectorS.Vector a
  -> m (Allocated 'Coherent a)
updateCoherentResize_ old xs =
  if newSize <= oldSize then
    updateCoherent xs old
  else do
    ctx <- ask
    destroy ctx old
    createCoherent (aLabel old) (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)

type ObserverCoherent a = Worker.ObserverIO (Allocated 'Coherent a)

newObserverCoherent
  :: ( MonadVulkan env m
     , Storable a
     )
  => "label" ::: Text
  -> Vk.BufferUsageFlagBits
  -> Int
  -> VectorS.Vector a
  -> Resource.ResourceT m (ObserverCoherent a)
newObserverCoherent label usage initialCapacity initialData = do
  initialBuffer <- createCoherent (Just label) usage initialCapacity initialData
  observer <- Worker.newObserverIO initialBuffer

  context <- ask
  void $! Resource.register do
    currentBuffer <- Worker.readObservedIO observer
    destroy context currentBuffer

  pure observer

{-# INLINE observeCoherentResize_ #-}
observeCoherentResize_
  :: ( MonadVulkan env m
     , Worker.HasOutput source
     , Worker.GetOutput source ~ VectorS.Vector output
     , Storable output
     )
  => source
  -> ObserverCoherent output
  -> m ()
observeCoherentResize_ source observer = do
  Worker.observeIO_ source observer updateCoherentResize_

{-# INLINE observeCoherentSingle #-}
observeCoherentSingle
  :: ( MonadVulkan env m
     , Worker.HasOutput source
     , Worker.GetOutput source ~ output
     , Storable output
     )
  => source
  -> ObserverCoherent output
  -> m ()
observeCoherentSingle source observer =
  Worker.observeIO_ source observer \a b ->
    pokeCoherent a 0 b $> a