Ktx1.hs

module Resource.Texture.Ktx1
  ( load
  , loadBytes
  , loadKtx1
  ) 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 GHC.Stack (withFrozenCallStack)
import UnliftIO.Resource (MonadResource)
import Vulkan.Core10 qualified as Vk
import Vulkan.Utils.FromGL qualified as FromGL
import VulkanMemoryAllocator qualified as VMA

import Engine.Vulkan.Types (HasVulkan(..), MonadVulkan, Queues)
import Resource.Image qualified as Image
import Resource.Source (Source(..))
import Resource.Source qualified as Source
import Resource.Texture (Texture(..), TextureLayers(..))
import Resource.Texture qualified as Texture

load
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     , Typeable a
     , HasCallStack
     )
  => Queues Vk.CommandPool
  -> Source
  -> m (Texture a)
load pool source =
  withFrozenCallStack $
    Source.load (loadBytes pool) source

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes pool bytes = do
  case Ktx1.fromByteString bytes of
    Left (offset, err) -> do
      logError $ "Texture load error: " <> fromString err
      throwM $ Texture.LoadError offset (Text.pack err)
    Right ktx1 ->
      loadKtx1 pool ktx1

loadKtx1
  :: forall a m env
  .  ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> Ktx1.Ktx
  -> m (Texture a)
loadKtx1 pool Ktx1.Ktx{header, images=ktxImages} = do
  context <- ask
  let vma = getAllocator context

  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

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

  dst <- Image.allocateDst
    pool
    Nothing -- XXX: Name the whole collection later, tagging source with its index.
    extent
    mipLevels
    numLayers
    format

  VMA.withBuffer vma (Texture.stageBufferCI totalSize) Texture.stageAllocationCI bracket \(staging, stage, stageInfo) -> do
    let mipImages = Vector.zip offsets images
    Vector.iforM_ mipImages \mipIx (offset, Ktx1.MipLevel{imageSize, arrayElements}) -> do
      Vector.iforM_ arrayElements \arrayIx Ktx1.ArrayElement{faces} -> do
        Vector.iforM_ faces \faceIx Ktx1.Face{zSlices} -> do
          Vector.iforM_ zSlices \sliceIx Ktx1.ZSlice{block} -> do
            let
              indices = mconcat
                [ "["
                , " mip:" <> display mipIx
                , " arr:" <> display arrayIx
                , " fac:" <> display faceIx
                , " slc:" <> display sliceIx
                , " ]"
                ]
            let blockOffset = offset + fromIntegral faceIx * imageSize
            let sectionPtr = Foreign.plusPtr (VMA.mappedData stageInfo) (fromIntegral blockOffset)
            logDebug $ mconcat
              [ indices
              , " base offset = "
              , display offset
              , " image offset = "
              , display $ fromIntegral 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

    final <- Image.copyBufferToDst
      pool
      staging
      dst
      offsets

    pure Texture
      { tFormat         = format
      , tMipLevels      = mipLevels
      , tLayers         = numLayers
      , tAllocatedImage = final
      }