{-# 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
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
        (\(_, devProps, 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)
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
  -> 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 ->
          (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_FIFO_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
      )

  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 <$> allocate
    (VMA.createImage allocator depthImageInfo depthImageAllocationInfo)
    (\(depthImage, depthAllocation, _) -> do
      putStrLn "destroying image"
      VMA.destroyImage allocator depthImage depthAllocation
      )

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

  depthImageview <- snd <$> allocate
    (Vk.createImageView logicalDevice depthImageView Nothing)
    (\depthImageview -> do
      putStrLn "destroying image view"
      Vk.destroyImageView logicalDevice depthImageview Nothing
      )

  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           = (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