Write.hs

module Codec.Ktx2.Write
  ( toFile
  , toChunks
  ) where

import Control.Monad (unless)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Binary (Binary(..))
import Data.Binary.Put (runPut, putLazyByteString, putByteString)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.Foldable (traverse_)
import Data.Map qualified as Map
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Word (Word64)

import Codec.Ktx.KeyValue (KeyValueData)
import Codec.Ktx.KeyValue qualified as KeyValueData
import Codec.Ktx2.DFD (DFD(..))
import Codec.Ktx2.DFD qualified as DFD
import Codec.Ktx2.Header (Header)
import Codec.Ktx2.Header qualified as Header
import Codec.Ktx2.Level qualified as Level

toFile
  :: MonadIO io
  => FilePath
  -> Header
  -> Vector DFD.Block
  -> KeyValueData
  -> ByteString
  -> [(Maybe Word64, ByteString)]
  -> io ()
toFile path headerBase dfdBlocks kvd sgd levels =
  liftIO . BSL.writeFile path $
    toChunks headerBase dfdBlocks kvd sgd levels

toChunks
  :: Header
  -> Vector DFD.Block
  -> KeyValueData
  -> ByteString
  -> [(Maybe Word64, ByteString)]
  -> BSL.ByteString
toChunks headerBase dfdBlocks kvd sgd levels =
  runPut do
    put header
    Vector.mapM_ put levelIndex
    put dfd

    putLazyByteString kvdBytes

    unless (BS.null sgd) do
      putByteString $ BS.replicate sgdPadding 0x00
      putByteString sgd

    traverse_ (putByteString . snd) (reverse levels)
  where
    header = headerBase
      { Header.levelCount =
          fromIntegral levelCount

      , Header.dfdByteOffset =
          fromIntegral dfdOffset
      , Header.dfdByteLength =
          fromIntegral dfdLength

      , Header.kvdByteOffset =
          fromIntegral kvdOffset
      , Header.kvdByteLength =
          fromIntegral kvdLength

      , Header.sgdByteOffset =
          fromIntegral sgdOffset
      , Header.sgdByteLength =
          fromIntegral sgdLength
      }

    levelIndex = Level.index levelBaseOffset levels
    levelCount = Vector.length levelIndex
    levelIndexOffset = 80
    levelIndexLength = levelCount * 8 * 3
    levelIndexEnd = levelIndexOffset + levelIndexLength

    (dfdOffset, dfdLength) =
      if Vector.null dfdBlocks then
        (0, 0)
      else
        ( levelIndexEnd
        , fromIntegral $ dfdTotalSize dfd
        )
    dfd = DFD
      { dfdTotalSize =
          4 + Vector.sum (Vector.map DFD.descriptorBlockSize dfdBlocks)
      , dfdBlocks = dfdBlocks
      }
    dfdEnd = levelIndexEnd + dfdLength

    (kvdOffset, kvdLength) =
      if Map.null kvd then
        (0, 0)
      else
        ( dfdEnd
        , fromIntegral $ BSL.length kvdBytes
        )
    kvdBytes = runPut $ KeyValueData.putDataLe kvd
    kvdEnd = dfdEnd + kvdLength

    (sgdPadding, sgdOffset, sgdLength) =
      if BS.null sgd then
        (0, 0, 0)
      else
        ( 7 - ((kvdEnd + 7) `rem` 8)
        , sgdPadding + kvdEnd
        , BS.length sgd
        )
    sgdEnd = kvdEnd + sgdLength

    levelBaseOffset = fromIntegral sgdEnd