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 :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env, Typeable a, HasCallStack) =>
Queues CommandPool -> Source -> m (Texture a)
load Queues CommandPool
pool Source
source =
  forall a. HasCallStack => (HasCallStack => a) -> a
withFrozenCallStack forall a b. (a -> b) -> a -> b
$
    forall a (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, Typeable a,
 HasCallStack) =>
(ByteString -> m a) -> Source -> m a
Source.load (forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Queues CommandPool
pool) Source
source

loadBytes
  :: ( TextureLayers a
     , MonadVulkan env m
     , MonadResource m
     , MonadThrow m
     , HasLogFunc env
     )
  => Queues Vk.CommandPool
  -> ByteString
  -> m (Texture a)
loadBytes :: forall a env (m :: * -> *).
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> ByteString -> m (Texture a)
loadBytes Queues CommandPool
pool ByteString
bytes = do
  case ByteString -> Either (ByteOffset, String) Ktx
Ktx1.fromByteString ByteString
bytes of
    Left (ByteOffset
offset, String
err) -> do
      forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logError forall a b. (a -> b) -> a -> b
$ Utf8Builder
"Texture load error: " forall a. Semigroup a => a -> a -> a
<> forall a. IsString a => String -> a
fromString String
err
      forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ ByteOffset -> Text -> TextureError
Texture.LoadError ByteOffset
offset (String -> Text
Text.pack String
err)
    Right Ktx
ktx1 ->
      forall a (m :: * -> *) env.
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> Ktx -> m (Texture a)
loadKtx1 Queues CommandPool
pool Ktx
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 :: forall a (m :: * -> *) env.
(TextureLayers a, MonadVulkan env m, MonadResource m, MonadThrow m,
 HasLogFunc env) =>
Queues CommandPool -> Ktx -> m (Texture a)
loadKtx1 Queues CommandPool
pool Ktx1.Ktx{Header
header :: Ktx -> Header
header :: Header
header, images :: Ktx -> MipLevels
images=MipLevels
ktxImages} = do
  env
context <- forall r (m :: * -> *). MonadReader r m => m r
ask
  let vma :: Allocator
vma = forall a. HasVulkan a => a -> Allocator
getAllocator env
context

  let
    skipMips :: Int
skipMips = Int
0 -- DEBUG: Vector.length ktxImages `div` 2
    mipsSkipped :: Int
mipsSkipped = forall a. Ord a => a -> a -> a
min (forall a. Vector a -> Int
Vector.length MipLevels
ktxImages forall a. Num a => a -> a -> a
- Int
1) Int
skipMips
    images :: MipLevels
images = forall a. Int -> Vector a -> Vector a
Vector.drop Int
mipsSkipped MipLevels
ktxImages
    mipLevels :: Word32
mipLevels = Header -> Word32
Ktx1.numberOfMipmapLevels Header
header forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
mipsSkipped

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Vector a -> Bool
Vector.null MipLevels
images) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ ByteOffset -> Text -> TextureError
Texture.LoadError ByteOffset
0 Text
"At least one image must be present in KTX"

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
mipLevels forall a. Eq a => a -> a -> Bool
== forall a. Vector a -> Int
Vector.length MipLevels
images) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Word32 -> Int -> TextureError
Texture.MipLevelsError Word32
mipLevels (forall a. Vector a -> Int
Vector.length MipLevels
images)

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vk_format.h#L676
  Format
format <- case Header -> Word32
Ktx1.glInternalFormat Header
header of
    -- XXX: Force all BC7s to SRGB
    Word32
36492 ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Format
Vk.FORMAT_BC7_SRGB_BLOCK
    Word32
other ->
      case forall a. (Eq a, Num a) => a -> Maybe Format
FromGL.internalFormat Word32
other of
        Maybe Format
Nothing ->
          forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Unexpected glInternalFormat: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Word32
other -- TODO: throwIo
        Just Format
fmt ->
          -- XXX: going in blind
          forall (f :: * -> *) a. Applicative f => a -> f a
