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 qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.Foldable (traverse_)
import Data.Map qualified as Map
import Data.Vector qualified as Vector

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

toFile :: MonadIO io => FilePath -> Ktx2 -> io ()
toFile :: forall (io :: * -> *). MonadIO io => FilePath -> Ktx2 -> io ()
toFile FilePath
path = IO () -> io ()
forall a. IO a -> io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> io ()) -> (Ktx2 -> IO ()) -> Ktx2 -> io ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> ByteString -> IO ()
BSL.writeFile FilePath
path (ByteString -> IO ()) -> (Ktx2 -> ByteString) -> Ktx2 -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ktx2 -> ByteString
toChunks

toChunks :: Ktx2 -> BSL.ByteString
toChunks :: Ktx2 -> ByteString
toChunks Ktx2.Ktx2{header :: Ktx2 -> Header
header=Header
headerBase, Vector Block
dfdBlocks :: Vector Block
dfdBlocks :: Ktx2 -> Vector Block
dfdBlocks, KeyValueData
kvd :: KeyValueData
kvd :: Ktx2 -> KeyValueData
kvd, ByteString
sgd :: ByteString
sgd :: Ktx2 -> ByteString
sgd, [(Maybe Word64, ByteString)]
levels :: [(Maybe Word64, ByteString)]
levels :: Ktx2 -> [(Maybe Word64, ByteString)]
levels} =
  Put -> ByteString
runPut do
    Header -> Put
forall t. Binary t => t -> Put
put Header
header
    (Level -> Put) -> Vector Level -> Put
forall (m :: * -> *) a b. Monad m => (a -> m b) -> Vector a -> m ()
Vector.mapM_ Level -> Put
forall t. Binary t => t -> Put
put Vector Level
levelIndex
    DFD -> Put
forall t. Binary t => t -> Put
put DFD
dfd

    ByteString -> Put
putLazyByteString ByteString
kvdBytes

    Bool -> Put -> Put
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
BS.null ByteString
sgd) do
      ByteString -> Put
putByteString (ByteString -> Put) -> ByteString -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString
BS.replicate Int
sgdPadding Word8
0x00
      ByteString -> Put
putByteString ByteString
sgd

    ((Maybe Word64, ByteString) -> Put)
-> [(Maybe Word64, ByteString)] -> Put
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (ByteString -> Put
putByteString (ByteString -> Put)
-> ((Maybe Word64, ByteString) -> ByteString)
-> (Maybe Word64, ByteString)
-> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Word64, ByteString) -> ByteString
forall a b. (a, b) -> b
snd) ([(Maybe Word64, ByteString)] -> [(Maybe Word64, ByteString)]
forall a. [a] -> [a]
reverse [(Maybe Word64, ByteString)]
levels)
  where
    header :: Header
header = 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 :: Vector Level
levelIndex = Word64 -> [(Maybe Word64, ByteString)] -> Vector Level
Level.index Word64
levelBaseOffset [(Maybe Word64, ByteString)]
levels
    levelCount :: Int
levelCount = Vector Level -> Int
forall a. Vector a -> Int
Vector.length Vector Level
levelIndex
    levelIndexOffset :: Int
levelIndexOffset = Int
80
    levelIndexLength :: Int
levelIndexLength = Int
levelCount Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
    levelIndexEnd :: Int
levelIndexEnd = Int
levelIndexOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
levelIndexLength

    (Int
dfdOffset, Int
dfdLength) =
      if Vector Block -> Bool
forall a. Vector a -> Bool
Vector.null Vector Block
dfdBlocks then
        (Int
0, Int
0)
      else
        ( Int
levelIndexEnd
        , Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Int) -> Word32 -> Int
forall a b. (a -> b) -> a -> b
$ DFD -> Word32
dfdTotalSize DFD
dfd
        )
    dfd :: DFD
dfd = DFD
      { dfdTotalSize :: Word32
dfdTotalSize =
          Word32
4 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Vector Word32 -> Word32
forall a. Num a => Vector a -> a
Vector.sum ((Block -> Word32) -> Vector Block -> Vector Word32
forall a b. (a -> b) -> Vector a -> Vector b
Vector.map Block -> Word32
DFD.descriptorBlockSize Vector Block
dfdBlocks)
      , dfdBlocks :: Vector Block
dfdBlocks = Vector Block
dfdBlocks
      }
    dfdEnd :: Int
dfdEnd = Int
levelIndexEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dfdLength

    (Int
kvdOffset, Int
kvdLength) =
      if KeyValueData -> Bool
forall k a. Map k a -> Bool
Map.null KeyValueData
kvd then
        (Int
0, Int
0)
      else
        ( Int
dfdEnd
        , Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
BSL.length ByteString
kvdBytes
        )
    kvdBytes :: ByteString
kvdBytes = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ KeyValueData -> Put
KeyValueData.putDataLe KeyValueData
kvd
    kvdEnd :: Int
kvdEnd = Int
dfdEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kvdLength

    (Int
sgdPadding, Int
sgdOffset, Int
sgdLength) =
      if ByteString -> Bool
BS.null ByteString
sgd then
        (Int
0, Int
0, Int
0)
      else
        ( Int
7 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ((Int
kvdEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
8)
        , Int
sgdPadding Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kvdEnd
        , ByteString -> Int
BS.length ByteString
sgd
        )
    sgdEnd :: Int
sgdEnd = Int
kvdEnd Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sgdLength

    levelBaseOffset :: Word64
levelBaseOffset = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sgdEnd