Ktx.hs

module Codec.Ktx where

import Data.Binary (Binary(..), decodeFileOrFail, decodeOrFail)
import Data.Binary.Get (Get, ByteOffset, getWord32le, getWord32be, getByteString, isolate, skip)
import Data.Binary.Put (Put, execPut, putByteString, putWord32le, putWord32be)
import Data.ByteString (ByteString)
import Data.ByteString.Builder (Builder, hPutBuilder)
import Data.Coerce (coerce)
import Data.Foldable (for_)
import Data.Map.Strict (Map)
import Data.Text (Text)
import Data.Vector (Vector)
import Data.Word (Word32)
import GHC.Generics (Generic)
import System.IO (IOMode(..), withBinaryFile)

import qualified Data.Text.Encoding as Text
import qualified Data.Map.Strict as Map
import qualified Data.Vector as Vector
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL

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 <- getKeyValueData header
    images <- getImages header
    pure Ktx{..}

  put Ktx{..} = do
    put header
    putKeyValueData 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
  ]

-- * Key-value data

type KeyValueData = Map Key Value

newtype Key = Key Text
  deriving (Eq, Ord, Show, Generic)

newtype Value = Value ByteString
  deriving (Show, Generic)

getKeyValueData :: Header -> Get KeyValueData
getKeyValueData Header{..} =
  isolate (fromIntegral bytesOfKeyValueData) $
    go (toInteger bytesOfKeyValueData) []
  where
    go remains acc
      | remains == 0 =
          pure $ Map.fromList acc

      | remains < 0 =
          fail "Attempted to read beyond bytesOfKeyValueData"

      | otherwise = do
          keyAndValueByteSize <- fmap toInteger getSize
          let paddingSize = 3 - ((keyAndValueByteSize + 3) `rem` 4)

          keyAndValue <- getByteString (fromInteger keyAndValueByteSize)
          skip (fromInteger paddingSize)

          {- XXX: Spec says:
              Any byte value is allowed.
              It is encouraged that the value be a NUL terminated UTF-8 string but this is not required.
              If the Value data is a string of bytes then the NUL termination
              should be included in the keyAndValueByteSize byte count
              (but programs that read KTX files must not rely on this).
          -}
          let
            (keyBS, valueBS) = BS.span (/= 0x00) keyAndValue
            key   = Key $ Text.decodeUtf8 keyBS
            value = Value $ BS.drop 1 valueBS

          go
            (remains - 4 - keyAndValueByteSize - paddingSize)
            ((key, value) : acc)

    getSize =
      if endianness == endiannessLE then
        getWord32le
      else
        getWord32be

putKeyValueData :: (Word32 -> Put) -> Map Key Value -> Put
putKeyValueData endianPut kvs =
  for_ (Map.toList kvs) \(Key key, Value value) -> do
    let
      keyAndValue = Text.encodeUtf8 key <> BS.singleton 0x00 <> value
      keyAndValueByteSize = BS.length keyAndValue
      paddingSize = 3 - ((keyAndValueByteSize + 3) `rem` 4)

    endianPut (fromIntegral keyAndValueByteSize)
    putByteString keyAndValue
    putByteString $ BS.replicate paddingSize 0

-- * 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

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