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

    , Dummy(..)
    , DecorrTerms(..)
    , DecorrWeights(..)
    , DecorrSamples(..)
    , EntropyVars(..)
    , WVBitstream(..)
    , RIFFHeader(..)
    , RIFFTrailer(..)
    , ConfigInfo(..)
    , Unknown(..)
    )
    where
import Codec.Audio.WavPack.Internal
import Control.Monad
import Data.Binary
import Data.Binary.BitPut (putBit, putNBits, runBitPut)
import Data.Binary.Get
import Data.Binary.Put
import Data.Binary.Strict.BitGet (getBit, runBitGet)
import qualified Data.Binary.Strict.BitGet as BG
import Data.Bits
import qualified Data.ByteString as S
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 0x0A = fmap SubBlock (get ∷ Get WVBitstream  )
          getSubBlock 0x21 = fmap SubBlock (get ∷ Get RIFFHeader   )
          getSubBlock 0x22 = fmap SubBlock (get ∷ Get RIFFTrailer  )
          getSubBlock 0x25 = fmap SubBlock (get ∷ Get ConfigInfo   )
          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)

-- | WV Bitstream
data WVBitstream
    = WVBitstream {
        wvStream ∷ !L.ByteString
      }
    deriving (Eq, Show, Typeable)

instance Metadata WVBitstream where
    metaID _ = 0x0A
    metaSize = fromIntegral ∘ L.length ∘ wvStream

instance Binary WVBitstream where
    put = putLazyByteString ∘ wvStream
    get = fmap WVBitstream getRemainingLazyByteString

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

instance Metadata RIFFHeader where
    metaID _ = 0x21
    metaSize = fromIntegral ∘ L.length ∘ riffHeader

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
    metaSize = fromIntegral ∘ L.length ∘ riffTrailer

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

-- | Configuration information.
data ConfigInfo
    = ConfigInfo {
      -- | fast mode
        cfgFast           ∷ !Bool
      -- | high quality mode
      , cfgHigh           ∷ !Bool
      -- | very high
      , cfgVeryHigh       ∷ !Bool
      -- | bitrate is kbps, not bits / sample
      , cfgBitrateKbps    ∷ !Bool
      -- | automatic noise shaping
      , cfgAutoShaping    ∷ !Bool
      -- | shaping mode specified
      , cfgShapeOverride  ∷ !Bool
      -- | joint-stereo mode specified
      , cfgJointOverride  ∷ !Bool
      -- | dynamic noise shaping
      , cfgDynamicShaping ∷ !Bool
      -- | create executable
      , cfgCreateEXE      ∷ !Bool
      -- | create correction file
      , cfgCreateWVC      ∷ !Bool
      -- | maximize hybrid compression
      , cfgOptimizeWVC    ∷ !Bool
      -- | calc noise in hybrid mode
      , cfgCalcNoise      ∷ !Bool
      -- | obsolete (for information)
      , cfgLossyMode      ∷ !Bool
      -- | extra processing mode level (1-6)
      , cfgExtraModeLevel ∷ !(S.Maybe Word8)
      -- | no wvx stream w/ floats & big ints
      , cfgSkipWVX        ∷ !Bool
      -- | compute & store MD5 signature
      , cfgMD5Checksum    ∷ !Bool
      -- | merge blocks of equal redundancy
      , cfgMergeBlocks    ∷ !Bool
      -- | optimize for mono streams posing
      , cfgOptimizeMono   ∷ !Bool
      }
    deriving (Eq, Show, Typeable)

instance Metadata ConfigInfo where
    metaID _ = 0x25
    metaSize ci
        | S.isJust $ cfgExtraModeLevel ci
            = 4
        | otherwise
            = 3

instance Binary ConfigInfo where
    put ci
        = let !bs = runBitPut $
                    do putBit $ cfgOptimizeMono   ci
                       putNBits 2 (0 ∷ Word8) -- unused
                       putBit $ cfgMergeBlocks    ci
                       putBit $ cfgMD5Checksum    ci
                       putBit $ cfgSkipWVX        ci
                       putBit $ S.isJust $ cfgExtraModeLevel ci
                       putBit $ cfgLossyMode      ci
                       putBit $ cfgCalcNoise      ci
                       putNBits 2 (0 ∷ Word8) -- unused
                       putBit $ cfgOptimizeWVC    ci
                       putBit $ cfgCreateWVC      ci
                       putBit $ cfgCreateEXE      ci
                       putBit $ cfgDynamicShaping ci
                       putBit $ cfgJointOverride  ci
                       putBit $ cfgShapeOverride  ci
                       putBit $ cfgAutoShaping    ci
                       putBit $ cfgBitrateKbps    ci
                       putBit $ cfgVeryHigh       ci
                       putBit $ cfgHigh           ci
                       putBit False -- unused
                       putBit $ cfgFast           ci
                       putBit False -- unused
          in
            do putLazyByteString (L.reverse bs)
               case cfgExtraModeLevel ci of
                 S.Nothing  → return ()
                 S.Just eml → putWord8 eml

    get = do bs  ← getBytes 3
             eml ← do xmode ← fmap (¬) isEmpty
                      if xmode then
                          fmap S.Just getWord8
                        else
                          return S.Nothing
             let r = runBitGet (S.reverse bs) $
                     do optimizeMono   ← getBit
                        BG.skip 2 -- unused
                        mergeBlocks    ← getBit
                        md5Checksum    ← getBit
                        skipWVX        ← getBit
                        extraMode      ← getBit
                        lossyMode      ← getBit
                        calcNoise      ← getBit
                        BG.skip 2 -- unused
                        optimizeWVC    ← getBit
                        createWVC      ← getBit
                        createEXE      ← getBit
                        dynamicShaping ← getBit
                        jointOverride  ← getBit
                        shapeOverride  ← getBit
                        autoShaping    ← getBit
                        bitrateKbps    ← getBit
                        veryHigh       ← getBit
                        high           ← getBit
                        BG.skip 1 -- unused
                        fast           ← getBit
                        BG.skip 1 -- unused
                        return ConfigInfo {
                                     cfgFast           = fast
                                   , cfgHigh           = high
                                   , cfgVeryHigh       = veryHigh
                                   , cfgBitrateKbps    = bitrateKbps
                                   , cfgAutoShaping    = autoShaping
                                   , cfgShapeOverride  = shapeOverride
                                   , cfgJointOverride  = jointOverride
                                   , cfgDynamicShaping = dynamicShaping
                                   , cfgCreateEXE      = createEXE
                                   , cfgCreateWVC      = createWVC
                                   , cfgOptimizeWVC    = optimizeWVC
                                   , cfgCalcNoise      = calcNoise
                                   , cfgLossyMode      = lossyMode
                                   , cfgExtraModeLevel = if extraMode then
                                                             eml
                                                         else
                                                             S.Nothing
                                   , cfgSkipWVX        = skipWVX
                                   , cfgMD5Checksum    = md5Checksum
                                   , cfgMergeBlocks    = mergeBlocks
                                   , cfgOptimizeMono   = optimizeMono
                                   }
                         
             case r of
               Left err → fail err
               Right ci → return ci

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