BasicV2.hs

module Codec.Ktx2.DFD.Khronos.BasicV2 where

import Data.Binary (Binary(..))
import Data.Binary.Get (Get, getWord8, getWord16le, getWord32le, runGetOrFail)
import Data.Binary.Put (PutM, putWord8, putWord16le, putWord32le, runPut)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.ByteString.Lazy qualified as LBS
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Word (Word8, Word16, Word32)
import GHC.Generics (Generic)

import Codec.Ktx2.DFD qualified as DFD

-- | Khronos
pattern VENDOR_ID :: (Eq a, Num a) => a
pattern VENDOR_ID = 0

-- | Basic DFD Block
pattern DESCRIPTOR_TYPE :: (Eq a, Num a) => a
pattern DESCRIPTOR_TYPE = 0

-- | KDF v1.3
pattern VERSION :: (Eq a, Num a) => a
pattern VERSION = 2

{- |
  A basic descriptor block is designed to encode common metadata associated with bulk data — especially image or texture data.

  While this descriptor holds more information about the data interpretation than is needed by many applications,
  a comprehensive encoding reduces the risk of metadata needed by different APIs being lost in translation.

  The format is described in terms of a repeating axis-aligned texel block composed of samples.
  Each sample contains a single channel of information with a single spatial offset within the texel block,
  and consists of an amount of contiguous data. This descriptor block consists of information about the interpretation
  of the texel block as a whole, supplemented by a description of a number of samples taken from one or more planes of contiguous memory.

  <https://registry.khronos.org/DataFormat/specs/1.3/dataformat.1.3.html>
-}
data BasicV2 = BasicV2
  { colorModel :: Word8
  , colorPrimaries :: Word8
  , transferFunction :: Word8
  , flags :: Word8

  , texelBlockDimension0 :: Word8
    {- ^
      The value held in each of these fields is one fewer than the size of the block in that dimension — 
      that is, a value of 0 represents a size of 1, a value of 1 represents a size of 2, etc.

      A texel block which covers fewer than four dimensions should have a size of 1 in each dimension
      that it lacks, and therefore the corresponding fields in the representation should be 0.
    -}
  , texelBlockDimension1 :: Word8
  , texelBlockDimension2 :: Word8
  , texelBlockDimension3 :: Word8

  , bytesPlane0 :: Word8
  , bytesPlane1 :: Word8
  , bytesPlane2 :: Word8
  , bytesPlane3 :: Word8
  , bytesPlane4 :: Word8
  , bytesPlane5 :: Word8
  , bytesPlane6 :: Word8
  , bytesPlane7 :: Word8

  , samples :: Vector Sample
  }
  deriving (Eq, Show, Generic)

getter :: Int -> Get BasicV2
getter numSamples = do
  colorModel <- getWord8
  colorPrimaries <- getWord8
  transferFunction <- getWord8
  flags <- getWord8

  texelBlockDimension0 <- fmap (+ 1) getWord8
  texelBlockDimension1 <- fmap (+ 1) getWord8
  texelBlockDimension2 <- fmap (+ 1) getWord8
  texelBlockDimension3 <- fmap (+ 1) getWord8

  bytesPlane0 <- getWord8
  bytesPlane1 <- getWord8
  bytesPlane2 <- getWord8
  bytesPlane3 <- getWord8
  bytesPlane4 <- getWord8
  bytesPlane5 <- getWord8
  bytesPlane6 <- getWord8
  bytesPlane7 <- getWord8

  samples <- Vector.replicateM numSamples get

  pure BasicV2{..}

putter :: BasicV2 -> PutM ()
putter BasicV2{..} = do
  putWord8 colorModel
  putWord8 colorPrimaries
  putWord8 transferFunction
  putWord8 flags

  putWord8 $ texelBlockDimension0 - 1
  putWord8 $ texelBlockDimension1 - 1
  putWord8 $ texelBlockDimension2 - 1
  putWord8 $ texelBlockDimension3 - 1

  putWord8 bytesPlane0
  putWord8 bytesPlane1
  putWord8 bytesPlane2
  putWord8 bytesPlane3
  putWord8 bytesPlane4
  putWord8 bytesPlane5
  putWord8 bytesPlane6
  putWord8 bytesPlane7

  Vector.mapM_ put samples

data Sample = Sample
  { bitOffset :: Word16
  , bitLength :: Word8
  , channelType :: Word8

  , samplePosition0 :: Word8
  , samplePosition1 :: Word8
  , samplePosition2 :: Word8
  , samplePosition3 :: Word8

  , sampleLower :: Word32
  , sampleUpper :: Word32
  }
  deriving (Eq, Show, Generic)

instance Binary Sample where
  get = do
    bitOffset <- getWord16le
    bitLength <- fmap (+1) getWord8
    channelType <- getWord8

    samplePosition0 <- getWord8
    samplePosition1 <- getWord8
    samplePosition2 <- getWord8
    samplePosition3 <- getWord8

    sampleLower <- getWord32le
    sampleUpper <- getWord32le

    pure Sample{..}

  put Sample{..} = do
    putWord16le bitOffset
    putWord8 $ bitLength - 1
    putWord8 channelType

    putWord8 samplePosition0
    putWord8 samplePosition1
    putWord8 samplePosition2
    putWord8 samplePosition3

    putWord32le sampleLower
    putWord32le sampleUpper

fromBlock :: DFD.Block -> Maybe BasicV2
fromBlock DFD.Block{descriptorBlockSize, body} =
  case runGetOrFail (getter numSamples) (LBS.fromChunks [body]) of
    Left _err ->
      Nothing
    Right (_bytes, _offset, ok) ->
      Just ok
  where
    numSamples =
      div (fromIntegral descriptorBlockSize - 24) 16

toBlock :: BasicV2 -> DFD.Block
toBlock v2 =
  DFD.Block
    { descriptorType =
        DESCRIPTOR_TYPE
    , vendorId =
        VENDOR_ID
    , versionNumber =
        VERSION
    , descriptorBlockSize =
        fromIntegral $ BS.length body + 8
    , body =
        body
    }
  where
    body = BSL.toStrict . runPut $ putter v2