{-# LANGUAGE OverloadedRecordDot #-}

-- | Physical device tools

module Engine.Setup.Device where

import RIO

import Control.Monad.Trans.Maybe (MaybeT(..))
import GHC.IO.Exception (IOException(..), IOErrorType(NoSuchThing))
import RIO.Text qualified as Text
import RIO.Vector qualified as V
import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Core11.Promoted_From_VK_KHR_multiview (PhysicalDeviceMultiviewFeatures(..))
import Vulkan.Core12.Promoted_From_VK_EXT_descriptor_indexing (PhysicalDeviceDescriptorIndexingFeatures(..))
import Vulkan.Core12.Promoted_From_VK_KHR_timeline_semaphore (PhysicalDeviceTimelineSemaphoreFeatures(..))
import Vulkan.CStruct.Extends ( SomeStruct(SomeStruct), pattern (:&), pattern (::&))
import Vulkan.Extensions.VK_KHR_get_physical_device_properties2 (getPhysicalDeviceFeatures2KHR)
import Vulkan.Extensions.VK_KHR_surface qualified as Khr
import Vulkan.Extensions.VK_KHR_swapchain (pattern KHR_SWAPCHAIN_EXTENSION_NAME)
import Vulkan.Extensions.VK_KHR_timeline_semaphore (pattern KHR_TIMELINE_SEMAPHORE_EXTENSION_NAME)
import Vulkan.Utils.Initialization (createDeviceFromRequirements, physicalDeviceName, pickPhysicalDevice)
import Vulkan.Utils.QueueAssignment (QueueSpec(..))
import Vulkan.Utils.QueueAssignment qualified as Utils
import Vulkan.Utils.Requirements.TH qualified as Utils
import Vulkan.Core10 (PhysicalDeviceFeatures(..))
import Vulkan.Zero (zero)

import Engine.Vulkan.Types (PhysicalDeviceInfo(..), Queues(..))

allocatePhysical
  :: ( MonadUnliftIO m, MonadThrow m
     , MonadReader env m
     , HasLogFunc env
     , MonadResource m
     )
  => Vk.Instance
  -> Maybe Khr.SurfaceKHR
  -> (PhysicalDeviceInfo -> Word64)
  -> m (PhysicalDeviceInfo, Vk.PhysicalDevice)
allocatePhysical vkInstance presentSurface score = do
  UnliftIO unliftIO <- askUnliftIO

  let
    create = unliftIO do
      logDebug "Picking physical device..."
      pickPhysicalDevice vkInstance (physicalDeviceInfo presentSurface) score >>= \case
        Nothing ->
          noSuchThing "Unable to find appropriate PhysicalDevice"
        Just res@(pdi, _dev) -> do
          logInfo $ mconcat
            [ "Using physical device: "
            , displayShow (pdiName pdi)
            ]
          pure res

    destroy _res = unliftIO $
      logDebug "Destroying physical device"

  fmap snd $ Resource.allocate create destroy

physicalDeviceInfo
  :: ( MonadIO m
     , MonadReader env m
     , HasLogFunc env
     )
  => Maybe Khr.SurfaceKHR
  -> Vk.PhysicalDevice
  -> m (Maybe PhysicalDeviceInfo)
physicalDeviceInfo presentSurface phys = runMaybeT do
  pdiName <- physicalDeviceName phys

  let
    ignoreDevice =
      "llvmpipe" `Text.isPrefixOf` pdiName

  if ignoreDevice then do
    logDebug $ "Ignoring " <> displayShow pdiName
    mzero
  else
    logDebug $ "Considering " <> displayShow pdiName

  hasTimelineSemaphores <- deviceHasTimelineSemaphores phys
  unless hasTimelineSemaphores do
    logWarn $ mconcat
      [ "Not using physical device "
      , displayShow pdiName
      , " because it doesn't support timeline semaphores"
      ]
    mzero

  hasSwapchainSupport <- deviceHasSwapchain phys
  unless hasSwapchainSupport do
    logWarn $ mconcat
      [ "Not using physical device "
      , displayShow pdiName
      , " because it doesn't support swapchains"
      ]
    mzero

  assigned <- Utils.assignQueues phys (queueRequirements phys presentSurface)
  (pdiQueueCreateInfos, pdiGetQueues) <- case assigned of
    Nothing -> do
      logDebug "Queue assignment failed"
      fallback <- Utils.assignQueues @_ @_ @IO phys (Identity $ QueueSpec 1.0 isFallbackQ)
      case fallback of
        Nothing -> do
          logWarn "Fallback assignment failed too"
          mzero
        Just (infos, getQueues) -> do
          logDebug "Fallback assignment succeeded"
          pure
            ( infos
            , \dev -> do
                Identity q <- getQueues dev
                pure $ Queues q q q
            )

    Just queues ->
      pure queues

  pdiTotalMemory <- do
    props <- Vk.getPhysicalDeviceMemoryProperties phys
    pure . sum $
      fmap (.size) (Vk.memoryHeaps props)

  pdiProperties <- Vk.getPhysicalDeviceProperties phys

  pure PhysicalDeviceInfo{..}
  where
    isFallbackQ _queueFamilyIndex queueFamilyProperties =
      pure $ Utils.isGraphicsQueueFamily queueFamilyProperties

{- |
  Requirements for a 'Queue' which has graphics support and can present to
  the specified surface.

  Priorities are ranged 0.0 to 1.0 with higher number means higher priority.
  https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#devsandqueues-priority
-}
queueRequirements
  :: MonadIO m
  => Vk.PhysicalDevice
  -> Maybe Khr.SurfaceKHR
  -> Queues (QueueSpec m)
queueRequirements phys presentSurface = Queues
  { qGraphics = QueueSpec 1.0 isGraphicsPresentQueue
  , qCompute  = QueueSpec 0.5 isComputeQueue
  , qTransfer = QueueSpec 0.0 isTransferQueue
  }
 where
  isGraphicsPresentQueue queueFamilyIndex queueFamilyProperties =
    case presentSurface of
      Just surf -> do
        pq <- Utils.isPresentQueueFamily phys surf queueFamilyIndex
        pure $ pq && gq
      Nothing ->
        pure gq
    where
      gq = Utils.isGraphicsQueueFamily queueFamilyProperties

  isTransferQueue _queueFamilyIndex queueFamilyProperties =
    pure $ Utils.isTransferQueueFamily queueFamilyProperties

  isComputeQueue _queueFamilyIndex queueFamilyProperties =
    pure $ Utils.isComputeQueueFamily queueFamilyProperties

deviceHasSwapchain :: MonadIO m => Vk.PhysicalDevice -> m Bool
deviceHasSwapchain dev = do
  (_, extensions) <- Vk.enumerateDeviceExtensionProperties dev Nothing
  pure $ V.any
    ((KHR_SWAPCHAIN_EXTENSION_NAME ==) . Vk.extensionName)
    extensions

deviceHasTimelineSemaphores :: MonadIO m => Vk.PhysicalDevice -> m Bool
deviceHasTimelineSemaphores phys = do
  (_, extensions) <- Vk.enumerateDeviceExtensionProperties phys Nothing
  let
    hasExt = V.any
      ((KHR_TIMELINE_SEMAPHORE_EXTENSION_NAME ==) . Vk.extensionName)
      extensions

  hasFeat <- getPhysicalDeviceFeatures2KHR phys >>= \case
    _ ::& (PhysicalDeviceTimelineSemaphoreFeatures hasTimelineSemaphores :& ()) ->
      pure hasTimelineSemaphores

  pure $ hasExt && hasFeat

allocateLogical
  :: ( MonadUnliftIO m
     , MonadReader env m, HasLogFunc env
     , MonadResource m
     )
  => PhysicalDeviceInfo -> Vk.PhysicalDevice -> m Vk.Device
allocateLogical pdi pd = do
  logDebug "Creating logical device"

  ld <- createDeviceFromRequirements
    [Utils.reqs|
      1.2

      VK_KHR_maintenance3
      VK_KHR_swapchain

      -- PhysicalDeviceFeatures.robustBufferAccess
      PhysicalDeviceFeatures.textureCompressionBC

      VK_KHR_multiview
      PhysicalDeviceMultiviewFeatures.multiview

      VK_EXT_descriptor_indexing
      PhysicalDeviceDescriptorIndexingFeatures.descriptorBindingPartiallyBound
      PhysicalDeviceDescriptorIndexingFeatures.descriptorBindingVariableDescriptorCount
      PhysicalDeviceDescriptorIndexingFeatures.runtimeDescriptorArray
      PhysicalDeviceDescriptorIndexingFeatures.shaderSampledImageArrayNonUniformIndexing

      VK_KHR_timeline_semaphore
      PhysicalDeviceTimelineSemaphoreFeatures.timelineSemaphore
    |]
    [Utils.reqs|
      PhysicalDeviceFeatures.samplerAnisotropy
      PhysicalDeviceFeatures.sampleRateShading
    |]
    pd
    deviceCI

  toIO (logDebug "Destroying logical device") >>=
    Resource.register

  pure ld

  where
    deviceCI = zero
      { Vk.queueCreateInfos = fmap SomeStruct (pdiQueueCreateInfos pdi)
      }

noSuchThing :: MonadThrow m => String -> m a
noSuchThing message =
  throwM $
    IOError Nothing NoSuchThing "" message Nothing Nothing