{-# 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
import qualified Data.Map.Strict as M
import Data.Maybe
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 as Vk
import qualified Vulkan.CStruct.Extends 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, textureDescriptorSetLayout, 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"
  textureFragmentShader <- loadShader vulkanLogicalDevice "shadersrc/textured_lit.frag" "frag"
  renderPass <- createRenderPass vulkanLogicalDevice (Khr.format surfaceFormat) Vk.FORMAT_D32_SFLOAT
  meshLayout <- createMeshPipelineLayout
    vulkanLogicalDevice
    (V.fromList
      [ descriptorSetLayout1
      , descriptorSetLayout2
      , textureDescriptorSetLayout
      ]
      )
  textureLayout <- createMeshPipelineLayout
    vulkanLogicalDevice
    (V.fromList
      [ descriptorSetLayout1
      , descriptorSetLayout2
      , textureDescriptorSetLayout
      ]
      )
  let meshContainer = ShaderContainer (Just meshVertexShader) (Just meshFragmentShader)
      textureContainer = ShaderContainer (Just meshVertexShader) (Just textureFragmentShader)
  createGraphicsPipelines
    vulkanLogicalDevice
    renderPass
    meshContainer
    dimensions
    meshLayout
    (Just (createDepthStencilStateCreateInfo True True Vk.COMPARE_OP_LESS_OR_EQUAL))
    "defaultMesh"
  createGraphicsPipelines
    vulkanLogicalDevice
    renderPass
    textureContainer
    dimensions
    textureLayout
    (Just (createDepthStencilStateCreateInfo True True Vk.COMPARE_OP_LESS_OR_EQUAL))
    "texturedMesh"

  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 vulkanLogicalDevice descriptorPool textureDescriptorSetLayout

  return $ EngineData
    window
    dimensions
    vulkanPhysicalDevice
    vulkanLogicalDevice
    vulkanInstance
    swapchain
    queue
    frameBuffers
    renderPass
    allocator
    depthImageView
    depthAllocatedImage
    Vk.FORMAT_D32_SFLOAT
    frames
    descriptorSetLayout1
    descriptorSetLayout2
    textureDescriptorSetLayout
    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, MonadResource m)
  => Vk.Device
  -> Vk.DescriptorPool
  -> Vk.DescriptorSetLayout
  -> m ()
initScene device descriptorPool texturedSetLayout = do
  let catMask = RenderObject
        { objectMesh     = "mask"
        , objectMaterial = "texturedMesh"
        , 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 [-50 .. 50])
          )
        (V.fromList [-50 .. 50])

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

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

  let samplerInfo = samplerCreateInfo Vk.FILTER_NEAREST Vk.SAMPLER_ADDRESS_MODE_REPEAT
  blockySampler <- snd <$> Vk.withSampler device samplerInfo Nothing allocate

  matLibraryContainer <- asks materialLibrary
  texturedMaterial <- do
    matLibrary <- liftIO $ STM.atomically $ STM.readTMVar matLibraryContainer
    return (fromMaybe (error "not texturedMesh") $ matLibrary M.!? "texturedMesh")

  let allocInfo = Vk.zero
        { Vk.descriptorPool = descriptorPool
        , Vk.setLayouts     = V.singleton texturedSetLayout
        }

  texturedDescriptorSet <- V.head . snd <$> Vk.withDescriptorSets device allocInfo allocate

  liftIO $ do
    matLibrary <- STM.atomically $ STM.readTMVar matLibraryContainer
    let material = texturedMaterial
          { materialTextureSet = texturedDescriptorSet
          }
        newMatLibrary = M.insert "texturedMesh" material matLibrary
    void $ STM.atomically $ STM.swapTMVar matLibraryContainer newMatLibrary

  textureContainer <- asks textureLibrary
  texture <- do
    texLibrary <- liftIO $ STM.atomically $ STM.readTMVar textureContainer
    return $ fromMaybe (error "no marksman texture found") $ texLibrary M.!? "marksman"

  let imageBufferInfo = Vk.zero
        { Vk.sampler = blockySampler
        , Vk.imageView = textureImageView texture
        , Vk.imageLayout = Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
        }
      texture1 = V.singleton $ Vk.SomeStruct $ writeDescriptorImage
        Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
        texturedDescriptorSet
        imageBufferInfo
        0

  Vk.updateDescriptorSets device texture1 V.empty

initDescriptors
  :: (MonadResource m)
  => Vk.Device
  -> m (Vk.DescriptorSetLayout, 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
      textureBind = descriptorsetLayoutBinding
        Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
        Vk.SHADER_STAGE_FRAGMENT_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
        }
      setInfo3 = Vk.zero
        { Vk.flags = Vk.DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT
        , Vk.bindings = V.singleton textureBind
        }
      poolInfo = Vk.zero
        { Vk.flags     =
            Vk.DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT .|.
            Vk.DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_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
            }
          , Vk.zero
            { Vk.type'           = Vk.DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER
            , Vk.descriptorCount = 10
            }
          ]
        }

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

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

  (_, textureDescriptorSetLayout) <-
    Vk.withDescriptorSetLayout device setInfo3 Nothing allocate

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

  return (descriptorSetLayout1, descriptorSetLayout2, textureDescriptorSetLayout, descriptorPool)
