{-# LANGUAGE
    DeriveDataTypeable
  , ExistentialQuantification
  , UnicodeSyntax
  #-}
-- | WavPack metadata sub-blocks
module Codec.Audio.WavPack.Metadata
    ( Metadata(..)
    , SubBlock

    , Dummy(..)
    , DecorrTerms(..)
    , DecorrWeights(..)
    , DecorrSamples(..)
    , EntropyVars(..)
    , RIFFHeader(..)
    , RIFFTrailer(..)
    , Unknown(..)
    )
    where
import Codec.Audio.WavPack.Internal
import Control.Monad
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import Data.Bits
import qualified Data.ByteString.Lazy as L
import Data.Int
import qualified Data.Strict as S
import Data.Typeable
import qualified Data.Vector.Unboxed as UV
import Prelude.Unicode

-- | Type class for every metadata sub-blocks.
class (Binary α, Eq α, Show α, Typeable α) ⇒ Metadata α where
    -- | Get the metadata ID without odd-size bit nor large-block bit
    -- (mandatory).
    metaID   ∷ α → Word8
    -- | Get the size of metadata, excluding the metadata header
    -- (optional).
    metaSize ∷ α → Word32
    metaSize = fromIntegral ∘ L.length ∘ runPut ∘ put
    -- | Cast a 'SubBlock' to this type of metadata (optional).
    fromSubBlock ∷ SubBlock → Maybe α
    fromSubBlock (SubBlock a) = cast a
    -- | Wrap the metadata into 'SubBlock' (optional).
    toSubBlock ∷ α → SubBlock
    toSubBlock = SubBlock

-- | An opaque sub-block container.
data SubBlock = ∀α. Metadata α ⇒ SubBlock α
                deriving Typeable

instance Metadata SubBlock where
    metaID   (SubBlock a) = metaID a
    metaSize (SubBlock a) = metaSize a
    fromSubBlock = Just
    toSubBlock   = id

instance Binary SubBlock where
    put (SubBlock a)
        = let size     = metaSize a
              size'    = size + 1
              oddBit   = if odd size     then 0x40 else 0
              largeBit = if size > 0x1FE then 0x80 else 0
              idWord   = metaID a .|. oddBit .|. largeBit
          in
            do putWord8 idWord
               if size > 0x1FE then
                   -- Don't forget about the endianness.
                   do putWord8 $ fromIntegral $ (size' `shiftR`  1) .&. 0xFF
                      putWord8 $ fromIntegral $ (size' `shiftR`  9) .&. 0xFF
                      putWord8 $ fromIntegral $ (size' `shiftR` 17) .&. 0xFF
                 else
                      putWord8 $ fromIntegral $ (size' `shiftR`  1) .&. 0xFF
               put a
               when (odd size) $ putWord8 0

    get = do idWord ← getWord8
             let isOdd   = idWord .&. 0x40 ≢ 0
                 isLarge = idWord .&. 0x80 ≢ 0
                 rawID   = idWord .&. (complement 0x40) .&. (complement 0x80)
                 adj     = if isOdd then -1 else 0
             size ← if isLarge then
                         do sz0 ← getWord8
                            sz1 ← getWord8
                            sz2 ← getWord8
                            return $ ( (fromIntegral sz2 `shiftL` 17) .|.
                                       (fromIntegral sz1 `shiftL`  9) .|.
                                       (fromIntegral sz0 `shiftL`  1)
                                     ) + adj
                     else
                         fmap ((+ adj) ∘ (`shiftL` 1) ∘ fromIntegral) getWord8
             subb ← getLazyByteString $ fromIntegral (size ∷ Word32)
             return $ runGet (getSubBlock rawID) subb
        where
          getSubBlock ∷ Word8 → Get SubBlock
          getSubBlock 0x00 = fmap SubBlock (get ∷ Get Dummy        )
          getSubBlock 0x02 = fmap SubBlock (get ∷ Get DecorrTerms  )
          getSubBlock 0x03 = fmap SubBlock (get ∷ Get DecorrWeights)
          getSubBlock 0x04 = fmap SubBlock (get ∷ Get DecorrSamples)
          getSubBlock 0x05 = fmap SubBlock (get ∷ Get EntropyVars  )
          getSubBlock 0x21 = fmap SubBlock (get ∷ Get RIFFHeader   )
          getSubBlock 0x22 = fmap SubBlock (get ∷ Get RIFFTrailer  )
          getSubBlock unknownID
              = if unknownID .&. 0x20 ≡ 0 then
                    fail ("Unknown WavPack metadata ID: " ⧺ show unknownID)
                else
                    -- It's unknown but optional. We can safely ignore it.
                    fmap (SubBlock ∘ Unknown unknownID) getRemainingLazyByteString

instance Eq SubBlock where
    (SubBlock a) == (SubBlock b)
        = Just a ≡ cast b

instance Show SubBlock where
    show (SubBlock a) = show a

-- | Dummy metadata to pad WavPack blocks.
data Dummy
    = Dummy {
        -- | Must be less than 2^25 bytes long due to the limitation
        -- of WavPack specification.
        dumSize ∷ Word32
      }
    deriving (Eq, Show, Typeable)

instance Metadata Dummy where
    metaID _ = 0x00
    metaSize = dumSize

instance Binary Dummy where
    put = putLazyByteString ∘ flip L.replicate 0x00 ∘ fromIntegral ∘ dumSize
    get = fmap (Dummy ∘ fromIntegral) remaining

-- | Decorrelation terms and deltas.
data DecorrTerms
    = DecorrTerms {
        -- | @[ (term, delta) ]@
        dectVec ∷ !(UV.Vector (Int8, Int8))
      }
    deriving (Eq, Show, Typeable)

instance Metadata DecorrTerms where
    metaID _ = 0x02
    metaSize = fromIntegral ∘ UV.length ∘ dectVec

instance Binary DecorrTerms where
    put = UV.mapM_ (putWord8 ∘ packDT) ∘ dectVec
        where
          packDT ∷ (Int8, Int8) → Word8
          packDT (term, δ)
              = fromIntegral ( (term + 5 .&. 0x1F)
                               .|.
                               ((δ `shiftL` 5) .&. 0xE0)
                             )
    get = do n   ← remaining
             vec ← UV.replicateM (fromIntegral n) $ fmap unpackDT getWord8
             -- THINKME: unpack.c(read_decorr_terms) reverses the
             -- order only when they are decoding. I don't know why so
             -- I leave it unreversed for now.
             return $ DecorrTerms vec
        where
          unpackDT ∷ Word8 → (Int8, Int8)
          unpackDT w
              = let term = (fromIntegral $ w .&. 0x1F) - 5
                    δ    = fromIntegral $ (w `shiftR` 5) .&. 0x07
                in
                  (term, δ)

-- | Decorrelation weights.
data DecorrWeights
    = DecorrWeights {
      -- | For mono blocks, this is a weight vector for the single
      --   channel. For stereo blocks, it's interleaved as A, B, A, B,
      --   ...
        decwVec ∷ !(UV.Vector Int16)
      }
    deriving (Eq, Show, Typeable)

instance Metadata DecorrWeights where
    metaID _ = 0x03
    metaSize = fromIntegral ∘ UV.length ∘ decwVec

instance Binary DecorrWeights where
    put = UV.mapM_ (putWord8 ∘ packWeight) ∘ decwVec
    get = do n   ← remaining
             vec ← UV.replicateM (fromIntegral n)
                    $ fmap unpackWeight getWord8
             -- THINKME: the same caution as DecorrTerms, but never
             -- try to simply reverse the vector. Think about the
             -- interleaving.
             return $ DecorrWeights vec

-- | Decorrelation samples
data DecorrSamples
    = DecorrSamples {
      -- | The decorrelation sample vector stored in the metadata
      --   as-is. Actual interpretation of the vector depends on the
      --   number of channels and each corresponding decorrelation
      --   terms.
        decsVec ∷ !(UV.Vector Int32)
      }
    deriving (Eq, Show, Typeable)

instance Metadata DecorrSamples where
    metaID _ = 0x04
    metaSize = fromIntegral ∘ (⋅ 2) ∘ UV.length ∘ decsVec

instance Binary DecorrSamples where
    put = UV.mapM_ (putWord16le ∘ fromIntegral ∘ log2s) ∘ decsVec
    get = do n   ← remaining
             vec ← UV.replicateM (fromIntegral $ n `div` 2)
                    $ fmap (exp2s ∘ fromIntegral) getWord16le
             return $ DecorrSamples vec

-- | Median log2 values.
data EntropyVars
    = EntropyVars {
      -- | Median log2 values for channel A, which always exists.
        entVarA ∷ !(Word32, Word32, Word32)
      -- | Median log2 values for channel B, which is absent when it's
      --   mono.
      , entVarB ∷ !(S.Maybe (Word32, Word32, Word32))
      }
    deriving (Eq, Show, Typeable)

instance Metadata EntropyVars where
    metaID _ = 0x05
    metaSize ev
        | S.isNothing $ entVarB ev =  6
        | otherwise                = 12

instance Binary EntropyVars where
    put ev
        = do putMedians $ entVarA ev
             case entVarB ev of
               S.Nothing    → return ()
               S.Just medsB → putMedians medsB
        where
          -- THINKME: words.c(write_entropy_vars) is a destructive
          -- subroutine. It calls read_entropy_vars() to read the
          -- values back to compensate for the loss through the log
          -- function.
          putMedians ∷ (Word32, Word32, Word32) → Put
          putMedians (med0, med1, med2)
              = do putWord16le $ log2 med0
                   putWord16le $ log2 med1
                   putWord16le $ log2 med2

    get = do medsA ← getMedians
             medsB ← do isMono ← isEmpty
                        if isMono then
                            return S.Nothing
                          else
                            fmap S.Just getMedians
             return $! EntropyVars medsA medsB
        where
          getMedians ∷ Get (Word32, Word32, Word32)
          getMedians
              = do med0 ← fmap exp2 getWord16le
                   med1 ← fmap exp2 getWord16le
                   med2 ← fmap exp2 getWord16le
                   return (med0, med1, med2)

-- | RIFF header for .wav files (before audio)
data RIFFHeader
    = RIFFHeader {
        riffHeader ∷ L.ByteString
      }
    deriving (Eq, Show, Typeable)

instance Metadata RIFFHeader where
    metaID _ = 0x21

instance Binary RIFFHeader where
    put = putLazyByteString ∘ riffHeader
    get = fmap RIFFHeader getRemainingLazyByteString

-- | RIFF trailer for .wav files (after audio)
data RIFFTrailer
    = RIFFTrailer {
        riffTrailer ∷ L.ByteString
      }
    deriving (Eq, Show, Typeable)

instance Metadata RIFFTrailer where
    metaID _ = 0x22

instance Binary RIFFTrailer where
    put = putLazyByteString ∘ riffTrailer
    get = fmap RIFFTrailer getRemainingLazyByteString

-- | Unknown but optional metadata found in the WavPack block.
data Unknown
    = Unknown {
        -- | The ID of this unknown metadata without odd-size bit nor
        -- large-block bit.
        unkID   ∷ Word8
        -- | Raw data; must be less than 2^25 bytes long.
      , unkData ∷ L.ByteString
      }
    deriving (Eq, Show, Typeable)

instance Metadata Unknown where
    metaID   = unkID

instance Binary Unknown where
    put = putLazyByteString ∘ unkData
    get = error "unsupported operation"