{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RecordWildCards #-}
module Main where

import Control.Monad
import Control.Monad.Loops
import Control.Monad.IO.Class
import Control.Monad.Trans.Resource
import qualified Data.ByteString as BS
import qualified Data.Vector as V
import Data.Word
import qualified SDL hiding (V2)
import qualified SDL.Video.Vulkan as SDL
import qualified Control.Concurrent.STM as STM
import Foreign.C.Types (CInt)
import Foreign.Ptr
import Linear
import qualified Vulkan as Vk
import qualified Vulkan.Zero as Vk
import qualified Vulkan.CStruct.Extends as Vk

myApiVersion :: Word32
myApiVersion = Vk.API_VERSION_1_3

main :: IO ()
main = runResourceT $ do
  (window, inst, surface, swapchain, images) <- initEngine

  SDL.showWindow window

  -- create abort condition for upcoming lop
  quit <- liftIO $ STM.newTMVarIO True

  -- main loop
  liftIO $ whileM_
    (STM.atomically $ STM.readTMVar quit
      )
    ( do
      -- poll events
      evs <- SDL.pollEvents
      -- flip abort condition on window close
      mapM_
        (\e -> case SDL.eventPayload e of
          SDL.WindowClosedEvent _ ->
            void $ STM.atomically $ STM.swapTMVar quit False
          _ -> return ()
          )
        evs
    )

initEngine
  :: MonadResource m
  => m (SDL.Window, Vk.Instance, Vk.SurfaceKHR, Vk.SwapchainKHR, V.Vector Vk.Image)
initEngine = do
  -- initialize SDL2 with all subsystems
  void $ allocate_
    SDL.initializeAll
    (do
      print ("This is the end!" :: String)
      SDL.quit
      )

  -- create a window configuration
  let windowConfig = SDL.defaultWindow
        { SDL.windowGraphicsContext = SDL.VulkanContext -- enable Vulkan support
        }

  -- load vulkan library
  SDL.vkLoadLibrary Nothing

  -- create the window
  (_, window) <-
    allocate
      (SDL.createWindow "Haskell ❤️ Vulkan" windowConfig)
      (\window -> do
        liftIO $ print ("destroying window" :: String)
        SDL.destroyWindow window
        )

  -- create vulkan surface and vulkan instance from window
  (surface, inst, swapchain, images) <- initVulkan window
  return (window, inst, surface, swapchain, images)

initVulkan
  :: MonadResource m
  => SDL.Window
  -> m (Vk.SurfaceKHR, Vk.Instance, Vk.SwapchainKHR, V.Vector Vk.Image)
initVulkan window = do

  vulkanInstance <- createInstance window

  vulkanPhysicalDevice <- pickPhysicalDevice vulkanInstance
  vulkanLogicalDevice <- createLogicalDevice vulkanPhysicalDevice

  (_, vulkanSurface) <- allocate
    (Vk.SurfaceKHR <$>
      SDL.vkCreateSurface window (castPtr (Vk.instanceHandle vulkanInstance))
      )
    (\vulkanSurface -> do
      print ("destroying Surface" :: String)
      Vk.destroySurfaceKHR vulkanInstance vulkanSurface Nothing
      )

  dimensions <- liftIO $ SDL.windowInitialSize <$> SDL.getWindowConfig window
  swapchain <- createSwapchain vulkanSurface dimensions vulkanLogicalDevice
  images <- getImageHandles swapchain vulkanLogicalDevice

  return (vulkanSurface, vulkanInstance, swapchain, images)

createInstance :: MonadResource m => SDL.Window -> m Vk.Instance
createInstance window = do
  windowExtensions <- liftIO $ fmap (Vk.EXT_DEBUG_UTILS_EXTENSION_NAME :) $
    traverse BS.packCString =<< SDL.vkGetInstanceExtensions window
  liftIO $ do
    print ("activating followinfg extensions:" :: String)
    print windowExtensions
  let createInfo = Vk.zero
        { Vk.applicationInfo = Just Vk.zero
          { Vk.applicationName = Nothing
          , Vk.apiVersion      = myApiVersion
          }
          -- enable debug util extension for debugging purposes
        , Vk.enabledExtensionNames = V.fromList windowExtensions
          -- enable validation layer for debugging purposes
        , Vk.enabledLayerNames = V.singleton "VK_LAYER_KHRONOS_validation"
        }

  (_, inst) <- allocate
    (Vk.createInstance createInfo Nothing)
    (\inst -> do
      print ("destroying instance" :: String)
      Vk.destroyInstance inst Nothing
      )

  return inst

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

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

  let discretePhysDevices = V.filter
        (\(_, devProp) -> Vk.deviceType devProp == Vk.PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
        devicesWithProps
      maxMemdiscretePhysDevice = fst $ 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