{-# LANGUAGE CPP #-}

module Engine.Setup where

import RIO

import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Extensions.VK_EXT_debug_utils qualified as Ext
import Vulkan.Extensions.VK_KHR_get_physical_device_properties2 qualified as Khr
import Vulkan.Extensions.VK_KHR_surface qualified as Khr
import Vulkan.Requirement (InstanceRequirement(..))
import Vulkan.Utils.Initialization (createInstanceFromRequirements)
import Vulkan.Utils.QueueAssignment (QueueFamilyIndex(..))
import Vulkan.Zero (zero)
import VulkanMemoryAllocator qualified as VMA
import Vulkan.Core12 (pattern API_VERSION_1_2)

#if MIN_VERSION_vulkan(3,15,0)
import Foreign.Ptr (castFunPtr)
import Vulkan.Dynamic qualified as VkDynamic
#endif

import Engine.Setup.Device (allocatePhysical, allocateLogical)
import Engine.Setup.Window qualified as Window
import Engine.Types (GlobalHandles(..))
import Engine.Types.Options (Options(..))
import Engine.Vulkan.Swapchain (SwapchainResources)
import Engine.Vulkan.Types (PhysicalDeviceInfo(..))
import Engine.Vulkan.Types qualified as Vulkan
import Engine.Worker qualified as Worker
import Engine.StageSwitch (newStageSwitchVar)

setup
  :: ( HasLogFunc env
     , MonadResource (RIO env)
     )
  => Options -> RIO env (GlobalHandles, Maybe SwapchainResources)
setup ghOptions = do
  logDebug $ displayShow ghOptions

  (windowReqs, ghWindow) <- Window.allocate
    (optionsFullscreen ghOptions)
    (optionsSize ghOptions)
    (optionsDisplay ghOptions)
    Window.pickLargest
    "Keid Engine"

  logDebug "Creating instance"
  let
    appInfo = Vk.ApplicationInfo
      { apiVersion = API_VERSION_1_2
      , applicationName = Nothing
      , applicationVersion = 0
      , engineName = Nothing
      , engineVersion = 0
      }
  ghInstance <- createInstanceFromRequirements
    (deviceProps : debugUtils : windowReqs)
    mempty
    (zero { Vk.applicationInfo = Just appInfo })

  logDebug "Creating surface"
  (_surfaceKey, ghSurface) <- Window.allocateSurface ghWindow ghInstance

  logDebug "Creating physical device"
  (ghPhysicalDeviceInfo, ghPhysicalDevice) <- allocatePhysical
    ghInstance
    (Just ghSurface)
    pdiTotalMemory

  logDebug "Creating logical device"
  ghDevice <- allocateLogical ghPhysicalDeviceInfo ghPhysicalDevice

  logDebug "Creating VMA"
  let
    allocatorCI :: VMA.AllocatorCreateInfo
    allocatorCI = zero
      { VMA.physicalDevice  = Vk.physicalDeviceHandle ghPhysicalDevice
      , VMA.device          = Vk.deviceHandle ghDevice
      , VMA.instance'       = Vk.instanceHandle ghInstance
      , VMA.vulkanFunctions = Just $ vmaVulkanFunctions ghDevice ghInstance
      }
  (_vmaKey, ghAllocator) <- VMA.withAllocator allocatorCI Resource.allocate
  toIO (logDebug "Releasing VMA") >>= Resource.register

  ghQueues <- liftIO $ pdiGetQueues ghPhysicalDeviceInfo ghDevice
  logDebug $ "Got command queues: " <> displayShow (fmap (unQueueFamilyIndex . fst) ghQueues)

  screen <- liftIO $ Window.getExtent2D ghWindow
  ghScreenVar <- Worker.newVar screen

  ghStageSwitch <- newStageSwitchVar

  pure (GlobalHandles{..}, Nothing)

vmaVulkanFunctions
  :: Vk.Device
  -> Vk.Instance
  -> VMA.VulkanFunctions
#if MIN_VERSION_vulkan(3,15,0)
vmaVulkanFunctions Vk.Device{deviceCmds} Vk.Instance{instanceCmds} =
  zero
    { VMA.vkGetInstanceProcAddr =
        castFunPtr $ VkDynamic.pVkGetInstanceProcAddr instanceCmds
    , VMA.vkGetDeviceProcAddr =
        castFunPtr $ VkDynamic.pVkGetDeviceProcAddr deviceCmds
    }
#else
vmaVulkanFunctions _device _instance = zero
#endif

setupHeadless
  :: ( HasLogFunc env
     , MonadResource (RIO env)
     )
  => Options
  -> RIO env Headless
setupHeadless opts = do
  logDebug $ displayShow opts

  logDebug "Creating instance"
  hInstance <- createInstanceFromRequirements
    (deviceProps : debugUtils : headlessReqs)
    mempty
    zero

  logDebug "Creating physical device"
  (hPhysicalDeviceInfo, hPhysicalDevice) <- allocatePhysical
    hInstance
    Nothing
    pdiTotalMemory

  logDebug "Creating logical device"
  hDevice <- allocateLogical hPhysicalDeviceInfo hPhysicalDevice

  logDebug "Creating VMA"
  let
    allocatorCI :: VMA.AllocatorCreateInfo
    allocatorCI = zero
      { VMA.physicalDevice  = Vk.physicalDeviceHandle hPhysicalDevice
      , VMA.device          = Vk.deviceHandle hDevice
      , VMA.instance'       = Vk.instanceHandle hInstance
      , VMA.vulkanFunctions = Just $ vmaVulkanFunctions hDevice hInstance
      }
  (_vmaKey, hAllocator) <- VMA.withAllocator allocatorCI Resource.allocate
  toIO (logDebug "Releasing VMA") >>= Resource.register

  hQueues <- liftIO $ pdiGetQueues hPhysicalDeviceInfo hDevice
  logDebug $ "Got command queues: " <> displayShow (fmap (unQueueFamilyIndex . fst) hQueues)

  pure Headless{..}

data Headless = Headless
  { hInstance           :: Vk.Instance
  , hPhysicalDeviceInfo :: Vulkan.PhysicalDeviceInfo
  , hPhysicalDevice     :: Vk.PhysicalDevice
  , hDevice             :: Vk.Device
  , hAllocator          :: VMA.Allocator
  , hQueues             :: Vulkan.Queues (QueueFamilyIndex, Vk.Queue)
  }

instance Vulkan.HasVulkan Headless where
  getInstance           = hInstance
  getQueues             = hQueues
  getPhysicalDevice     = hPhysicalDevice
  getPhysicalDeviceInfo = hPhysicalDeviceInfo
  getDevice             = hDevice
  getAllocator          = hAllocator

deviceProps :: InstanceRequirement
deviceProps = RequireInstanceExtension
  Nothing
  Khr.KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME
  minBound

debugUtils :: InstanceRequirement
debugUtils = RequireInstanceExtension
  Nothing
  Ext.EXT_DEBUG_UTILS_EXTENSION_NAME
  minBound

headlessReqs :: [InstanceRequirement]
headlessReqs =
  [ RequireInstanceExtension
      Nothing
      Khr.KHR_SURFACE_EXTENSION_NAME
      minBound
  ]