{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedRecordDot #-}
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 Foreign
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
import qualified VulkanMemoryAllocator as VMA

-- internal imports

import Image
import Types

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 devProps == 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, Word32)
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
        , Vk.enabledFeatures       = Just (Vk.zero
          { Vk.multiDrawIndirect = True
          } :: Vk.PhysicalDeviceFeatures)
        }
      shaderDrawParametersFeatures = Vk.zero
        { Vk.shaderDrawParameters  = True
        } :: Vk.PhysicalDeviceShaderDrawParametersFeatures

  (_, logDevice) <-
    Vk.withDevice
      physDevice
      (deviceCreateInfo Vk.::& shaderDrawParametersFeatures Vk.:& ())
      Nothing
      allocate

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

  return (logDevice, formats, graphicsQueueFamily)

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

  liftIO $ do
    putStrLn "available formats:"
    print surfaceFormats

  (surfaceFormat, colorSpace) <-
    if
      V.any
        (\a ->
          a.format ==
            Vk.FORMAT_B8G8R8A8_SRGB
          && a.colorSpace ==
            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_FIFO_KHR
        , Vk.preTransform     = Vk.SURFACE_TRANSFORM_IDENTITY_BIT_KHR
        , Vk.compositeAlpha   = Vk.COMPOSITE_ALPHA_OPAQUE_BIT_KHR
        , Vk.clipped          = True
        }

  swapchain <- snd <$>
    Vk.withSwapchainKHR logicalDevice createInfo Nothing allocate

  let depthImageExtent =
        (\(V2 w h) -> Vk.Extent3D (fromIntegral w) (fromIntegral h) 1)
        windowDimension

  let depthImageInfo = imageCreate
        Vk.FORMAT_D32_SFLOAT
        depthImageExtent
        Vk.IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT

  let depthImageAllocationInfo = Vk.zero
        { VMA.usage         = VMA.MEMORY_USAGE_GPU_ONLY
        , VMA.requiredFlags = Vk.MEMORY_PROPERTY_DEVICE_LOCAL_BIT
        }

  (depthImage, depthAllocation, _) <- snd <$>
    VMA.withImage allocator depthImageInfo depthImageAllocationInfo allocate

  let allocatedImage = AllocatedImage depthImage depthAllocation
      depthImageView = imageviewCreate
        (image allocatedImage)
        Vk.FORMAT_D32_SFLOAT
        Vk.IMAGE_ASPECT_DEPTH_BIT

  depthImageview <- snd <$>
    Vk.withImageView logicalDevice depthImageView Nothing allocate

  return (swapchain, Vk.SurfaceFormatKHR surfaceFormat colorSpace, depthImageview, allocatedImage)

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"

  liftIO $
    putStrLn ("number of images: " <> show (V.length handles))

  V.mapM
    (\tmpImage -> do
      let createInfo = Vk.zero
            { Vk.image            = tmpImage
            , Vk.viewType         = Vk.IMAGE_VIEW_TYPE_2D
            , Vk.format           = surfaceFormat.format
            , 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 <$>
        Vk.withImageView logicalDevice createInfo Nothing allocate
      )
    handles
