{-# 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 Control.Monad.Trans.Maybe
import Data.Bits
import qualified Data.ByteString as BS
import Data.Maybe
import qualified Data.Text as T
import qualified Data.Vector as V
import Data.Word
import qualified SDL
import qualified SDL.Video.Vulkan as SDL
import qualified Control.Concurrent.STM as STM
import Foreign.Ptr
import qualified Vulkan as Vk
import qualified Vulkan.CStruct.Extends as Vk
import qualified Vulkan.Requirement as Vk
import qualified Vulkan.Utils.Initialization as Vk
import qualified Vulkan.Zero as Vk
import qualified VulkanMemoryAllocator as Vk

myApiVersion :: Word32
myApiVersion = Vk.API_VERSION_1_2

main :: IO ()
main = runResourceT $ do
  (window, inst, surface) <- 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
    )

  quitEngine window inst surface

initEngine :: MonadResource m => m (SDL.Window, Vk.Instance, Vk.SurfaceKHR)
initEngine = do
  -- initialize SDL2 with all subsystems
  SDL.initializeAll

  -- 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 <- SDL.createWindow "Haskell ❤️ Vulkan" windowConfig

  (surface, inst) <- initVulkan window

  return (window, inst, surface)

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

  vulkanInstance <- createInstance window
  -- vulkanPhysicalDevice <- pickPhysicalDevice vulkanInstance
  -- vulkanLogicalDevice <- createLogicalDevice vulkanPhysicalDevice

  vulkanSurface <- Vk.SurfaceKHR <$>
    SDL.vkCreateSurface window (castPtr (Vk.instanceHandle vulkanInstance))
  DeviceParams
    devName
    physicalDevice
    logicalDevice
    graphicsQueue
    graphicsQueueFamilyIndex <- createDevice vulkanInstance vulkanSurface

  let commandPoolCreateInfo = Vk.zero
        { Vk.queueFamilyIndex = graphicsQueueFamilyIndex
        } :: Vk.CommandPoolCreateInfo
  commandPools <-
    V.replicateM
      3
      (snd <$>
        Vk.withCommandPool logicalDevice commandPoolCreateInfo Nothing allocate
        )

  allocator <- createVMA vulkanInstance physicalDevice logicalDevice

  liftIO $ print ("Using device: " <> devName)

  return (vulkanSurface, vulkanInstance)


quitEngine
  :: MonadResource m
  => SDL.Window
  -> Vk.Instance
  -> Vk.SurfaceKHR
  -> m ()
quitEngine window inst surface = liftIO $ do
  Vk.destroySurfaceKHR inst surface Nothing

  Vk.destroyInstance inst Nothing

  -- destroy window
  SDL.destroyWindow window

  -- quit SDL
  SDL.quit

createInstance :: MonadResource m => SDL.Window -> m Vk.Instance
createInstance window = do
  windowExtensions <- liftIO $
    traverse BS.packCString =<< SDL.vkGetInstanceExtensions window
  let createInfo = Vk.zero
        { Vk.applicationInfo = Just Vk.zero
          { Vk.applicationName = Nothing
          , Vk.apiVersion      = myApiVersion
          }
        }
      requirements =
        [Vk.RequireInstanceExtension Nothing n minBound | n <- windowExtensions ]
  Vk.createDebugInstanceFromRequirements requirements [] createInfo

data DeviceParams = DeviceParams
  { dpDevicename               :: T.Text
  , dpPhysicalDevice           :: Vk.PhysicalDevice
  , dpDevice                   :: Vk.Device
  , dpGraphicsQueue            :: Vk.Queue
  , dpGraphicsQueueFamilyIndex :: Word32
  }

createDevice
  :: MonadResource m
  => Vk.Instance
  -> Vk.SurfaceKHR
  -> m DeviceParams
createDevice inst surface = do
  (pdi, physDevice) <-
    fromMaybe (error "Unable to find suitable physival device!") <$>
    Vk.pickPhysicalDevice inst (physicalDeviceInfo surface) id

  devName <- Vk.physicalDeviceName physDevice

  --
  -- get logical device
  --

  let graphicsQueueFamilyIndex = pdiGraphicsQueueFamilyIndex pdi
      deviceCreateInfo = Vk.zero
        { Vk.queueCreateInfos      =
          [ Vk.SomeStruct Vk.zero
            { Vk.queueFamilyIndex = graphicsQueueFamilyIndex
            , Vk.queuePriorities = [ 1 ]
            }
          ]
        , Vk.enabledExtensionNames = [ Vk.KHR_SWAPCHAIN_EXTENSION_NAME ]
        }
  (_, device) <- Vk.withDevice physDevice deviceCreateInfo Nothing allocate
  graphicsQueue <-
    Vk.getDeviceQueue device graphicsQueueFamilyIndex 0

  return $ DeviceParams
    devName
    physDevice
    device
    graphicsQueue
    graphicsQueueFamilyIndex

data PhysicalDeviceInfo = PhysicalDeviceInfo
  { pdiTotalMemory              :: Word64
  , pdiGraphicsQueueFamilyIndex :: Word32
  }
  deriving (Eq, Ord)

physicalDeviceInfo
  :: MonadIO m
  => Vk.SurfaceKHR
  -> Vk.PhysicalDevice
  -> m (Maybe PhysicalDeviceInfo)
physicalDeviceInfo surface phys = runMaybeT $ do
  -- We need swapchains
  guard =<< deviceHasSwapchains phys

  -- We also need graphics and present queues
  pdiGraphicsQueueFamilyIndex <- do
    queueFamilyProperties <- Vk.getPhysicalDeviceQueueFamilyProperties phys
    let isGraphicsQueue q =
          ((Vk.QUEUE_GRAPHICS_BIT .&. Vk.queueFlags q) > Vk.QueueFlagBits 0) &&
          (Vk.queueCount q > 0)
        graphicsQueueIndices = fromIntegral . fst <$> V.filter
          (isGraphicsQueue . snd)
          (V.indexed queueFamilyProperties)
        isPresentQueue i = Vk.getPhysicalDeviceSurfaceSupportKHR phys i surface
    presentQueueIndices <- V.filterM isPresentQueue graphicsQueueIndices
    MaybeT (return $ presentQueueIndices V.!? 0)

  -- Score based on total memory
  pdiTotalMemory <- do
    heaps <- Vk.memoryHeaps <$> Vk.getPhysicalDeviceMemoryProperties phys
    return $ sum (( Vk.size :: Vk.MemoryHeap -> Vk.DeviceSize) <$> heaps)

  return $ PhysicalDeviceInfo {..}

deviceHasSwapchains :: MonadIO m => Vk.PhysicalDevice -> m Bool
deviceHasSwapchains phys = do
  (_, extensions) <- Vk.enumerateDeviceExtensionProperties phys Nothing
  return $
    V.any ((Vk.KHR_SWAPCHAIN_EXTENSION_NAME ==) . Vk.extensionName) extensions

createVMA :: MonadResource m => Vk.Instance -> Vk.PhysicalDevice -> Vk.Device -> m Vk.Allocator
createVMA inst physicalDevice logicalDevice = snd <$>
  Vk.withAllocator
    Vk.zero
      { Vk.flags            = Vk.zero
      , Vk.physicalDevice   = Vk.physicalDeviceHandle physicalDevice
      , Vk.device           = Vk.deviceHandle logicalDevice
      , Vk.instance'        = Vk.instanceHandle inst
      , Vk.vulkanApiVersion = myApiVersion
      }
    allocate