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

import Control.Monad
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Resource
import qualified Data.Vector as V
import Data.Maybe (fromMaybe, isNothing)
import Data.Bits
import Data.Word (Word32)
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, props,_ ) = 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

  liftIO $ do
    putStrLn "picked the following device:"
    print (Vk.deviceName props)

  return maxMemdiscretePhysDevice

getQueueFamily
  :: (MonadResource m)
  => Vk.PhysicalDevice
  -> Vk.QueueFlagBits
  -> m Word32
getQueueFamily physicalDevice distinctBit = do

  queueFamilyProperties <- Vk.getPhysicalDeviceQueueFamilyProperties physicalDevice

  liftIO $ do
    putStrLn "There are currently following queue families available"
    print queueFamilyProperties

  let queueFamily = V.foldl
        (\acc (index, qFamily) ->
          if Vk.queueFlags qFamily .&. distinctBit /= zeroBits && isNothing acc
          then Just index
          else acc
          )
        Nothing
        (V.zip (V.fromList [0..]) queueFamilyProperties)

  return $ fromMaybe
    (error "getQueueFamilies: The requested queueFamily doesn't seem to exist")
    queueFamily

createLogicalDevice
  :: (MonadResource m)
  => Vk.PhysicalDevice
  -> Vk.SurfaceKHR
  -> m (Vk.Device, V.Vector Vk.SurfaceFormatKHR)
createLogicalDevice physDevice surface = do

  graphicsQueueFamily <- getQueueFamily physDevice Vk.QUEUE_GRAPHICS_BIT

  let priorities = V.singleton 1
      extensionNames = V.singleton Vk.KHR_SWAPCHAIN_EXTENSION_NAME
      queueCreateInfo = Vk.zero
        { Vk.queueFamilyIndex = graphicsQueueFamily
        , 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
      putStrLn "destroying logical device"
      Vk.destroyDevice logDevice Nothing
      )

  (result, formats) <- Vk.getPhysicalDeviceSurfaceFormatsKHR physDevice surface
  unless (result == Vk.SUCCESS) $
    error "createLogicalDevice: Failed retrieving surface image formats"

  return (logDevice, formats)

createSwapchain
  :: (MonadResource m)
  => Vk.SurfaceKHR
  -> V.Vector Vk.SurfaceFormatKHR
  -> V2 CInt
  -> Vk.Device
  -> m (Vk.SwapchainKHR, Vk.SurfaceFormatKHR)
createSwapchain surface surfaceFormats windowDimension logicalDevice = do

  liftIO $ do
    putStrLn "available formats:"
    print surfaceFormats

  (surfaceFormat, colorSpace) <-
    if
      V.any
        (\a ->
          (Vk.format :: Vk.SurfaceFormatKHR -> Vk.Format) a ==
            Vk.FORMAT_B8G8R8A8_SRGB
          && (Vk.colorSpace :: Vk.SurfaceFormatKHR -> Vk.ColorSpaceKHR) a ==
            Vk.COLOR_SPACE_SRGB_NONLINEAR_KHR)
        surfaceFormats
    then
      return (Vk.FORMAT_B8G8R8A8_SRGB, Vk.COLOR_SPACE_SRGB_NONLINEAR_KHR)
    else
      error ("createSwapchain: No suitable surface formats available" :: String)

  liftIO $ do
    putStrLn "picking following format and color space:"
    print (surfaceFormat, colorSpace)

  let createInfo = Vk.zero
        { Vk.surface          = surface
        , Vk.minImageCount    = 4
        , Vk.imageFormat      = surfaceFormat
        , Vk.imageColorSpace  = colorSpace
        , 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_OPAQUE_BIT_KHR
        , Vk.clipped          = True
        }

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

  return (swapchain, Vk.SurfaceFormatKHR surfaceFormat colorSpace)

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

  V.mapM
    (\image -> do
      let createInfo = Vk.zero
            { Vk.image            = image
            , Vk.viewType         = Vk.IMAGE_VIEW_TYPE_2D
            , Vk.format           = (Vk.format :: Vk.SurfaceFormatKHR -> Vk.Format) surfaceFormat
            , Vk.components       = Vk.ComponentMapping
              { Vk.r = Vk.COMPONENT_SWIZZLE_IDENTITY
              , Vk.g = Vk.COMPONENT_SWIZZLE_IDENTITY
              , Vk.b = Vk.COMPONENT_SWIZZLE_IDENTITY
              , Vk.a = Vk.COMPONENT_SWIZZLE_IDENTITY
              }
            , Vk.subresourceRange = Vk.ImageSubresourceRange
              { aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
              , baseMipLevel   = 0
              , levelCount     = 1
              , baseArrayLayer = 0
              , layerCount     = 1
              }
            }
      snd <$> allocate
        (Vk.createImageView logicalDevice createInfo Nothing)
        (\imageView -> do
          putStrLn "destroying image view"
          Vk.destroyImageView logicalDevice imageView Nothing
          )
      )
    handles