{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
module Devices where

import Control.Monad
import Control.Monad.Trans.Resource
import qualified Data.Vector as V
import Foreign.C.Types (CInt)
import Linear
import qualified Vulkan as Vk
import qualified Vulkan.Zero as Vk
import qualified Vulkan.CStruct.Extends as Vk

pickPhysicalDevice
  :: (MonadResource m)
  => Vk.Instance
  -> m Vk.PhysicalDevice
pickPhysicalDevice vkInstance = do
  physDevices <- snd <$> Vk.enumeratePhysicalDevices vkInstance

  devicesWithPropsAndFeatures <- V.mapM
    (\device -> do
      devProps <- Vk.getPhysicalDeviceProperties device
      devFeatures <- Vk.getPhysicalDeviceFeatures device
      return (device, devProps, devFeatures)
    )
    physDevices

  let discretePhysDevices = V.filter
        (\(_, _, devFeatures) ->
          -- Vk.deviceType devProp == Vk.PHYSICAL_DEVICE_TYPE_DISCRETE_GPU &&
          Vk.geometryShader devFeatures
          )
        devicesWithPropsAndFeatures
      maxMemdiscretePhysDevice = (\(x, _, _) -> x) $ V.foldl
        (\acc@(_, devProp1, _) ndev@(_, devProp2, _) ->
          let dimension = Vk.maxImageDimension2D . Vk.limits
          in
            if dimension devProp1 < dimension devProp2
            then ndev
            else acc
          )
        (V.head discretePhysDevices)
        discretePhysDevices

  return maxMemdiscretePhysDevice

createLogicalDevice
  :: (MonadResource m)
  => Vk.PhysicalDevice
  -> m Vk.Device
createLogicalDevice physDevice = do
  let priorities = V.singleton 1
      extensionNames = V.singleton Vk.KHR_SWAPCHAIN_EXTENSION_NAME
      queueCreateInfo = Vk.zero
        { Vk.queueFamilyIndex = 0
        , Vk.queuePriorities = priorities
        }
      deviceCreateInfo = Vk.zero
        { Vk.queueCreateInfos = V.singleton (Vk.SomeStruct queueCreateInfo)
        , Vk.enabledExtensionNames = extensionNames
        }

  (_, logDevice) <- allocate
    (Vk.createDevice physDevice deviceCreateInfo Nothing)
    (\logDevice -> do
      print ("destroying logical device" :: String)
      Vk.destroyDevice logDevice Nothing
      )

  return logDevice

createSwapchain
  :: (MonadResource m)
  => Vk.SurfaceKHR
  -> V2 CInt
  -> Vk.Device
  -> m Vk.SwapchainKHR
createSwapchain surface windowDimension logicalDevice = do
  let createInfo = Vk.zero
        { Vk.surface = surface
        , Vk.minImageCount = 4
        , Vk.imageFormat = Vk.FORMAT_B8G8R8A8_SRGB
        , Vk.imageColorSpace = Vk.COLOR_SPACE_SRGB_NONLINEAR_KHR
        , Vk.imageExtent =
            (\(V2 w h) -> Vk.Extent2D (fromIntegral w) (fromIntegral h)) windowDimension
        , Vk.imageArrayLayers = 1
        , Vk.imageUsage = Vk.IMAGE_USAGE_COLOR_ATTACHMENT_BIT
        , Vk.presentMode = Vk.PRESENT_MODE_MAILBOX_KHR
        , Vk.preTransform = Vk.SURFACE_TRANSFORM_IDENTITY_BIT_KHR
        , Vk.compositeAlpha = Vk.COMPOSITE_ALPHA_INHERIT_BIT_KHR
        }

  snd <$> allocate
    (Vk.createSwapchainKHR logicalDevice createInfo Nothing)
    (\swapchain -> do
      print ("destroying Swapchain" :: String)
      Vk.destroySwapchainKHR logicalDevice swapchain Nothing
      )

getImageHandles
  :: (MonadResource m)
  => Vk.SwapchainKHR
  -> Vk.Device
  -> m (V.Vector Vk.Image)
getImageHandles swapchain logicalDevice = do
  (result, handles) <- Vk.getSwapchainImagesKHR logicalDevice swapchain
  when (result /= Vk.SUCCESS) $
    error "getImageHandles: Failed acquiring images from swapchain"
  return handles