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

import Codec.Wavefront

import qualified Control.Concurrent.STM as STM
import qualified Control.Monad as CM (void)
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Trans.Resource

import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import qualified Data.Vector as V

import Foreign

import Linear (V2(..), V3(..), V4(..))

import qualified VulkanMemoryAllocator as VMA
import qualified Vulkan.Core10 as Vk
import qualified Vulkan.Zero as Vk

-- internal imports

import Types
import Sync

loadMeshes
  :: (MonadReader ReadState m, MonadResource m, MonadFail m)
  => VMA.Allocator
  -> UploadContext
  -> Vk.Queue
  -> Vk.Device
  -> m ()
loadMeshes allocator uploadContext queue device = do
  let triangle = V.fromList
        [ Vertex (V3 0.5 0 0)    (V3 0 0 0) (V4 0 1 0 1) (V2 0 0)
        , Vertex (V3 (-0.5) 0 0) (V3 0 0 0) (V4 0 1 0 1) (V2 0 0)
        , Vertex (V3 0 1 0)   (V3 0 0 0) (V4 0 1 0 1) (V2 0 0)
        ]

  triMesh <- uploadMesh triangle allocator uploadContext queue device
  maskMesh <- loadFromObj
    "./assets/models/jaeger.obj"
    allocator
    uploadContext
    queue
    device

  meshLib <- asks meshLibrary

  liftIO $ STM.atomically $ do
    meshMap <- STM.readTMVar meshLib
    let newMap =
          M.insert "triangle" triMesh $
          M.insert "mask" maskMesh meshMap
    CM.void $ STM.swapTMVar meshLib newMap

uploadMesh
  :: (MonadResource m, MonadFail m)
  => V.Vector Vertex
  -> VMA.Allocator
  -> UploadContext
  -> Vk.Queue
  -> Vk.Device
  -> m Mesh
uploadMesh vertices allocator uploadContext queue device = do

  let bufferSize = fromIntegral $ V.length vertices * sizeOf (undefined :: Vertex)
      stagingBufferCrateInfo = Vk.zero
        { Vk.size  = bufferSize
        , Vk.usage = Vk.BUFFER_USAGE_TRANSFER_SRC_BIT
        }
      stagingAllocationCreateInfo = Vk.zero
        { VMA.usage = VMA.MEMORY_USAGE_CPU_ONLY
        } :: VMA.AllocationCreateInfo

  (stagingReleaseKey, stagingBuffer) <- do
    (releaseKey, (buffer, bAllocation, bInfo)) <-
      VMA.withBuffer allocator stagingBufferCrateInfo stagingAllocationCreateInfo allocate
    return (releaseKey, AllocatedBuffer
      { bufferInfo = bInfo
      , bufferAllocation = bAllocation
      , allocatedBuffer = buffer
      })

  (dataReleaseKey, dataPtr) <-
    VMA.withMappedMemory allocator (bufferAllocation stagingBuffer) allocate

  liftIO $ V.mapM_ (\(idx, vertex) ->
    poke
      (castPtr dataPtr `plusPtr` (idx * sizeOf vertex))
      vertex
    )
    (V.zip
      (V.fromList [0..])
      vertices
      )

  liftIO $ release dataReleaseKey

  let vertexBufferCreateInfo = Vk.zero
        { Vk.size = bufferSize
        , Vk.usage = Vk.BUFFER_USAGE_VERTEX_BUFFER_BIT .|. Vk.BUFFER_USAGE_TRANSFER_DST_BIT
        }
      allocationCreateInfo = Vk.zero
        { VMA.usage = VMA.MEMORY_USAGE_GPU_ONLY
        } :: VMA.AllocationCreateInfo

  mesh <- do
    (buffer, bAllocation, bInfo) <- snd <$>
      VMA.withBuffer allocator vertexBufferCreateInfo allocationCreateInfo allocate
    return $ Mesh
      { meshVertices = vertices
      , meshBuffer = AllocatedBuffer
        { bufferInfo = bInfo
        , bufferAllocation = bAllocation
        , allocatedBuffer = buffer
        }
      }

  immediateSubmit uploadContext queue device$ \cmd -> do
    let copy = V.singleton $ Vk.BufferCopy
          { size = bufferSize
          , dstOffset = 0
          , srcOffset = 0
          }

    Vk.cmdCopyBuffer cmd (allocatedBuffer stagingBuffer) (allocatedBuffer $ meshBuffer mesh) copy

  release stagingReleaseKey

  return mesh

loadFromObj
  :: (MonadResource m, MonadFail m)
  => FilePath
  -> VMA.Allocator
  -> UploadContext
  -> Vk.Queue
  -> Vk.Device
  -> m Mesh
loadFromObj filepath vma uploadContext queue device = do
  eitherObj <- liftIO $ fromFile filepath
  case eitherObj of
    Left err ->
      error ("loadFromObj:  loading mesh data: " <> err)
    Right obj -> do
      let vertices = digestFaces (objFaces obj) (objLocations obj) (objTexCoords obj)
      uploadMesh vertices vma uploadContext queue device
  where
    digestFaces
      :: V.Vector (Element Face) -> V.Vector Location -> V.Vector TexCoord -> V.Vector Vertex
    digestFaces faces locations texCoords =
      V.foldl
        (\acc (Element _ _ _ _ (Face fa fb fc others)) ->
          if not (null others)
          then
            error "loadFromObj: Non-triangular face detected"
          else
            let newVertices = V.map
                  (\faceIndex ->
                    let position = faceIndexToPosition faceIndex locations
                        normal   = V3 0 0 0
                        color    = V4 1 1 0 1
                        uv       = faceIndexToUV faceIndex texCoords
                    in
                      Vertex position normal color uv
                    )
                  (V.fromList [ fa, fb, fc ])
            in
              acc V.++ newVertices
          )
        V.empty
        faces
    faceIndexToPosition :: FaceIndex -> V.Vector Location -> V3 Float
    faceIndexToPosition index locations =
      let Location x y z _ = locations V.! (faceLocIndex index - 1)
      in
        V3 x y z
    faceIndexToUV :: FaceIndex -> V.Vector TexCoord -> V2 Float
    faceIndexToUV index texCoords =
      let TexCoord r s _ =
            texCoords V.!
              (fromMaybe (error "no UV coordinates present") (faceTexCoordIndex index) - 1)
      in
        V2 r s

compactDraws
  :: M.Map String Mesh
  -> M.Map String Material
  -> V.Vector RenderObject
  -> V.Vector IndirectBatch
compactDraws meshLib materialLib objects =

  V.foldl
    (\acc (index, RenderObject mesh material _) ->

      let newDraw = IndirectBatch
            { batchMesh     = meshLib M.! mesh
            , batchMaterial = materialLib M.! material
            , batchFirst    = index
            , batchCount    = 1
            }
          sameMesh = not (V.null acc) && (meshLib M.! mesh == batchMesh (V.last acc))
          sameMaterial = not (V.null acc) && (materialLib M.! material == batchMaterial (V.last acc))

      in
        if sameMesh && sameMaterial
        then
          V.init acc `V.snoc` ((V.last acc)
            { batchCount = batchCount (V.last acc) + 1
            })
        else
          acc `V.snoc` newDraw
      )
    V.empty
    (V.zip
      (V.fromList [0..])
      objects
      )