{-# 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 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

-- internal imports

import Instance
import Devices
import GraphicsPipeline
import Framebuffers
import Memory
import Types
import Mesh
import CommandBuffer
import Sync
import Textures
import Util

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, queueIndex) <-
    createLogicalDevice vulkanPhysicalDevice vulkanSurface

  uploadContext <- initSyncStructures vulkanLogicalDevice queueIndex

  (descriptorSetLayout1, descriptorSetLayout2, 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/tri_mesh.vert" "vert"
  meshFragmentShader <- loadShader vulkanLogicalDevice "shadersrc/rainbow_lit.frag" "frag"
  renderPass <- createRenderPass vulkanLogicalDevice (Khr.format surfaceFormat) Vk.FORMAT_D32_SFLOAT
  meshLayout <- createMeshPipelineLayout
    vulkanLogicalDevice
    (V.fromList
      [ descriptorSetLayout1
      , descriptorSetLayout2
      ]
      )
  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
      vulkanLogicalDevice
      allocator
      descriptorPool
      descriptorSetLayout1
      descriptorSetLayout2
      queueIndex

  loadMeshes allocator uploadContext queue vulkanLogicalDevice

  loadImages allocator uploadContext queue vulkanLogicalDevice

  initScene

  return $ EngineData
    window
    dimensions
    vulkanPhysicalDevice
    vulkanLogicalDevice
    vulkanInstance
    swapchain
    queue
    frameBuffers
    renderPass
    allocator
    depthImageView
    depthAllocatedImage
    Vk.FORMAT_D32_SFLOAT
    frames
    descriptorSetLayout1
    descriptorSetLayout2
    descriptorPool
    (GPUSceneData
      (V4 0 0 0 0)
      (V4 0 0 0 0)
      (V4 0 0 0 0)
      (V4 0 0 0 0)
      (V4 0 0 0 0)
      )
    uploadContext


initScene :: (MonadReader ReadState m, MonadIO m) => m ()
initScene = do
  let catMask = 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` catMask) V.++ triangles

initDescriptors
  :: (MonadResource m)
  => Vk.Device
  -> m (Vk.DescriptorSetLayout, Vk.DescriptorSetLayout, Vk.DescriptorPool)
initDescriptors device = do
  let cameraBind = descriptorsetLayoutBinding
        Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
        Vk.SHADER_STAGE_VERTEX_BIT
        0
      sceneBind = descriptorsetLayoutBinding
        Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER -- _DYNAMIC
        Vk.SHADER_STAGE_FRAGMENT_BIT
        1
      objectBind = descriptorsetLayoutBinding
        Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER
        Vk.SHADER_STAGE_VERTEX_BIT
        0
      setInfo = Vk.zero
        { Vk.flags = Vk.DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT
        , Vk.bindings = V.fromList
          [ cameraBind
          , sceneBind
          ]
        }
      setInfo2 = Vk.zero
        { Vk.flags = Vk.DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT
        , Vk.bindings = V.singleton objectBind
        }
      poolInfo = Vk.zero
        { Vk.flags     = Vk.DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT
        , Vk.maxSets   = 10
        , Vk.poolSizes = V.fromList
          [ Vk.zero
            { Vk.type'           = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            , Vk.descriptorCount = 10
            }
          , Vk.zero
            { Vk.type'           = Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC
            , Vk.descriptorCount = 10
            }
          , Vk.zero
            { Vk.type'           = Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER
            , Vk.descriptorCount = 10
            }
          ]
        }

  (_, descriptorSetLayout1) <-
    Vk.withDescriptorSetLayout device setInfo Nothing allocate

  (_, descriptorSetLayout2) <-
    Vk.withDescriptorSetLayout device setInfo2 Nothing allocate

  (_, descriptorPool) <-
    Vk.withDescriptorPool device poolInfo Nothing allocate

  return (descriptorSetLayout1, descriptorSetLayout2, descriptorPool)