module Resource.CommandBuffer
  ( allocatePools
  , withPools
  , oneshot_
  ) where

import RIO

import Data.Vector qualified as Vector
import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.CStruct.Extends (SomeStruct(..))
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))
import Vulkan.Zero (zero)

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues(..))

allocatePools
  :: (HasVulkan context, Resource.MonadResource m)
  => context
  -> m (Resource.ReleaseKey, Queues Vk.CommandPool)
allocatePools context = do
  -- context <- ask
  bootstrapQueues <- for (getQueues context) \(QueueFamilyIndex ix, _queue) -> do
    let
      commandPoolCI = Vk.CommandPoolCreateInfo
        { flags            = zero
        , queueFamilyIndex = ix
        }
    Vk.withCommandPool (getDevice context) commandPoolCI Nothing Resource.allocate

  bqKey <- Resource.register $
    traverse_ (Resource.release . fst) bootstrapQueues
  pure (bqKey, fmap snd bootstrapQueues)

withPools
  :: (MonadVulkan env m, MonadResource m)
  => (Queues Vk.CommandPool -> m a)
  -> m a
withPools action = do
  context <- ask
  bracket (allocatePools context) (Resource.release . fst) (action . snd)

-- | Scratch command buffer for transfer operations.
-- The simple fence makes it unusable for rendering.
oneshot_
  :: (HasVulkan context, MonadUnliftIO m)
  => context
  -> Queues Vk.CommandPool
  -> (forall a . Queues a -> a)
  -> (Vk.CommandBuffer -> m ())
  -> m ()
oneshot_ context commandQueues pickQueue action =
  Vk.withCommandBuffers device commandBufferAllocateInfo bracket $ \case
    (Vector.toList -> [buf]) -> do
      let
        oneTime :: Vk.CommandBufferBeginInfo '[]
        oneTime = zero
          { Vk.flags = Vk.COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT
          }

      Vk.useCommandBuffer buf oneTime $ action buf

      Vk.withFence device zero Nothing bracket $ \fence -> do
        Vk.queueSubmit transferQueue (wrap buf) fence
        Vk.waitForFences device (pure fence) True maxBound >>= \case
          Vk.SUCCESS ->
            -- traceM "oneshot_: transfer finished"
            pure ()
          err ->
            error $ "copyBuffer failed: " <> show err
    _ ->
      error "assert: exactly the requested buffer was given"
  where
    device = getDevice context
    transferQueue = snd . pickQueue $ getQueues context

    commandBufferAllocateInfo :: Vk.CommandBufferAllocateInfo
    commandBufferAllocateInfo = zero
      { Vk.commandPool        = pickQueue commandQueues
      , Vk.level              = Vk.COMMAND_BUFFER_LEVEL_PRIMARY
      , Vk.commandBufferCount = 1
      }

    wrap buf = Vector.singleton $ SomeStruct zero
      { Vk.commandBuffers =
          Vector.singleton (Vk.commandBufferHandle buf)
      }