{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Textures where

import qualified Codec.Picture as P
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 qualified Data.Vector.Storable as VS
import Foreign hiding (void)
import qualified Control.Concurrent.STM as STM
import qualified Vulkan as Vk
import qualified Vulkan.CStruct.Extends as Vk
import qualified Vulkan.Zero as Vk
import qualified VulkanMemoryAllocator as VMA

-- internal imports

import Types
import Memory
import Sync
import Image

loadImageFromFile
  :: (MonadResource m, MonadIO m)
  => VMA.Allocator
  -> UploadContext
  -> Vk.Queue
  -> Vk.Device
  -> FilePath
  -> m AllocatedImage
loadImageFromFile allocator uploadContext queue device file = do

  pixels <- liftIO $ either error P.convertRGBA8 <$> P.readImage file

  let texWidth = P.imageWidth pixels
      texHeight = P.imageHeight pixels
      imageSize = 4 * texWidth * texHeight
      imageFormat = Vk.FORMAT_R8G8B8A8_SRGB

  stagingBuffer <-
    createAllocatedBuffer
    allocator
    imageSize
    Vk.BUFFER_USAGE_TRANSFER_SRC_BIT
    VMA.MEMORY_USAGE_CPU_ONLY

  let dataPtr = VMA.mappedData (bufferInfo stagingBuffer)

  liftIO $ pokeArray (castPtr dataPtr) (VS.toList $ P.imageData pixels)

  let imageExtent = Vk.Extent3D
        { width  = fromIntegral texWidth
        , height = fromIntegral texHeight
        , depth  = 1
        }
      dImgAllocCreateInfo = Vk.zero
        { VMA.usage = VMA.MEMORY_USAGE_GPU_ONLY
        } :: VMA.AllocationCreateInfo
      dImgCreateInfo = imageCreate
        imageFormat
        imageExtent
        (Vk.IMAGE_USAGE_SAMPLED_BIT .|. Vk.IMAGE_USAGE_TRANSFER_DST_BIT)

  (newImage, newAllocation, _) <- snd <$>
    VMA.withImage allocator dImgCreateInfo dImgAllocCreateInfo allocate

  immediateSubmit
    uploadContext
    queue
    device
    (\cmd -> do
      let range = Vk.zero
            { Vk.aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
            , Vk.baseMipLevel   = 0
            , Vk.levelCount     = 1
            , Vk.baseArrayLayer = 0
            , Vk.layerCount     = 1
            }
          imageBarrierToTransfer = V.singleton $ Vk.SomeStruct (Vk.zero
            { Vk.oldLayout        = Vk.IMAGE_LAYOUT_UNDEFINED
            , Vk.newLayout        = Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
            , Vk.image            = newImage
            , Vk.subresourceRange = range
            -- , Vk.srcAccessMask    = bit 0
            , Vk.dstAccessMask    = Vk.ACCESS_TRANSFER_WRITE_BIT
            } :: Vk.ImageMemoryBarrier '[])

      Vk.cmdPipelineBarrier
        cmd
        Vk.PIPELINE_STAGE_TOP_OF_PIPE_BIT
        Vk.PIPELINE_STAGE_TRANSFER_BIT
        (bit 0)
        V.empty
        V.empty
        imageBarrierToTransfer

      let copyRegion = V.singleton (Vk.zero
            { Vk.bufferOffset      = 0
            , Vk.bufferRowLength   = 0
            , Vk.bufferImageHeight = 0
            , Vk.imageExtent       = imageExtent
            , Vk.imageSubresource  = Vk.zero
              { Vk.aspectMask     = Vk.IMAGE_ASPECT_COLOR_BIT
              , Vk.mipLevel       = 0
              , Vk.baseArrayLayer = 0
              , Vk.layerCount     = 1
              } :: Vk.ImageSubresourceLayers
            } :: Vk.BufferImageCopy)

      Vk.cmdCopyBufferToImage
        cmd
        (allocatedBuffer stagingBuffer)
        newImage
        Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
        copyRegion

      let imageBarrierToReadable = V.singleton $ Vk.SomeStruct (Vk.zero
            { Vk.oldLayout = Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
            , Vk.newLayout = Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            , Vk.image            = newImage
            , Vk.subresourceRange = range
            , Vk.srcAccessMask = Vk.ACCESS_TRANSFER_WRITE_BIT
            , Vk.dstAccessMask = Vk.ACCESS_SHADER_READ_BIT
            } :: Vk.ImageMemoryBarrier '[])

      Vk.cmdPipelineBarrier
        cmd
        Vk.PIPELINE_STAGE_TRANSFER_BIT
        Vk.PIPELINE_STAGE_FRAGMENT_SHADER_BIT
        (bit 0)
        V.empty
        V.empty
        imageBarrierToReadable
      )

  return $ AllocatedImage
    { image = newImage
    , allocation = newAllocation
    }

loadImages
  :: (MonadResource m, MonadReader ReadState m)
  => VMA.Allocator
  -> UploadContext
  -> Vk.Queue
  -> Vk.Device
  -> m ()
loadImages allocator uploadContext queue device = do

  loadedImage <- loadImageFromFile allocator uploadContext queue device "assets/textures/jaeger.texture.tga"

  let imageInfo = imageviewCreate (image loadedImage) Vk.FORMAT_R8G8B8A8_SRGB Vk.IMAGE_ASPECT_COLOR_BIT

  imageView <- snd <$> Vk.withImageView device imageInfo Nothing allocate

  let texture = Texture
        { textureImage     = loadedImage
        , textureImageView = imageView
        }

  textureLib <- asks textureLibrary

  void $ liftIO $ STM.atomically $ do
    texMap <- STM.readTMVar textureLib

    let newMap = M.insert "marksman" texture texMap

    STM.swapTMVar textureLib newMap