pure Format
fmt
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
    [ Utf8Builder
"Loading format "
    , forall a. Display a => a -> Utf8Builder
display forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.glInternalFormat Header
header
    , Utf8Builder
" as "
    , forall a. Show a => a -> Utf8Builder
displayShow Format
format
    ]

  -- XXX: https://github.com/KhronosGroup/KTX-Software/blob/bf849b7f/lib/vkloader.c#L552
  let
    extent :: Extent3D
extent = Vk.Extent3D
      { $sel:width:Extent3D :: Word32
Vk.width  = Header -> Word32
Ktx1.pixelWidth Header
header forall a. Bits a => a -> Int -> a
`Foreign.shiftR` Int
mipsSkipped
      , $sel:height:Extent3D :: Word32
Vk.height = Header -> Word32
Ktx1.pixelHeight Header
header forall a. Bits a => a -> Int -> a
`Foreign.shiftR` Int
mipsSkipped
      , $sel:depth:Extent3D :: Word32
Vk.depth  = forall a. Ord a => a -> a -> a
max Word32
1 forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.pixelDepth Header
header
      }
    arrayLayers :: Word32
arrayLayers = forall a. Ord a => a -> a -> a
max Word32
1 forall a b. (a -> b) -> a -> b
$ Header -> Word32
Ktx1.numberOfArrayElements Header
header

  -- TODO: basisu can encode movies as arrays, this could be handy
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word32
arrayLayers forall a. Eq a => a -> a -> Bool
== Word32
1) do
    forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logError Utf8Builder
"TODO: arrayLayers > 1"
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> TextureError
Texture.ArrayError Word32
1 Word32
arrayLayers

  let
    numLayers :: Word32
numLayers = Header -> Word32
Ktx1.numberOfFaces Header
header
    mipSizes :: Vector Word32
mipSizes = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
(*) Word32
numLayers forall b c a. (b -> c) -> (a -> b) -> a -> c
. MipLevel -> Word32
Ktx1.imageSize) MipLevels
images

    offsets' :: Vector Word32
offsets' = forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
Vector.scanl' forall a. Num a => a -> a -> a
(+) Word32
0 Vector Word32
mipSizes
    totalSize :: Word32
totalSize = forall a. Vector a -> a
Vector.last Vector Word32
offsets'
    offsets :: Vector Word32
offsets = forall a. Vector a -> Vector a
Vector.init Vector Word32
offsets'

  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ Utf8Builder
"mipSizes: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> Utf8Builder
displayShow Vector Word32
mipSizes
  forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ Utf8Builder
"offsets: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> Utf8Builder
displayShow Vector Word32
offsets

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word32
numLayers forall a. Eq a => a -> a -> Bool
== forall a. TextureLayers a => Word32
textureLayers @a) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> TextureError
Texture.ArrayError (forall a. TextureLayers a => Word32
textureLayers @a) Word32
numLayers

  DstImage
dst <- forall env (m :: * -> *).
(MonadVulkan env m, MonadResource m) =>
Queues CommandPool
-> Maybe Text
-> Extent3D
-> Word32
-> Word32
-> Format
-> m DstImage
Image.allocateDst
    Queues CommandPool
pool
    forall a. Maybe a
Nothing -- XXX: Name the whole collection later, tagging source with its index.
    Extent3D
extent
    Word32
mipLevels
    Word32
numLayers
    Format
format

  forall (a :: [*]) (io :: * -> *) r.
(Extendss BufferCreateInfo a, PokeChain a, MonadIO io) =>
Allocator
-> BufferCreateInfo a
-> AllocationCreateInfo
-> (io (Buffer, Allocation, AllocationInfo)
    -> ((Buffer, Allocation, AllocationInfo) -> io ()) -> r)
-> r
VMA.withBuffer Allocator
vma (forall a. Integral a => a -> BufferCreateInfo '[]
Texture.stageBufferCI Word32
totalSize) AllocationCreateInfo
Texture.stageAllocationCI forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket \(Buffer
staging, Allocation
stage, AllocationInfo
stageInfo) -> do
    let mipImages :: Vector (Word32, MipLevel)
mipImages = forall a b. Vector a -> Vector b -> Vector (a, b)
Vector.zip Vector Word32
offsets MipLevels
images
    forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector (Word32, MipLevel)
