{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module CommandBuffer where

import qualified Control.Concurrent.STM as STM
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Trans.Resource
import qualified Data.Map.Strict as M
import qualified Data.Vector as V
import Foreign hiding (void)
import Foreign.C.Types (CInt)
import Linear
import qualified Vulkan as Vk
import qualified Vulkan.Zero as Vk
import qualified VulkanMemoryAllocator as VMA

-- internal imports

import Types
import Memory
import Util
import Mesh (compactDraws)

frameOverlap :: Int
frameOverlap = 2

maxObjects :: Int
maxObjects = 10 ^ 6

createFrames
  :: (MonadResource m, MonadFail m)
  => Vk.Device
  -> VMA.Allocator
  -> Vk.DescriptorPool
  -> Vk.DescriptorSetLayout
  -> Vk.DescriptorSetLayout
  -> Word32
  -> m (V.Vector FrameData, Vk.Queue)
createFrames
  logicalDevice
  allocator
  descriptorPool
  setLayout
  objectSetLayout
  queueFamilyIndex
  = do

  let poolCreateInfo = Vk.zero
        { Vk.flags            = Vk.COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT
        , Vk.queueFamilyIndex = queueFamilyIndex
        } :: Vk.CommandPoolCreateInfo

  graphicsQueue <- Vk.getDeviceQueue logicalDevice queueFamilyIndex 0

  frames <- V.mapM
    (\_ -> do

      commandPool <- snd <$>
        Vk.withCommandPool logicalDevice poolCreateInfo Nothing allocate

      let commandBufferAllocationInfo = Vk.zero
            { Vk.commandPool        = commandPool
            , Vk.level              = Vk.COMMAND_BUFFER_LEVEL_PRIMARY
            , Vk.commandBufferCount = 1
            }

      (_, commandBuffer) <- Vk.withCommandBuffers logicalDevice commandBufferAllocationInfo allocate

      let semaphoreCreateInfo = Vk.zero
          fenceCreateInfo = Vk.zero
            { Vk.flags = Vk.FENCE_CREATE_SIGNALED_BIT
            } :: Vk.FenceCreateInfo '[]

      [presentSemaphore, renderSemaphore] <- replicateM 2
        (snd <$>
          Vk.withSemaphore logicalDevice semaphoreCreateInfo Nothing allocate
          )
      renderFence <- snd <$>
        Vk.withFence logicalDevice fenceCreateInfo Nothing allocate

      cameraBuffer <- createAllocatedBuffer
        allocator
        (sizeOf (GPUCameraData identity identity identity))
        Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT
        VMA.MEMORY_USAGE_CPU_TO_GPU

      sceneBuffer <- createAllocatedBuffer
        allocator
        (sizeOf undefinedGPUSceneData * frameOverlap)
        Vk.BUFFER_USAGE_UNIFORM_BUFFER_BIT
        VMA.MEMORY_USAGE_CPU_TO_GPU

      objectBuffer <- createAllocatedBuffer
        allocator
        (sizeOf undefinedGPUObjectData * maxObjects)
        Vk.BUFFER_USAGE_STORAGE_BUFFER_BIT
        VMA.MEMORY_USAGE_CPU_TO_GPU

      indirectBuffer <- createAllocatedBuffer
        allocator
        (maxObjects * sizeOf (undefined :: Vk.DrawIndirectCommand))
        (Vk.BUFFER_USAGE_TRANSFER_DST_BIT .|.
          Vk.BUFFER_USAGE_STORAGE_BUFFER_BIT .|.
          Vk.BUFFER_USAGE_INDIRECT_BUFFER_BIT
          )
        VMA.MEMORY_USAGE_CPU_TO_GPU

      let allocationInfo = Vk.zero
            { Vk.descriptorPool = descriptorPool
            , Vk.setLayouts     = V.singleton setLayout
            }
          objectSetAllocInfo = Vk.zero
            { Vk.descriptorPool = descriptorPool
            , Vk.setLayouts     = V.singleton objectSetLayout
            }

      globalDescriptor <- Vk.allocateDescriptorSets logicalDevice allocationInfo
      objectDescriptor <- Vk.allocateDescriptorSets logicalDevice objectSetAllocInfo

      let cameraInfo = Vk.zero
            { Vk.buffer = allocatedBuffer cameraBuffer
            , Vk.offset = 0
            , Vk.range  = fromIntegral $ sizeOf (undefined :: GPUCameraData)
            } :: Vk.DescriptorBufferInfo
          sceneInfo = Vk.zero
            { Vk.buffer = allocatedBuffer sceneBuffer
            , Vk.offset = 0
            -- , Vk.offset = padUniformBufferSize (fromIntegral $ sizeOf undefinedGPUSceneData) deviceProperties * fromIntegral index
            , Vk.range  = fromIntegral $ sizeOf undefinedGPUSceneData
            } :: Vk.DescriptorBufferInfo
          objectInfo = Vk.zero
            { Vk.buffer = allocatedBuffer objectBuffer
            , Vk.offset = 0
            , Vk.range  = fromIntegral $ sizeOf undefinedGPUObjectData * maxObjects
            } :: Vk.DescriptorBufferInfo
          cameraWrite = writeDescriptorBuffer
            Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            (V.head globalDescriptor)
            cameraInfo
            0
          sceneWrite = writeDescriptorBuffer
            Vk.DESCRIPTOR_TYPE_UNIFORM_BUFFER
            (V.head globalDescriptor)
            sceneInfo
            1
          objectWrite = writeDescriptorBuffer
            Vk.DESCRIPTOR_TYPE_STORAGE_BUFFER
            (V.head objectDescriptor)
            objectInfo
            0
          setWrites = V.fromList
            [ cameraWrite
            , sceneWrite
            , objectWrite
            ]


      Vk.updateDescriptorSets logicalDevice setWrites V.empty

      return FrameData
        { framePresentSemaphore  = presentSemaphore
        , frameRenderSemaphore   = renderSemaphore
        , frameRenderFence       = renderFence
        , frameCommandPool       = commandPool
        , frameMainCommandBuffer = V.head commandBuffer
        , frameCameraBuffer      = cameraBuffer
        , frameSceneBuffer       = sceneBuffer
        , frameObjectBuffer      = objectBuffer
        , frameIndirectBuffer    = indirectBuffer
        , frameGlobalDescriptor  = V.head globalDescriptor
        , frameObjectDescriptor  = V.head objectDescriptor
        }

      )
    (V.fromList [1 .. frameOverlap])

  return (frames, graphicsQueue)

recordCommandBuffer
  :: (MonadResource m, MonadReader ReadState m, MonadIO m)
  => VMA.Allocator
  -> Vk.CommandBuffer
  -> Vk.RenderPass
  -> Vk.Framebuffer
  -> V2 CInt
  -> Vk.Pipeline
  -> FrameData
  -> m ()
recordCommandBuffer
  allocator
  commandBuffer
  renderPass
  frameBuffer
  dimensions@(V2 width height)
  graphicsPipeline
  frame
  = do

    let commandBufferBeginInfo = Vk.zero
          { Vk.flags = Vk.COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT
          , Vk.inheritanceInfo = Nothing
          }

    Vk.useCommandBuffer commandBuffer commandBufferBeginInfo $ do

      let renderPassInfo = Vk.zero
            { Vk.renderPass  = renderPass
            , Vk.framebuffer = frameBuffer
            , Vk.renderArea  = Vk.Rect2D
              { offset = Vk.Offset2D
                { Vk.x = 0
                , Vk.y = 0
                }
              , extent = Vk.Extent2D
                { Vk.width  = fromIntegral width
                , Vk.height = fromIntegral height
                }
              }
            , Vk.clearValues = V.fromList
              [ Vk.Color (Vk.Float32 0 0 0 1)
              , Vk.DepthStencil (Vk.ClearDepthStencilValue 1 0)
              ]
            }

      Vk.cmdUseRenderPass commandBuffer renderPassInfo Vk.SUBPASS_CONTENTS_INLINE $ do

        liftIO $
          prepareRecording dimensions commandBuffer graphicsPipeline

        renderObjects <- (liftIO . STM.atomically . STM.readTMVar) =<< asks renderables
        meshMap <- (liftIO . STM.atomically . STM.readTMVar) =<< asks meshLibrary
        materialMap <- (liftIO . STM.atomically . STM.readTMVar) =<< asks materialLibrary

        let draws = compactDraws meshMap materialMap renderObjects
            drawCommands = V.map
              (\(i, RenderObject mesh _ _) ->
                Vk.DrawIndirectCommand
                  { vertexCount   = fromIntegral (V.length $ meshVertices $ meshMap M.! mesh)
                  , firstVertex   = 0
                  , instanceCount = 1
                  , firstInstance = i
                  }
                )
              (V.zip
                (V.fromList [0..])
                renderObjects
                )

        (memRelease, memPtr) <- VMA.withMappedMemory
          allocator
          (bufferAllocation $ frameIndirectBuffer frame)
          allocate

        V.mapM_
          (\(idx, command) -> liftIO $
            poke
              (castPtr memPtr `plusPtr` (idx * sizeOf command))
              command
          )
          (V.zip
            (V.fromList [0 ..])
            drawCommands
           )

        release memRelease

        V.mapM_
          (\(IndirectBatch mesh material first count) -> do

            when
              (materialTextureSet material /= Vk.NULL_HANDLE)
              (Vk.cmdBindDescriptorSets
                commandBuffer
                Vk.PIPELINE_BIND_POINT_GRAPHICS
                (materialPipelineLayout material)
                2
                (V.singleton $ materialTextureSet material)
                V.empty
                )

            Vk.cmdBindPipeline commandBuffer Vk.PIPELINE_BIND_POINT_GRAPHICS (materialPipeline material)

            Vk.cmdBindDescriptorSets
              commandBuffer
              Vk.PIPELINE_BIND_POINT_GRAPHICS
              (materialPipelineLayout material)
              0
              (V.singleton $ frameGlobalDescriptor frame)
              V.empty

            Vk.cmdBindDescriptorSets
              commandBuffer
              Vk.PIPELINE_BIND_POINT_GRAPHICS
              (materialPipelineLayout material)
              1
              (V.singleton $ frameObjectDescriptor frame)
              V.empty

            Vk.cmdBindVertexBuffers
              commandBuffer
              0
              (V.fromList [ allocatedBuffer $ meshBuffer mesh ])
              (V.fromList [ 0 ])

            let indirectOffset = fromIntegral $
                  first * fromIntegral ( sizeOf (undefined :: Vk.DrawIndirectCommand))
                drawStride = fromIntegral $ sizeOf (undefined :: Vk.DrawIndirectCommand)

            Vk.cmdDrawIndirect
              commandBuffer
              (allocatedBuffer $ frameIndirectBuffer frame)
              indirectOffset
              count
              drawStride
            )
          draws

prepareRecording
  :: V2 CInt
  -> Vk.CommandBuffer
  -> Vk.Pipeline
  -> IO ()
prepareRecording (V2 width height) commandBuffer graphicsPipeline = do
  Vk.cmdBindPipeline commandBuffer Vk.PIPELINE_BIND_POINT_GRAPHICS graphicsPipeline

  -- Passing negative height and nonzero y origin to the viewport flips the models back in order
  let viewport = Vk.zero
        { Vk.x        = 0
        , Vk.y        = fromIntegral height
        , Vk.width    = fromIntegral width
        , Vk.height   = - fromIntegral height
        , Vk.minDepth = 0
        , Vk.maxDepth = 1
        }
      scissor = Vk.Rect2D
        { Vk.offset = Vk.Offset2D
          { Vk.x = 0
          , Vk.y = 0
          }
        , Vk.extent = Vk.Extent2D
          { Vk.width = fromIntegral width
          , Vk.height = fromIntegral height
          }
        }

  Vk.cmdSetViewport commandBuffer 0 (V.singleton viewport)
  Vk.cmdSetScissor commandBuffer 0 (V.singleton scissor)

pokeData
  :: (MonadResource m, MonadReader ReadState m, MonadIO m)
  => FrameData
  -> Int
  -> GPUSceneData
  -> GPUCameraData
  -> m ()
pokeData frame frameNumber sceneParameters cameraParameters = do
  renderObjects <- (liftIO . STM.atomically . STM.readTMVar) =<< asks renderables

  let scenePointer = VMA.mappedData (bufferInfo $ frameSceneBuffer frame)
      memoryPointer = VMA.mappedData (bufferInfo $ frameCameraBuffer frame)
      objectPointer = VMA.mappedData (bufferInfo $ frameObjectBuffer frame)
      framed = fromIntegral frameNumber / 120
      ambient = normalize $ V4 (sin framed) 0 (cos framed) 1
      -- ambient = V4 1 0 1 1
      ambientParams = sceneParameters
        { ambientColor = ambient
        }

  V.mapM_
    (\(index, obj) -> do

      let modelMatrix = GPUObjectData (objectMatrix obj)

      liftIO $ poke
        (castPtr objectPointer `plusPtr` (index * sizeOf (undefined :: GPUObjectData)))
        modelMatrix
      )
    (V.zip
      (V.fromList [0 ..])
      renderObjects
      )

  liftIO $ do
    poke
      (castPtr scenePointer)
      ambientParams

    poke
      (castPtr memoryPointer)
      cameraParameters
