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

    , Dummy(..)
    , DecorrTerms(..)
    , DecorrWeights(..)
    , RIFFHeader(..)
    , RIFFTrailer(..)
    , Unknown(..)
    )
    where
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 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 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 ∘ packW) ∘ decwVec
        where
          packW ∷ Int16 → Word8
          packW w
              = let w'   | w  >  1024 =  1024
                         | w  < -1024 = -1024
                         | otherwise  =     w
                    w''  | w' >     0 = w' - ((w' + 64) `shiftR` 7)
                         | otherwise  = w'
                    w'''              = (w'' + 4) `shiftR` 3
                in
                  fromIntegral w'''
    get = do n   ← remaining
             vec ← UV.replicateM (fromIntegral n) $ fmap unpackW getWord8
             -- THINKME: the same caution as DecorrTerms, but never
             -- try to reverse the vector simply. Think about the
             -- interleaving.
             return $ DecorrWeights vec
        where
          unpackW ∷ Word8 → Int16
          unpackW w
              = let w'  ∷ Int8
                    w'  = fromIntegral w
                    w'' ∷ Int16
                    w'' = (fromIntegral w') `shiftL` 3
                    w''' | w'' > 0   = w'' + ((w'' + 64) `shiftR` 7)
                         | otherwise = w''
                in
                  w'''

-- | 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"