{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Init where

import qualified Control.Concurrent.STM as STM
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Trans.Resource
import Data.Bits (Bits(bit))
import qualified Data.Vector as V
import qualified SDL hiding (V2)
import qualified SDL.Video.Vulkan as SDL
import Foreign.Ptr
import Linear as L
import qualified Vulkan.Core10 as Vk
import qualified Vulkan.Extensions.VK_KHR_swapchain as Khr
import qualified Vulkan.Extensions.VK_KHR_surface as Khr
import qualified Vulkan.Zero as Vk
import qualified VulkanMemoryAllocator as VMA

-- internal imports

import Instance
import Devices
import GraphicsPipeline
import Framebuffers
import Memory
import Types
import Mesh
import CommandBuffer

initEngine
  :: Render EngineData
initEngine = do
  -- initialize SDL2 with all subsystems
  void $ allocate_
    SDL.initializeAll
    (do
      putStrLn "This is the end!"
      SDL.quit
      )

  -- create a window configuration
  let windowConfig = SDL.defaultWindow
        { SDL.windowGraphicsContext = SDL.VulkanContext -- enable Vulkan support
        , SDL.windowInitialSize     = V2 1280 1024
        }

  -- load vulkan library
  SDL.vkLoadLibrary Nothing

  -- create the window
  (_, window) <-
    allocate
      (do
        eexcept <- try (SDL.createWindow "Haskell ❤️ Vulkan" windowConfig)
        case eexcept of
          Left err -> error $ "here's your exception: " <> show (err :: SDL.SDLException)
          Right a  -> return a
        )
      (\window -> do
        liftIO $ putStrLn "destroying window"
        SDL.destroyWindow window
        )

  -- initialize vúlkan data structures
  initVulkan window


initVulkan
  :: SDL.Window
  -> Render EngineData
initVulkan window = do

  vulkanInstance <- createInstance window

  (_, vulkanSurface) <- allocate
    (Khr.SurfaceKHR <$>
      SDL.vkCreateSurface window (castPtr (Vk.instanceHandle vulkanInstance))
      )
    (\vulkanSurface -> do
      putStrLn "destroying surface"
      Khr.destroySurfaceKHR vulkanInstance vulkanSurface Nothing
      )

  vulkanPhysicalDevice <- pickPhysicalDevice vulkanInstance
  (vulkanLogicalDevice, surfaceFormats) <- createLogicalDevice vulkanPhysicalDevice vulkanSurface

  (descriptorSetLayout, descriptorPool) <- initDescriptors vulkanLogicalDevice

  allocator <- initAllocator vulkanPhysicalDevice vulkanLogicalDevice vulkanInstance

  dimensions <- liftIO $ SDL.windowInitialSize <$> SDL.getWindowConfig window
  (swapchain, surfaceFormat, depthImageView, depthAllocatedImage) <-
    createSwapchain vulkanSurface surfaceFormats dimensions vulkanLogicalDevice allocator
  imageViews <- getImageViewHandles swapchain surfaceFormat vulkanLogicalDevice
  meshVertexShader <- loadShader vulkanLogicalDevice "shadersrc/mesh.vert" "vert"
  meshFragmentShader <- loadShader vulkanLogicalDevice "shadersrc/rainbow.frag" "frag"
  renderPass <- createRenderPass vulkanLogicalDevice (Khr.format surfaceFormat) Vk.FORMAT_D32_SFLOAT
  meshLayout <- createMeshPipelineLayout vulkanLogicalDevice descriptorSetLayout
  let meshContainer = ShaderContainer (Just meshVertexShader) (Just meshFragmentShader)
  createGraphicsPipelines
    vulkanLogicalDevice
    renderPass
    meshContainer
    dimensions
    meshLayout
    (Just (createDepthStencilStateCreateInfo True True Vk.COMPARE_OP_LESS_OR_EQUAL))

  frameBuffers <- createFramebuffer vulkanLogicalDevice renderPass imageViews depthImageView dimensions

  (frames, queue) <- createFrames vulkanPhysicalDevice vulkanLogicalDevice allocator descriptorPool descriptorSetLayout

  loadMeshes allocator

  initScene

  EngineData
    window
    dimensions
    vulkanPhysicalDevice
    vulkanLogicalDevice
    vulkanInstance
    swapchain
    queue
    frameBuffers
    meshLayout
    renderPass
    allocator
    depthImageView
    depthAllocatedImage
    Vk.FORMAT_D32_SFLOAT
    frames
    descriptorSetLayout
    descriptorPool
    <$> Vk.getPhysicalDeviceProperties vulkanPhysicalDevice

initScene :: (MonadReader ReadState m, MonadIO m) => m ()
initScene = do
  let mask = RenderObject
        { objectMesh     = "mask"
        , objectMaterial = "defaultMesh"
        , objectMatrix   = identity
        }
      triangles = V.concatMap
        (\y ->
          V.map
            (\x ->
              let
                translationMatrix = mkTransformation (Quaternion 1 (V3 0 0 0)) (V3 x 0 y)
                scaleMatrix       = scaled (V4 0.2 0.2 0.2 1)
              in
                RenderObject
                  { objectMesh     = "triangle"
                  , objectMaterial = "defaultMesh"
                  , objectMatrix   = translationMatrix !*! scaleMatrix
                  }
              )
            (V.fromList [-20 .. 20])
          )
        (V.fromList [-20 .. 20])

  renderableContainer <- asks renderables
  renderableVector <- liftIO $ STM.atomically $ STM.readTMVar renderableContainer

  void $ liftIO $ STM.atomically $ STM.swapTMVar renderableContainer $
    (renderableVector `V.snoc` mask) V.++ triangles

initDescriptors
  :: (MonadResource m)
  => Vk.Device
  -> m (Vk.DescriptorSetLayout, Vk.DescriptorPool)
initDescriptors device = do
  let cameraBufferBinding = Vk.zero
        { Vk.binding         = 0
        , Vk.descriptorCount = 1
        , Vk.descriptorType  = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
        , Vk.stageFlags      = Vk.SHADER_STAGE_VERTEX_BIT
        }
      setInfo = Vk.zero
        { Vk.flags = Vk.DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT
        , Vk.bindings = V.singleton cameraBufferBinding
        }
      poolInfo = Vk.zero
        { Vk.flags     = Vk.DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT
        , Vk.maxSets   = 10
        , Vk.poolSizes = V.singleton
          ( Vk.zero
            { Vk.type'           = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            , Vk.descriptorCount = 10
            }
          )
        }

  (_, descriptorSetLayout) <- allocate
    (Vk.createDescriptorSetLayout device setInfo Nothing)
    (\descriptorLayout -> do
      putStrLn "destroying descriptor layout"
      Vk.destroyDescriptorSetLayout device descriptorLayout Nothing
      )

  (_, descriptorPool) <- allocate
    (Vk.createDescriptorPool device poolInfo Nothing)
    (\descriptorPoolInfo -> do
      putStrLn "destroying descriptor pool"
      Vk.destroyDescriptorPool device descriptorPoolInfo Nothing
      )

  return (descriptorSetLayout, descriptorPool)