module Vulkan.Utils.QueueAssignment
  ( assignQueues
  , QueueSpec(..)
  , QueueFamilyIndex(..)
  , QueueIndex(..)
  -- * Queue Family Predicates
  , isComputeQueueFamily
  , isGraphicsQueueFamily
  , isTransferQueueFamily
  , isTransferOnlyQueueFamily
  , isPresentQueueFamily
  ) where

import           Control.Applicative
import           Control.Category               ( (>>>) )
import           Control.Monad                  ( filterM )
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class      ( MonadTrans(lift) )
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State.Strict
                                                ( evalState
                                                , evalStateT
                                                , get
                                                , put
                                                )
import           Data.Bits
import           Data.Foldable
import           Data.Functor                   ( (<&>) )
import           Data.Traversable
import qualified Data.Vector                   as V
import           Data.Vector                    ( Vector )
import           Data.Word
import           GHC.Stack                      ( HasCallStack )
import           Vulkan.Core10
import           Vulkan.Extensions.VK_KHR_surface
                                                ( SurfaceKHR
                                                , getPhysicalDeviceSurfaceSupportKHR
                                                )
import           Vulkan.Utils.Misc
import           Vulkan.Zero

----------------------------------------------------------------
-- Device Queue creation
----------------------------------------------------------------

-- | Requirements for a 'Queue' to be assigned a family by 'assignQueues'.
--
-- To assign to a specific queue family index @f@:
--
-- @
-- queueSpecFamilyPredicate = \i _ -> i == f
-- @
--
-- To assign to any queue family which supports compute operations:
--
-- @
-- let isComputeQueue q = QUEUE_COMPUTE_BIT .&&. queueFlags q
-- in QueueSpec priority (\_index q -> pure (isComputeQueue q))
-- @
data QueueSpec m = QueueSpec
  { queueSpecQueuePriority :: Float
  , queueSpecFamilyPredicate
      :: QueueFamilyIndex -> QueueFamilyProperties -> m Bool
  }

newtype QueueFamilyIndex = QueueFamilyIndex { unQueueFamilyIndex :: Word32 }
  deriving (Eq, Ord, Enum, Show)

newtype QueueIndex = QueueIndex { unQueueIndex :: Word32 }
  deriving (Eq, Ord, Enum, Show)

