module Codec.EBML.Get where

import Data.Binary.Get (Get, bytesRead, getByteString, isEmpty, lookAheadM)
import Data.Text.Encoding (decodeUtf8)
import Data.Word (Word64)

import Codec.EBML.Element
import Codec.EBML.Schema
import Data.Bits (Bits)

getElement :: EBMLSchemas -> Get EBMLElement
getElement :: EBMLSchemas -> Get EBMLElement
getElement EBMLSchemas
schemas = do
    EBMLElementHeader
elth <- Get EBMLElementHeader
getElementHeader
    EBMLSchemas -> EBMLElementHeader -> Get EBMLElement
getElementValue EBMLSchemas
schemas EBMLElementHeader
elth

getElementValue :: EBMLSchemas -> EBMLElementHeader -> Get EBMLElement
getElementValue :: EBMLSchemas -> EBMLElementHeader -> Get EBMLElement
getElementValue EBMLSchemas
schemas EBMLElementHeader
elth = do
    -- here is a good place to add traceM debug
    EBMLValue
val <- case EBMLID -> EBMLSchemas -> Maybe EBMLSchema
lookupSchema EBMLElementHeader
elth.eid EBMLSchemas
schemas of
        Maybe EBMLSchema
Nothing -> EBMLElementHeader -> Get EBMLValue
getBinary EBMLElementHeader
elth
        Just EBMLSchema
schema -> EBMLSchema
schema.decode EBMLSchemas
schemas EBMLElementHeader
elth
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ EBMLElementHeader -> EBMLValue -> EBMLElement
EBMLElement EBMLElementHeader
elth EBMLValue
val

getDocument :: EBMLSchemas -> Get EBMLDocument
getDocument :: EBMLSchemas -> Get EBMLDocument
getDocument EBMLSchemas
schemas = [EBMLElement] -> EBMLDocument
EBMLDocument forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get [EBMLElement]
go
  where
    go :: Get [EBMLElement]
go = do
        Bool
empty <- Get Bool
isEmpty
        if Bool
empty
            then forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            else do
                EBMLElement
elt <- EBMLSchemas -> Get EBMLElement
getElement EBMLSchemas
schemas
                [EBMLElement]
elts <- Get [EBMLElement]
go
                forall (f :: * -> *) a. Applicative f => a -> f a
pure (EBMLElement
elt forall a. a -> [a] -> [a]
: [EBMLElement]
elts)

getBinary :: EBMLElementHeader -> Get EBMLValue
getBinary :: EBMLElementHeader -> Get EBMLValue
getBinary EBMLElementHeader
elth = case EBMLElementHeader
elth.size of
    Maybe Word64
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Invalid binary header size"
    Just Word64
sz -> ByteString -> EBMLValue
EBMLBinary forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
getByteString (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sz)

getText :: EBMLElementHeader -> Get EBMLValue
getText :: EBMLElementHeader -> Get EBMLValue
getText EBMLElementHeader
elth = case EBMLElementHeader
elth.size of
    Maybe Word64
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Invalid text header size"
    Just Word64
sz -> Text -> EBMLValue
EBMLText forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
decodeUtf8 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
getByteString (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sz)

getUnsignedInteger :: EBMLElementHeader -> Get EBMLValue
getUnsignedInteger :: EBMLElementHeader -> Get EBMLValue
getUnsignedInteger EBMLElementHeader
elth = Word64 -> EBMLValue
EBMLUnsignedInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. (Bits a, Integral a) => Maybe Word64 -> Get a
getInt EBMLElementHeader
elth.size

getInteger :: EBMLElementHeader -> Get EBMLValue
getInteger :: EBMLElementHeader -> Get EBMLValue
getInteger EBMLElementHeader
elth = Word64 -> EBMLValue
EBMLUnsignedInteger forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. (Bits a, Integral a) => Maybe Word64 -> Get a
getInt EBMLElementHeader
elth.size

getInt :: (Bits a, Integral a) => Maybe Word64 -> Get a
getInt :: forall a. (Bits a, Integral a) => Maybe Word64 -> Get a
getInt Maybe Word64
size = forall a. (Num a, Bits a) => Int -> a -> Get a
getVar Int
sz a
0
  where
    -- TODO: check the value is in the [0..8] range
    sz :: Int
