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 path = liftIO . BSL.writeFile path . toChunks

toChunks :: Ktx2 -> BSL.ByteString
toChunks Ktx2.Ktx2{header=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