{-# LANGUAGE CPP #-}

module Engine.Setup where

import RIO

import Foreign.Ptr (castFunPtr)
import GHC.Clock (getMonotonicTimeNSec)
import UnliftIO.Resource (MonadResource)
import UnliftIO.Resource qualified as Resource
import Vulkan.Core10 qualified as Vk
import Vulkan.Core12 (pattern API_VERSION_1_2)
import Vulkan.Dynamic qualified as VkDynamic
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_portability_enumeration 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.Utils.Requirements (checkInstanceRequirements, requirementReport)
import Vulkan.Zero (zero)
import VulkanMemoryAllocator qualified as VMA

import Engine.Setup.Device (allocatePhysical, allocateLogical)
import Engine.Setup.Window qualified as Window
import Engine.StageSwitch (newStageSwitchVar)
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

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"

  let
    iReqs =
#ifdef darwin_HOST_OS
      portabilityEnum :
#endif
      deviceProps : debugUtils : windowReqs
    oReqs = []
    appInfo = Vk.ApplicationInfo
      { apiVersion = API_VERSION_1_2
      , applicationName = Nothing
      , applicationVersion = 0
      , engineName = Nothing
      , engineVersion = 0
      }
    instanceCI = zero
      { Vk.applicationInfo = Just appInfo
      , Vk.flags =
#ifdef darwin_HOST_OS
          Vk.INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR
#else
          zero
#endif
      }
  (instanceCI', reqResult, optResult) <- checkInstanceRequirements iReqs oReqs instanceCI
  let report = fmap fromString $ requirementReport reqResult optResult
  ghInstance <- case instanceCI' of
    Nothing -> do
      logWarn "Instance check failed"
      traverse_ logWarn report
      createInstanceFromRequirements iReqs oReqs instanceCI
    Just ci -> do
      traverse_ logDebug report
      createInstanceFromRequirements iReqs oReqs ci

  setupWith ghOptions ghWindow ghInstance

setupWith
  :: ( HasLogFunc env
     , MonadResource (RIO env)
     )
  => Options
  -> Window.Window
  -> Vk.Instance
  -> RIO env (GlobalHandles, Maybe SwapchainResources)
setupWith ghOptions ghWindow ghInstance = do
  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
  size <- liftIO (Window.getSize ghWindow)
  ghWindowSize <- Worker.newVar size

  ghStageSwitch <- newStageSwitchVar

  ghMonotonicStart <- liftIO getMonotonicTimeNSec

  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

  let
    iReqs =
#ifdef darwin_HOST_OS
      portabilityEnum :
#endif
      deviceProps : debugUtils : headlessReqs
    oReqs = []
    appInfo = Vk.ApplicationInfo
      { apiVersion = API_VERSION_1_2
      , applicationName = Nothing
      , applicationVersion = 0
      , engineName = Nothing
      , engineVersion = 0
      }
    instanceCI = zero
      { Vk.applicationInfo = Just appInfo
      , Vk.flags =
#ifdef darwin_HOST_OS
          Vk.INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR
#else
          zero
#endif
      }

  (instanceCI', reqResult, optResult) <- checkInstanceRequirements iReqs oReqs instanceCI
  let report = fmap fromString $ requirementReport reqResult optResult
  hInstance <- case instanceCI' of
    Nothing -> do
      logWarn "Instance check failed"
      traverse_ logWarn report
      createInstanceFromRequirements iReqs oReqs instanceCI
    Just ci -> do
      traverse_ logDebug report
      createInstanceFromRequirements iReqs oReqs ci

  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

portabilityEnum :: InstanceRequirement
portabilityEnum = RequireInstanceExtension
  Nothing
  Khr.KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME
  minBound

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
  ]

-- deriving instance Show InstanceRequirement