sz = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 forall a b. (Integral a, Num b) => a -> b
fromIntegral Maybe Word64
size

getRoot :: EBMLSchemas -> EBMLElementHeader -> Get EBMLValue
getRoot :: EBMLSchemas -> EBMLElementHeader -> Get EBMLValue
getRoot EBMLSchemas
schemas EBMLElementHeader
elth = case EBMLElementHeader
elth.size of
    Maybe Word64
Nothing -> [EBMLElement] -> EBMLValue
EBMLRoot forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EBMLSchemas -> EBMLID -> Get [EBMLElement]
getUntil EBMLSchemas
schemas EBMLElementHeader
elth.eid
    Just Word64
sz -> EBMLSchemas -> Word64 -> Get EBMLValue
getRootFixed EBMLSchemas
schemas Word64
sz

getUntil :: EBMLSchemas -> EBMLID -> Get [EBMLElement]
getUntil :: EBMLSchemas -> EBMLID -> Get [EBMLElement]
getUntil EBMLSchemas
schemas EBMLID
eid = Get [EBMLElement]
go
  where
    getChild :: Get (Maybe EBMLElement)
    getChild :: Get (Maybe EBMLElement)
getChild = do
        -- This is not exactly correct. The rfc-8794 spec (chapter 6.2) says we should decode until
        -- any valid parent or global element. Because the EBMLSchema doesn't yet contain this information,
        -- and because in practice such unknown-sized element are segment/cluster, we simply decode until
        -- we find a matching element.
        EBMLElementHeader
elth <- Get EBMLElementHeader
getElementHeader
        if EBMLElementHeader
elth.eid forall a. Eq a => a -> a -> Bool
== EBMLID
eid
            then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
            else forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EBMLSchemas -> EBMLElementHeader -> Get EBMLElement
getElementValue EBMLSchemas
schemas EBMLElementHeader
elth

    go :: Get [EBMLElement]
go = do
        Bool
empty <- Get Bool
isEmpty
        if Bool
empty
            then forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            else Get [EBMLElement]
goGet

    goGet :: Get [EBMLElement]
goGet =
        forall a. Get (Maybe a) -> Get (Maybe a)
lookAheadM Get (Maybe EBMLElement)
getChild forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Just EBMLElement
elt -> do
                [EBMLElement]
elts <- Get [EBMLElement]
go
                forall (f :: * -> *) a. Applicative f => a -> f a
pure (EBMLElement
elt forall a. a -> [a] -> [a]
: [EBMLElement]
elts)
            Maybe EBMLElement
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure []

getRootFixed :: EBMLSchemas -> Word64 -> Get EBMLValue
getRootFixed :: EBMLSchemas -> Word64 -> Get EBMLValue
getRootFixed EBMLSchemas
schemas Word64
sz = do
    Int64
startPosition <- Get Int64
bytesRead
    let maxPosition :: Int64
maxPosition = Int64
startPosition forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
sz
        getChilds :: Get [EBMLElement]
getChilds = do
            Int64
currentPosition <- Get Int64
bytesRead
            if
                    | Int64
currentPosition forall a. Ord a => a -> a -> Bool
> Int64
maxPosition ->
                        forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Element decode position " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Int64
currentPosition forall a. Semigroup a => a -> a -> a
<> String
" exceed parent size " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Word64
sz
                    | Int64
currentPosition forall a. Eq a => a -> a -> Bool
== Int64
maxPosition ->
                        forall (f :: * -> *) a. Applicative f => a -> f a
pure []
                    | Bool
otherwise -> do
                        EBMLElement
elt <- EBMLSchemas -> Get EBMLElement
getElement EBMLSchemas
schemas
                        [EBMLElement]
elts <- Get [EBMLElement]
getChilds
                        forall (f :: * -> *) a. Applicative f => a -> f a
pure (EBMLElement
elt forall a. a -> [a] -> [a]
: [EBMLElement]
elts)
    [EBMLElement] -> EBMLValue
EBMLRoot forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get [EBMLElement]
getChilds