{-# LANGUAGE CPP #-}
{-# 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.Requirement (DeviceRequirement)
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 reqs opts pd deviceCI

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

  pure ld

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

reqs :: [DeviceRequirement]
reqs =
#ifdef darwin_HOST_OS
  [Utils.req| VK_KHR_portability_subset |] :
#endif
  [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
  |]

opts :: [DeviceRequirement]
opts =
  [Utils.reqs|
    PhysicalDeviceFeatures.samplerAnisotropy
    PhysicalDeviceFeatures.sampleRateShading
  |]

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