{-# LANGUAGE OverloadedRecordDot #-}

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.Core10.CommandBuffer qualified as CommandBuffer
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 pools pickQueue action =
  Vk.withCommandBuffers device commandBufferAllocateInfo bracket \case
    (Vector.toList -> [buf]) -> do
      let
        oneTime = zero
          { CommandBuffer.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 queue (wrap buf) fence
        Vk.waitForFences device (pure fence) True maxBound >>= \case
          Vk.SUCCESS ->
            -- traceM "oneshot_: queue finished"
            pure ()
          err ->
            error $ "oneshot_ failed: " <> show err
    _ ->
      error "assert: exactly the requested buffer was given"
  where
    device = getDevice context
    queue = snd . pickQueue $ getQueues context

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

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