Ktx1.hs

module Resource.Texture.Ktx1
  ( createTexture
  ) where

import RIO

import Codec.Ktx qualified as Ktx1
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Data.Text qualified as Text
import Data.Vector qualified as Vector
import Foreign qualified
import Vulkan.Core10 qualified as Vk
import Vulkan.Utils.FromGL qualified as FromGL
import VulkanMemoryAllocator qualified as VMA

import Engine.Types (StageRIO)
import Engine.Vulkan.Types (HasVulkan(..), Queues)
import Resource.Compressed.Zstd qualified as Zstd
import Resource.Image (AllocatedImage(..))
import Resource.Image qualified as Image
import Resource.Texture (Texture(..), TextureLayers(..))
import Resource.Texture qualified as Texture

createTexture
  :: forall a st . (TextureLayers a)
  => Queues Vk.CommandPool
  -> FilePath
  -> StageRIO st (Texture a)
createTexture pool path = do
  context <- ask
  let vma = getAllocator context

  logInfo $ "Loading " <> fromString path
  loader path >>= \case
    Left (offset, err) -> do
      logError $ "Texture load error: " <> fromString err
      throwM $ Texture.LoadError offset (Text.pack err)
    Right Ktx1.Ktx{header, images=ktxImages} -> do
      let
        skipMips = 0 -- DEBUG: Vector.length ktxImages `div` 2
        mipsSkipped = min (Vector.length ktxImages - 1) skipMips
        images = Vector.drop mipsSkipped ktxImages
        mipLevels = Ktx1.numberOfMipmapLevels header - fromIntegral mipsSkipped

      when (Vector.null images) $
        throwM $ Texture.LoadError 0 "At least one image must be present in KTX"

      unless (fromIntegral mipLevels == Vector.length images) $
        throwM $ Texture.MipLevelsError mipLevels (Vector.length images)

      -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vk_format.h#L676
      format <- case Ktx1.glInternalFormat header of
        -- XXX: Force all BC7s to SRGB
        36492 ->
          pure Vk.FORMAT_BC7_SRGB_BLOCK
        other ->
          case FromGL.internalFormat other of
            Nothing ->
              error $ "Unexpected glInternalFormat: " <> show other -- TODO: throwIo
            Just fmt ->
              -- XXX: going in blind
              pure fmt
      logDebug $ mconcat
        [ "Loading format "
        , display $ Ktx1.glInternalFormat header
        , " as "
        , displayShow format
        ]

      -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vkloader.c#L552
      let
        extent = Vk.Extent3D
          { Vk.width  = Ktx1.pixelWidth header `Foreign.shiftR` mipsSkipped
          , Vk.height = Ktx1.pixelHeight header `Foreign.shiftR` mipsSkipped
          , Vk.depth  = max 1 $ Ktx1.pixelDepth header
          }
        arrayLayers = max 1 $ Ktx1.numberOfArrayElements header

      -- TODO: basisu can encode movies as arrays, this could be handy
      unless (arrayLayers == 1) do
        logError "TODO: arrayLayers > 1"
        throwM $ Texture.ArrayError 1 arrayLayers

      let
        numLayers = Ktx1.numberOfFaces header
        mipSizes = fmap ((*) numLayers . Ktx1.imageSize) images

        offsets' = Vector.scanl' (+) 0 mipSizes
        totalSize = Vector.last offsets'
        offsets = Vector.init offsets'

      logDebug $ "mipSizes: " <> displayShow mipSizes
      logDebug $ "offsets: " <> displayShow offsets

      {- XXX:
        Image created before staging buffer due to `Image.copyBufferToImage`
        issued inside `VMA.withBuffer` bracket.
      -}
      (image, allocation, _info) <- VMA.createImage
        vma
        (Texture.imageCI format extent mipLevels numLayers)
        Texture.imageAllocationCI

      Image.transitionImageLayout
        context
        pool
        image
        mipLevels
        numLayers -- XXX: arrayLayers is always 0 for now
        format
        Vk.IMAGE_LAYOUT_UNDEFINED
        Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL

      unless (numLayers == textureLayers @a) $
        throwM $ Texture.ArrayError (textureLayers @a) numLayers

      VMA.withBuffer vma (Texture.stageBufferCI totalSize) Texture.stageAllocationCI bracket \(staging, stage, stageInfo) -> do
        let ixMipImages = Vector.zip3 (Vector.fromList [0..]) offsets images
        Vector.forM_ ixMipImages \(mipIx, offset, Ktx1.MipLevel{imageSize, arrayElements}) -> do
          let ixArrayElements = Vector.zip (Vector.fromList [0..]) arrayElements
          Vector.forM_ ixArrayElements \(arrayIx, Ktx1.ArrayElement{faces}) -> do
            let ixFaces = Vector.zip (Vector.fromList [0..]) faces
            Vector.forM_ ixFaces \(faceIx, Ktx1.Face{zSlices}) -> do
              let ixSlices = Vector.zip (Vector.fromList [0..]) zSlices
              Vector.forM_ ixSlices \(sliceIx, Ktx1.ZSlice{block}) -> do
                let
                  indices = mconcat
                    [ "["
                    , " mip:" <> display @Word32 mipIx
                    , " arr:" <> display @Word32 arrayIx
                    , " fac:" <> display @Word32 faceIx
                    , " slc:" <> display @Word32 sliceIx
                    , " ]"
                    ]
                let blockOffset = offset + faceIx * imageSize
                let sectionPtr = Foreign.plusPtr (VMA.mappedData stageInfo) (fromIntegral blockOffset)
                logDebug $ mconcat
                  [ indices
                  , " base offset = "
                  , display offset
                  , " image offset = "
                  , display $ faceIx * imageSize
                  , " image size = "
                  , display imageSize
                  ]
                liftIO $ unsafeUseAsCStringLen block \(pixelsPtr, pixelBytes) -> do
                  if pixelBytes /= fromIntegral imageSize then
                    error "assert: MipLevel.imageSize matches block.pixelBytes"
                  else
                    -- traceShowM (sectionPtr, pixelBytes)
                    Foreign.copyBytes sectionPtr (Foreign.castPtr pixelsPtr) pixelBytes

        VMA.flushAllocation vma stage 0 Vk.WHOLE_SIZE

        -- XXX: copying to image while the staging buffer is still alive
        Image.copyBufferToImage
          context
          pool
          staging
          image
          extent
          offsets
          numLayers

      -- XXX: staging buffer is gone

      Image.transitionImageLayout
        context
        pool
        image
        mipLevels
        numLayers -- XXX: arrayLayers is always 0 for now
        format
        Vk.IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
        Vk.IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL

      imageView <- Texture.createImageView
        context
        image
        format
        mipLevels
        numLayers

      let
        allocatedImage = AllocatedImage
          { aiAllocation = allocation
          , aiImage      = image
          , aiImageView  = imageView
          }
      pure Texture
        { tFormat         = format
        , tMipLevels      = mipLevels
        , tLayers         = numLayers
        , tAllocatedImage = allocatedImage
        }
  where
    loader = liftIO . Zstd.fromFileWith (pure . Ktx1.fromByteString) Ktx1.fromFile