Ktx.hs

module Codec.Ktx where

import Data.Binary (Binary(..), decodeFileOrFail, decodeOrFail)
import Data.Binary.Get (Get, ByteOffset, getWord32le, getWord32be, getByteString)
import Data.Binary.Put (Put, execPut, putByteString, putWord32le, putWord32be)
import Data.ByteString (ByteString)
import Data.ByteString qualified  as BS
import Data.ByteString.Builder (Builder, hPutBuilder)
import Data.ByteString.Lazy qualified  as BSL
import Data.Coerce (coerce)
import Data.Vector (Vector)
import Data.Vector qualified  as Vector
import Data.Word (Word32)
import GHC.Generics (Generic)
import System.IO (IOMode(..), withBinaryFile)

import Codec.Ktx.KeyValue (KeyValueData)
import Codec.Ktx.KeyValue qualified as KeyValue

fromByteStringLazy :: BSL.ByteString -> Either (ByteOffset, String) Ktx
fromByteStringLazy bsl =
  case decodeOrFail bsl of
    Right (_leftovers, _bytesLeft, ktx) ->
      Right ktx
    Left (_leftovers, bytesLeft, err) ->
      Left (bytesLeft, err)

fromByteString :: ByteString -> Either (ByteOffset, String) Ktx
fromByteString = fromByteStringLazy . BSL.fromStrict

fromFile :: FilePath -> IO (Either (ByteOffset, String) Ktx)
fromFile = decodeFileOrFail

toBuilder :: Ktx -> Builder
toBuilder = execPut . put

toFile :: FilePath -> Ktx -> IO ()
toFile dest ktx =
  withBinaryFile dest WriteMode $ \handle ->
    hPutBuilder handle (toBuilder ktx)

data Ktx = Ktx
  { header :: Header
  , kvs    :: KeyValueData
  , images :: MipLevels
  } deriving (Show, Generic)

instance Binary Ktx where
  get = do
    header <- get
    kvs <- KeyValue.getData
      (mkGetWord32 $ endianness header)
      (fromIntegral $ bytesOfKeyValueData header)
    images <- getImages header
    pure Ktx{..}

  put Ktx{..} = do
    put header
    KeyValue.putData putWord32 kvs
    putImages putWord32 images
    where
      putWord32 = mkPutWord32 $ endianness header

-- * Header

data Header = Header
  { identifier            :: ByteString
  , endianness            :: Word32
  , glType                :: Word32
  , glTypeSize            :: Word32
  , glFormat              :: Word32
  , glInternalFormat      :: Word32
  , glBaseInternalFormat  :: Word32
  , pixelWidth            :: Word32
  , pixelHeight           :: Word32
  , pixelDepth            :: Word32
  , numberOfArrayElements :: Word32
  , numberOfFaces         :: Word32
  , numberOfMipmapLevels  :: Word32
  , bytesOfKeyValueData   :: Word32
  } deriving (Show, Generic)

instance Binary Header where
  get = do
    identifier <- getByteString 12
    if identifier == canonicalIdentifier then
      pure ()
    else
      fail $ "KTX identifier mismatch: " <> show identifier

    endianness <- getWord32le
    let
      getNext =
        if endianness == endiannessLE then
          getWord32le
        else
          getWord32be

    glType                <- getNext
    glTypeSize            <- getNext
    glFormat              <- getNext
    glInternalFormat      <- getNext
    glBaseInternalFormat  <- getNext
    pixelWidth            <- getNext
    pixelHeight           <- getNext
    pixelDepth            <- getNext
    numberOfArrayElements <- getNext
    numberOfFaces         <- getNext
    numberOfMipmapLevels  <- getNext
    bytesOfKeyValueData   <- getNext

    pure Header{..}

  put Header{..} = do
    putByteString identifier
    let putWord32 = mkPutWord32 endianness

    putWord32 endianness
    putWord32 glType
    putWord32 glTypeSize
    putWord32 glFormat
    putWord32 glInternalFormat
    putWord32 glBaseInternalFormat
    putWord32 pixelWidth
    putWord32 pixelHeight
    putWord32 pixelDepth
    putWord32 numberOfArrayElements
    putWord32 numberOfFaces
    putWord32 numberOfMipmapLevels
    putWord32 bytesOfKeyValueData

endiannessLE :: Word32
endiannessLE = 0x04030201

canonicalIdentifier :: ByteString
canonicalIdentifier = BS.pack
  [ 0xAB, 0x4B, 0x54, 0x58, 0x20, 0x31, 0x31, 0xBB -- «KTX 11»
  , 0x0D, 0x0A, 0x1A, 0x0A                         -- \r\n\x1A\n
  ]

-- * Images

type MipLevels = Vector MipLevel

data MipLevel = MipLevel
  { imageSize     :: Word32
  , arrayElements :: Vector ArrayElement
  }
  deriving (Show, Generic)

newtype ArrayElement = ArrayElement
  { faces :: Vector Face
  }
  deriving (Show, Generic)

newtype Face = Face
  { zSlices :: Vector ZSlice
  }
  deriving (Show, Generic)

newtype ZSlice = ZSlice
  { block :: ByteString
  }
  deriving (Generic)

instance Show ZSlice where
  show ZSlice{..} =
    let
      size = BS.length block
    in
      mconcat
        [ "ZSlice ("
        , show size
        , ") "
        , show (BS.take 32 block)
        ]

getImages :: Header -> Get MipLevels
getImages Header{..} =
  some_ numberOfMipmapLevels' do
    imageSize <- getImageSize

    let
      sliceSize = fromIntegral $
        if numberOfFaces == 6 then
          imageSize
        else
          imageSize
            `div` numberOfArrayElements'
            `div` numberOfFaces
            `div` pixelDepth'

    elements <- some_ numberOfArrayElements' $
      some_ numberOfFaces $
        some_ pixelDepth' $
          ZSlice <$> getByteString sliceSize

    pure MipLevel
      { imageSize     = imageSize
      , arrayElements = coerce elements
      }

  where
    some_ n action = Vector.forM (Vector.fromList [1..n]) \_ix -> action

    numberOfMipmapLevels'
      | numberOfMipmapLevels == 0 = 1
      | otherwise                 = numberOfMipmapLevels

    numberOfArrayElements'
      | numberOfArrayElements == 0 = 1
      | otherwise                  = numberOfArrayElements

    pixelDepth'
      | pixelDepth == 0 = 1
      | otherwise       = pixelDepth

    getImageSize =
      if endianness == endiannessLE then
        getWord32le
      else
        getWord32be

putImages :: (Word32 -> Put) -> MipLevels -> Put
putImages putWord32 mipLevels = Vector.forM_ mipLevels \MipLevel{..} -> do
  putWord32 imageSize
  Vector.forM_ arrayElements \ArrayElement{..} ->
    Vector.forM_ faces \Face{..} ->
      Vector.forM_ zSlices \ZSlice{..} ->
        putByteString block

-- * Utils

mkGetWord32 :: Word32 -> Get Word32
mkGetWord32 someEndianness =
  if someEndianness == endiannessLE then
    getWord32le
  else
    getWord32be

mkPutWord32 :: Word32 -> (Word32 -> Put)
mkPutWord32 someEndianness =
  if someEndianness == endiannessLE then
    putWord32le
  else
    putWord32be