-- | Given a 'PhysicalDevice' and a set of requirements for queues, calculate an
-- assignment of queues to queue families and return information with which to
-- create a 'Device' and also a function to extract the requested 'Queue's from
-- the device.
--
-- You may want to create a custom type with a 'Traversable' instance to store
-- your queues like:
--
-- @
-- data MyQueues q = MyQueues
--   { computeQueue            :: q
--   , graphicsAndPresentQueue :: q
--   , transferQueue           :: q
--   }
--
-- myQueueSpecs :: MyQueues QueueSpec
-- myQueueSpecs = MyQueues
--   { computeQueue            = QueueSpec 0.5 isComputeQueueFamily
--   , graphicsAndPresentQueue = QueueSpec 1   isPresentQueueFamily
--   , transferQueue           = QueueSpec 1   isTransferOnlyQueueFamily
--   }
-- @
--
-- Note, this doesn't permit differentiating queue family assignment based on
-- whether or not the queue is protected.
assignQueues
  :: forall f m n
   . (Traversable f, MonadIO m, MonadIO n)
  => PhysicalDevice
  -> f (QueueSpec m)
  -- ^ A set of requirements for 'Queue's to be created
  -> m
       ( Maybe
           ( Vector (DeviceQueueCreateInfo '[])
           , Device -> n (f (QueueFamilyIndex, Queue))
           )
       )
  -- ^
  -- - A set of 'DeviceQueueCreateInfo's to pass to 'createDevice'
  -- - A function to extract the requested 'Queue's from the 'Device' created
  --   with the 'DeviceQueueCreateInfo's
  --
  -- 'Nothing' if it wasn't possible to satisfy all the 'QueueSpec's
assignQueues phys specs = runMaybeT $ do
  queueFamilyProperties <-
    zip [QueueFamilyIndex 0 ..]
    .   V.toList
    <$> getPhysicalDeviceQueueFamilyProperties phys

  -- For each QueueSpec find the list of applicable families
  specsWithFamilies <- for specs $ \spec -> do
    families <- filterM (lift . uncurry (queueSpecFamilyPredicate spec))
                        queueFamilyProperties
    pure (spec, fst <$> families)

  let -- Get the number of available queues for each family
      familiesWithCapacities :: [(QueueFamilyIndex, Word32)]
      familiesWithCapacities =
        [ (i, queueCount)
        | (i, QueueFamilyProperties {..}) <- queueFamilyProperties
        ]

  -- Assign each QueueSpec to a queue family
  specsWithFamily :: f (QueueSpec m, QueueFamilyIndex) <- headMay
    (assign
      familiesWithCapacities
      (specsWithFamilies <&> \(spec, indices) index ->
        if index `elem` indices then Just (spec, index) else Nothing
      )
    )

  let maxFamilyIndex :: Maybe QueueFamilyIndex
      maxFamilyIndex = maximumMay (snd <$> toList specsWithFamily)

      -- Assign each QueueSpec an index within its queue family
      specsWithQueueIndex :: f (QueueSpec m, QueueFamilyIndex, QueueIndex)
      specsWithQueueIndex =
        flip evalState (repeat (QueueIndex 0))
          $ for specsWithFamily
          $ \(spec, familyIndex) -> do
              indices <- get
              let (index, indices') =
                    incrementAt (unQueueFamilyIndex familyIndex) indices
              put indices'
              pure (spec, familyIndex, index)

      -- Gather the priorities for each queue in each queue family
      queuePriorities :: [[Float]]
      queuePriorities = foldr
        (\(QueueSpec {..}, QueueFamilyIndex i) ps ->
          prependAt i queueSpecQueuePriority ps
        )
        (replicate
          (maybe 0 (fromIntegral . unQueueFamilyIndex . succ) maxFamilyIndex)
          []
        )
        specsWithFamily

      -- Make 'DeviceQueueCreateInfo's for the required queue families and
      -- priorities.
      queueCreateInfos :: Vector (DeviceQueueCreateInfo '[])
      queueCreateInfos = V.fromList
        [ zero { queueFamilyIndex = familyIndex
               , queuePriorities  = V.fromList ps
               }
        | (familyIndex, ps) <- zip [0 ..] queuePriorities
        , not (null ps)
        ]

      -- Get
      extractQueues :: Device -> n (f (QueueFamilyIndex, Queue))
      extractQueues dev =
        for specsWithQueueIndex
          $ \(_, i@(QueueFamilyIndex familyIndex), QueueIndex index) ->
              (i, ) <$> getDeviceQueue dev familyIndex index

  pure (queueCreateInfos, extractQueues)

----------------------------------------------------------------
-- Queue Predicates
----------------------------------------------------------------

isComputeQueueFamily :: QueueFamilyProperties -> Bool
isComputeQueueFamily q = QUEUE_COMPUTE_BIT .&&. queueFlags q

isGraphicsQueueFamily :: QueueFamilyProperties -> Bool
isGraphicsQueueFamily q = QUEUE_GRAPHICS_BIT .&&. queueFlags q

isTransferQueueFamily :: QueueFamilyProperties -> Bool
isTransferQueueFamily q = QUEUE_TRANSFER_BIT .&&. queueFlags q

-- | Does this queue have 'QUEUE_TRANSFER_BIT' set and not 'QUEUE_COMPUTE_BIT'
-- or 'QUEUE_GRAPHICS_BIT'
isTransferOnlyQueueFamily :: QueueFamilyProperties -> Bool
isTransferOnlyQueueFamily q =
  (   queueFlags q
    .&. (QUEUE_TRANSFER_BIT .|. QUEUE_GRAPHICS_BIT .|. QUEUE_COMPUTE_BIT)
    )
    == QUEUE_TRANSFER_BIT

-- | Can this queue family present to this surface on this device
isPresentQueueFamily
  :: MonadIO m => PhysicalDevice -> SurfaceKHR -> QueueFamilyIndex -> m Bool
isPresentQueueFamily phys surf (QueueFamilyIndex i) =
  getPhysicalDeviceSurfaceSupportKHR phys i surf

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

-- | Find all possible valid assignments for elements of a 'Traversable' with
-- some limited resources.
--
-- >>> assign @[] @_ @() [("a", 1)] []
-- [[]]
--
-- >>> assign @[] [("hi", 1), ("foo", 3)] [Just, Just . reverse, Just . take 1 ]
-- [["hi","oof","f"],["foo","ih","f"],["foo","oof","h"],["foo","oof","f"]]
--
-- >>> assign @[] [("a", 1), ("b", 2)] [\case {"a" -> Just 1; "b" -> Just 2; _ -> Nothing}, \case {"b" -> Just 3; _ -> Nothing}, \case {"a" -> Just 4; _ -> Nothing}]
-- [[2,3,4]]
assign
  :: forall f a b
   . Traversable f
  => [(a, Word32)]
  -- ^ How many of each 'a' are available
  -> f (a -> Maybe b)
  -- ^ Which 'a's can each element use
  -> [f b]
  -- ^ A list of assignments, each element in this list has the length of the
  -- requirements list
assign capacities = flip evalStateT capacities . traverse
  (\p -> do
    cs            <- get
    (choice, cs') <- lift (select p cs)
    put cs'
    pure choice
  )

-- | Select an element from the list according to some predicate, and return
-- that element along with the decremented list.
select :: (a -> Maybe b) -> [(a, Word32)] -> [(b, [(a, Word32)])]
select p = \case
  [] -> []
  x : xs ->
    let hit b = (b, if snd x == 1 then xs else (pred <$> x) : xs)
        miss = do
          (selected, xs') <- select p xs
          pure (selected, x : xs')
    in  if snd x == 0
          then miss
          else case p (fst x) of
            Nothing -> miss
            Just b  -> hit b : miss

headMay :: Alternative f => [a] -> f a
headMay = \case
  []    -> empty
  x : _ -> pure x

maximumMay :: (Foldable f, Ord a) => f a -> Maybe a
maximumMay f = if null f then Nothing else Just (maximum f)

incrementAt :: (HasCallStack, Enum a) => Word32 -> [a] -> (a, [a])
incrementAt index = modAt index succ

prependAt :: HasCallStack => Word32 -> a -> [[a]] -> [[a]]
prependAt index p = snd . modAt index (p :)

modAt :: HasCallStack => Word32 -> (a -> a) -> [a] -> (a, [a])
modAt index f = splitAt (fromIntegral index) >>> \case
  (_ , []    ) -> error "modAt, out of bounds"
  (xs, y : ys) -> (y, xs <> (f y : ys))