mipImages \Int
mipIx (Word32
offset, Ktx1.MipLevel{Word32
imageSize :: Word32
imageSize :: MipLevel -> Word32
imageSize, Vector ArrayElement
arrayElements :: MipLevel -> Vector ArrayElement
arrayElements :: Vector ArrayElement
arrayElements}) -> do
      forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector ArrayElement
arrayElements \Int
arrayIx Ktx1.ArrayElement{Vector Face
faces :: ArrayElement -> Vector Face
faces :: Vector Face
faces} -> do
        forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector Face
faces \Int
faceIx Ktx1.Face{Vector ZSlice
zSlices :: Face -> Vector ZSlice
zSlices :: Vector ZSlice
zSlices} -> do
          forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
Vector.iforM_ Vector ZSlice
zSlices \Int
sliceIx Ktx1.ZSlice{ByteString
block :: ZSlice -> ByteString
block :: ByteString
block} -> do
            let
              indices :: Utf8Builder
indices = forall a. Monoid a => [a] -> a
mconcat
                [ Utf8Builder
"["
                , Utf8Builder
" mip:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
mipIx
                , Utf8Builder
" arr:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
arrayIx
                , Utf8Builder
" fac:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
faceIx
                , Utf8Builder
" slc:" forall a. Semigroup a => a -> a -> a
<> forall a. Display a => a -> Utf8Builder
display Int
sliceIx
                , Utf8Builder
" ]"
                ]
            let blockOffset :: Word32
blockOffset = Word32
offset forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
faceIx forall a. Num a => a -> a -> a
* Word32
imageSize
            let sectionPtr :: Ptr Any
sectionPtr = forall a b. Ptr a -> Int -> Ptr b
Foreign.plusPtr (AllocationInfo -> Ptr ()
VMA.mappedData AllocationInfo
stageInfo) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
blockOffset)
            forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) =>
Utf8Builder -> m ()
logDebug forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
              [ Utf8Builder
indices
              , Utf8Builder
" base offset = "
              , forall a. Display a => a -> Utf8Builder
display Word32
offset
              , Utf8Builder
" image offset = "
              , forall a. Display a => a -> Utf8Builder
display forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
faceIx forall a. Num a => a -> a -> a
* Word32
imageSize
              , Utf8Builder
" image size = "
              , forall a. Display a => a -> Utf8Builder
display Word32
imageSize
              ]
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
block \(Ptr CChar
pixelsPtr, Int
pixelBytes) -> do
              if Int
pixelBytes forall a. Eq a => a -> a -> Bool
/= forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
imageSize then
                forall a. HasCallStack => String -> a
error String
"assert: MipLevel.imageSize matches block.pixelBytes"
              else
                -- traceShowM (sectionPtr, pixelBytes)
                forall a. Ptr a -> Ptr a -> Int -> IO ()
Foreign.copyBytes Ptr Any
sectionPtr (forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr CChar
pixelsPtr) Int
pixelBytes

    forall (io :: * -> *).
MonadIO io =>
Allocator
-> Allocation
-> ("offset" ::: DeviceSize)
-> ("offset" ::: DeviceSize)
-> io ()
VMA.flushAllocation Allocator
vma Allocation
stage "offset" ::: DeviceSize
0 "offset" ::: DeviceSize
Vk.WHOLE_SIZE

    AllocatedImage
final <- forall env (m :: * -> *) deviceSize (t :: * -> *).
(MonadVulkan env m, Integral deviceSize, Foldable t) =>
Queues CommandPool
-> Buffer
-> DstImage
-> ("mip offsets" ::: t deviceSize)
-> m AllocatedImage
Image.copyBufferToDst
      Queues CommandPool
pool
      Buffer
staging
      DstImage
dst
      Vector Word32
offsets

    pure Texture
      { $sel:tFormat:Texture :: Format
tFormat         = Format
format
      , $sel:tMipLevels:Texture :: Word32
tMipLevels      = Word32
mipLevels
      , $sel:tLayers:Texture :: Word32
tLayers         = Word32
numLayers
      , $sel:tAllocatedImage:Texture :: AllocatedImage
tAllocatedImage = AllocatedImage
final
      }