{-# LANGUAGE OverloadedStrings, FlexibleInstances, TemplateHaskell #-}
{-# LANGUAGE RecordWildCards, BangPatterns, FlexibleContexts #-}
module Data.Avro where

import Bio.Iteratee
import Control.Applicative
import Control.Monad
import Control.Monad.ST ( runST, ST )
import Data.Aeson hiding ((.=))
import Data.Array.MArray
import Data.Array.ST ( STUArray )
import Data.Array.Unsafe ( castSTUArray )
import Data.Binary.Get
import Data.Bits
import Data.Binary.Builder
import Data.Foldable ( foldMap )
import Data.Int ( Int64 )
import Data.Maybe
import Data.Monoid
import Data.Scientific
import Data.Text.Encoding
import Data.Word ( Word32, Word64 )
import Foreign.Storable ( Storable, sizeOf )
import Language.Haskell.TH
import System.Random

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.HashMap.Strict as H
import qualified Data.ListLike as LL
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U

-- ^ Support for Avro.
-- Current status is that we can generate schemas for certain Haskell
-- values, serialize to binary and JSON representations, and write
-- Container files using the null codec.  The C implementation likes
-- some, but not all of these containers; it's unclear if that's the
-- fault of the C implementation, though.
--
-- Meanwhile, serialization works for nested sums-of-products, as long as the
-- product uses record syntax and the top level is a plain record.
-- The obvious primitives are supported.

(.=) :: ToJSON a => String -> a -> (T.Text, Value)
k .= v = (T.pack k, toJSON v)

string :: String -> Value
string = String . T.pack

-- | This is the class of types we can embed into the Avro
-- infrastructure.  Right now, we can derive a schema, encode to
-- the Avro binary format, and encode to the Avro JSON encoding.
class Avro a where
    -- | Produces the schema for this type.  Schemas are represented as
    -- JSON values.  The monad is used to keep a table of already
    -- defined types, so the schema can refer to them by name.  (The
    -- concrete argument serves to specify the type, it is not actually
    -- used.)
    toSchema :: a -> MkSchema Value

    -- | Serializes a value to the binary representation.  The schema is
    -- implied, serialization to related schemas is not supported.
    toBin :: a -> Builder

    -- | Deserializzes a value from binary representation.  Right now,
    -- no attempt at schema matching is done, the schema must match the
    -- expected one exactly.
    fromBin :: Get a

    -- | Serializes a value to the JSON representation.  Note that even
    -- the JSON format needs a schema for successful deserialization,
    -- and here we support only the one implied schema.
    toAvron :: a -> Value


newtype MkSchema a = MkSchema
    { mkSchema :: (a -> H.HashMap T.Text Value -> Value) -> H.HashMap T.Text Value -> Value }

instance Functor MkSchema where fmap f m = MkSchema (\k -> mkSchema m (k . f))
instance Applicative MkSchema where pure a = MkSchema (\k -> k a)
                                    u <*> v = MkSchema (\k -> mkSchema u (\a -> mkSchema v (k . a)))
instance Monad MkSchema where return a = MkSchema (\k -> k a)
                              a >>= m = MkSchema (\k -> mkSchema a (\a' -> mkSchema (m a') k))

memoObject :: String -> [(T.Text,Value)] -> MkSchema Value
memoObject nm ps = MkSchema $ \k h ->
    let nm' = T.pack nm
        obj = object $ ("name" .= nm) : ps
    in case H.lookup nm' h of
        Nothing -> k obj $! H.insert nm' obj h
        Just obj' | obj == obj' -> k (String nm') h
                  | otherwise -> error $ "same type name, different schema: " ++ nm

runMkSchema :: MkSchema Value -> Value
runMkSchema x = mkSchema x postproc H.empty
  where
    -- Objects are fine as is.
    postproc (Object  o) _ = Object o
    -- Top level can't be a string, can it?  Need to wrap into the long form.
    postproc (String tp) _ = object [ "type" .= String tp ]
    -- Top level Array should be fine, too.
    postproc (Array a) _ = Array a
    -- reject anything else
    postproc v _ = error $ "Not allowed as toplevel schema: " ++ show v

-- instances for primitive types

-- | The Avro \"null\" type is represented as the empty tuple.
instance Avro () where
    toSchema _ = return $ String "null"
    toBin   () = mempty
    fromBin    = return ()
    toAvron () = Null

instance Avro Bool where
    toSchema _ = return $ String "boolean"
    toBin      = singleton . fromIntegral . fromEnum
    fromBin    = toEnum . fromIntegral <$> getWord8
    toAvron    = Bool

instance Avro Int where
    toSchema _ = return $ String "long"
    toBin      = encodeIntBase128
    fromBin    = decodeIntBase128
    toAvron    = Number . fromIntegral

instance Avro Int64 where
    toSchema _ = return $ String "long"
    toBin      = encodeIntBase128
    fromBin    = decodeIntBase128
    toAvron    = Number . fromIntegral

instance Avro Float where
    toSchema _ = return $ String "float"
    toBin      = putWord32le . floatToWord
    fromBin    = wordToFloat <$> getWord32le
    toAvron    = Number . fromFloatDigits

instance Avro Double where
    toSchema _ = return $ String "double"
    toBin      = putWord64le . doubleToWord
    fromBin    = wordToDouble <$> getWord64le
    toAvron    = Number . fromFloatDigits

instance Avro B.ByteString where
    toSchema _ = return $ String "bytes"
    toBin    s = encodeIntBase128 (B.length s) <> fromByteString s
    fromBin    = decodeIntBase128 >>= getByteString
    toAvron    = String . decodeLatin1

instance Avro T.Text where
    toSchema _ = return $ String "string"
    toBin      = toBin . encodeUtf8
    fromBin    = decodeUtf8 <$> fromBin
    toAvron    = String


-- Integer<->Float conversions, stolen from cereal.

{-# INLINE wordToFloat #-}
wordToFloat :: Word32 -> Float
wordToFloat x = runST (cast x)

{-# INLINE wordToDouble #-}
wordToDouble :: Word64 -> Double
wordToDouble x = runST (cast x)

{-# INLINE floatToWord #-}
floatToWord :: Float -> Word32
floatToWord x = runST (cast x)

{-# INLINE doubleToWord #-}
doubleToWord :: Double -> Word64
doubleToWord x = runST (cast x)

{-# INLINE cast #-}
cast :: ( MArray (STUArray s) b (ST s), MArray (STUArray s) a (ST s) ) => a -> ST s b
cast x = (newArray (0 :: Int, 0) x >>= castSTUArray >>= flip readArray 0)

-- | Implements Zig-Zag-Coding like in Protocol Buffers and Avro.
zig :: (Storable a, Bits a) => a -> a
zig x = (x `shiftL` 1) `xor` (x `shiftR` (8 * sizeOf x -1))

-- | Reverses Zig-Zag-Coding like in Protocol Buffers and Avro.
zag :: (Storable a, Bits a, Num a) => a -> a
zag x = negate (x .&. 1) `xor` ((x .&. complement 1) `rotateR` 1)

-- | Encodes a word of any size using a variable length "base 128"
-- encoding.
encodeWordBase128 :: (Integral a, Bits a) => a -> Builder
encodeWordBase128 x | x' == 0   = singleton (fromIntegral (x .&. 0x7f))
                    | otherwise = singleton (fromIntegral (x .&. 0x7f .|. 0x80))
                                  <> encodeWordBase128 x'
  where x' = x `shiftR` 7

decodeWordBase128 :: (Integral a, Bits a) => Get a
decodeWordBase128 = go 0 0
  where
    go acc sc = do x <- getWord8
                   let !acc' = acc .|. fromIntegral x `shiftL` sc
                   if x .&. 0x80 == 0
                        then return acc'
                        else go acc' (sc+7)

-- | Encodes an int of any size by combining the zig-zag coding with the
-- base 128 encoding.
encodeIntBase128 :: (Integral a, Bits a, Storable a) => a -> Builder
encodeIntBase128 = encodeWordBase128 . zig

-- | Decodes an int of any size by combining the zig-zag decoding with
-- the base 128 decoding.
decodeIntBase128 :: (Integral a, Bits a, Storable a) => Get a
decodeIntBase128 = zag <$> decodeWordBase128

zigInt :: Int -> Builder
zigInt = encodeIntBase128

zagInt :: Get Int
zagInt = decodeWordBase128

-- Complex Types

-- | A list becomes an Avro array
-- The chunked encoding for lists may come in handy.  How to select the
-- chunk size is not obvious, though.
instance Avro a => Avro [a] where
    toSchema as = do sa <- toSchema (head as)
                     return $ object [ "type" .= String "array", "items" .= sa ]
    toBin    [] = singleton 0
    toBin    as = toBin (length as) <> foldMap toBin as <> singleton 0
    toAvron     = Array . V.fromList . map toAvron

    -- This is not suitable for incremental processing.
    fromBin     = get_blocks []
      where
        get_blocks acc = zagInt >>= \l -> if l == 0 then return $ reverse acc
                                                    else get_block acc l >>= get_blocks
        get_block acc l = if l == 0 then return acc
                                    else fromBin >>= \a -> get_block (a:acc) (l-1)


-- | A generic vector becomes an Avro array
instance Avro a => Avro (V.Vector a) where
    toSchema as = do sa <- toSchema (V.head as)
                     return $ object [ "type" .= String "array", "items" .= sa ]
    toBin    as | V.null as = singleton 0
                | otherwise = toBin (V.length as) <> foldMap toBin as <> singleton 0
    toAvron     = Array . V.map toAvron

    -- This is not suitable for incremental processing.
    fromBin     = get_blocks []
      where
        get_blocks acc = zagInt >>= \l -> if l == 0 then return $ V.concat $ reverse acc
                                                    else get_block [] l >>=
                                                         get_blocks . (: acc) . V.fromListN l . reverse
        get_block acc l = if l == 0 then return acc
                                    else fromBin >>= \a -> get_block (a:acc) (l-1)

-- | An unboxed vector becomes an Avro array
instance (Avro a, U.Unbox a) => Avro (U.Vector a) where
    toSchema as = do sa <- toSchema (U.head as)
                     return $ object [ "type" .= String "array", "items" .= sa ]
    toBin    as | U.null as = singleton 0
                | otherwise = toBin (U.length as) <> U.foldr ((<>) . toBin) mempty as <> singleton 0
    toAvron     = Array . V.map toAvron . U.convert

    -- This is not suitable for incremental processing.
    fromBin     = get_blocks []
      where
        get_blocks acc = zagInt >>= \l -> if l == 0 then return $ U.concat $ reverse acc
                                                    else get_block [] l >>=
                                                         get_blocks . (: acc) . U.fromListN l . reverse
        get_block acc l = if l == 0 then return acc
                                    else fromBin >>= \a -> get_block (a:acc) (l-1)


-- | A map from Text becomes an Avro map.
instance Avro a => Avro (H.HashMap T.Text a) where
    toSchema   m = do sa <- toSchema (m H.! T.empty)
                      return $ object [ "type" .= String "map", "values" .= sa ]
    toBin     as | H.null as = singleton 0
                 | otherwise = toBin (H.size as) <> H.foldrWithKey (\k v b -> toBin k <> toBin v <> b) (singleton 0) as
    toAvron      = Object . H.map toAvron

    -- This is not suitable for incremental processing.
    fromBin     = get_blocks H.empty
      where
        get_blocks !acc = zagInt >>= \l -> if l == 0 then return acc
                                                     else get_block acc l >>= get_blocks
        get_block !acc l = if l == 0 then return acc
                                     else fromBin >>= \k -> fromBin >>= \v -> get_block (H.insert k v acc) (l-1)



-- * Some(!) complex types.
--
-- Enums:  Enumerated symbols.  This is generated automatically for sums
-- of empty alternatives.  Constructor names become enum symbols.

-- Records:  This is generated automatically for product types using
-- Haskell record syntax.
--
-- Unions:  For Haskell sum-of-product types using record syntax for
-- every arm, an Avro instance resolving to a union of record can be
-- generated automatically.  The constructor names become record type
-- names, their fields become record fields.

-- XXX Sometimes we build sum types containing sum types, Maybe being the
-- most obvious example.  A (Maybe a) where a itself yields a union,
-- should probably yield a union with one more alternative (the null).


deriveAvros :: [Name] -> Q [Dec]
deriveAvros = liftM concat . mapM deriveAvro

deriveAvro :: Name -> Q [Dec]
deriveAvro nm = reify nm >>= case_info
  where
    err m = fail $ "cannot derive Avro for " ++ show nm ++ ", " ++ m

    case_info (TyConI dec) = case_dec dec
    case_info            _ = err "it is not a type constructor"

    simple_cons (NormalC _ []) = True
    simple_cons _              = False

    record_cons (RecC _ _) = True
    record_cons _          = False

    case_dec (NewtypeD _cxt _name _tyvarbndrs  _con _) = err $ "don't know what to do for NewtypeD"
    case_dec (DataD    _cxt _name _tyvarbndrs cons _)
        | all simple_cons cons = mk_enum_inst [ nm1 | NormalC nm1 [] <- cons ]
        | all record_cons cons = mk_record_inst [ (nm1, vsts) | RecC nm1 vsts <- cons ]
        | otherwise            = err $ "don't know how to make an instance with these constructors"
    case_dec _ = fail $ "is not a data or newtype declaration"

    tolit = litE . StringL . nameBase
    tolitlist (x:xs) = [| T.pack $(tolit x) : $(tolitlist xs) |]
    tolitlist [    ] = [| [] |]

    -- enum instance from list of names
    mk_enum_inst :: [Name] -> Q [Dec]
    mk_enum_inst nms =
        [d| instance Avro $(conT nm) where
                toSchema _ = return $ object [ "type" .= string "enum"
                                             , "name" .= string $(tolit nm)
                                             , "symbols" .= $(tolitlist nms) ]
                toBin x = $(
                    return $ CaseE (VarE 'x)
                        [ Match (ConP nm1 [])
                                (NormalB (AppE (VarE 'zigInt)
                                               (LitE (IntegerL i)))) []
                        | (i,nm1) <- zip [0..] nms ] )

                fromBin = zagInt >>= \x -> $(
                    return $ CaseE (VarE 'x)
                        [ Match (LitP (IntegerL i))
                                (NormalB (AppE (VarE 'return)
                                               (ConE nm1))) []
                        | (i,nm1) <- zip [0..] nms ] )

                toAvron x = $(
                    return $ CaseE (VarE 'x)
                        [ Match (ConP nm1 [])
                                (NormalB (AppE (VarE 'string)
                                               (LitE (StringL (nameBase nm1))))) []
                        | nm1 <- nms ] )
        |]

    -- record instance from record-like constructors
    -- XXX maybe allow empty "normal" constructors, too
    mk_record_inst :: [ (Name, [(Name, Strict, Type)]) ] -> Q [Dec]
    mk_record_inst [(nm1,fs1)] =
        [d| instance Avro $(conT nm) where
                toSchema _ = $(mk_product_schema nm1 fs1)
                toBin      = $(to_bin_product fs1)
                fromBin    = $(from_bin_product [| return $(conE nm1) |] fs1)
                toAvron    = $(to_avron_product fs1)
        |]

    mk_record_inst arms =
        [d| instance Avro $(conT nm) where
                toSchema _ = Array . V.fromList <$> sequence
                             $( foldr (\(nm1,fs) k -> [| $(mk_product_schema nm1 fs) : $k |])
                                      [| [] |] arms )
                toBin =
                    $( do x <- newName "x"
                          LamE [VarP x] . CaseE (VarE x)
                             <$> sequence [ ($ []) . Match (RecP nm1 []) . NormalB
                                                <$> [| zigInt $(litE (IntegerL i)) <> $(to_bin_product fs) $(varE x) |]
                                          | (i,(nm1,fs)) <- zip [0..] arms ] )

                fromBin = zagInt >>=
                    $( do x <- newName "x"
                          LamE [VarP x] . CaseE (VarE x)
                            <$> sequence [ ($ []) . Match (LitP (IntegerL i)) . NormalB
                                                <$> from_bin_product [| return $(conE nm1) |] fs
                                         | (i,(nm1,fs)) <- zip [0..] arms ] )

                toAvron =
                    $( do x <- newName "x"
                          LamE [VarP x] . CaseE (VarE x)
                             <$> sequence [ ($ []) . Match (RecP nm1 []) . NormalB
                                                <$> [| object [ $(tolit nm1) .= $(to_avron_product fs) $(varE x) ] |]
                                          | (nm1,fs) <- arms ] )
        |]

    -- create schema for a product from a name and a list of fields
    mk_product_schema nm1 tps =
        [| $( fieldlist tps ) >>= \flds ->
           memoObject $( tolit nm1 )
               [ "type" .= string "record"
               , "fields" .= Array (V.fromList flds) ] |]

    fieldlist = foldr go [| return [] |]
        where
            go (nm1,_,tp) k =
                [| do sch <- toSchema $(sigE (varE 'undefined) (return tp))
                      obs <- $k
                      return $ object [ "name" .= string $(tolit nm1)
                                      , "type" .= sch ]
                             : obs |]

    -- binary encoding of records: field by field.
    to_bin_product nms =
        [| \x -> $( foldr (\(nm1,_,_) k -> [| mappend (toBin ($(varE nm1) x)) $k |] )
                          [| mempty |] nms ) |]

    from_bin_product =
        foldl (\expr (_,_,_) -> [| $expr <*> fromBin |])

    -- json encoding of records: fields in an object
    to_avron_product nms =
        [| \x -> object $(
            foldr (\(nm1,_,_) k -> [| ($(tolit nm1) .= toAvron ($(varE nm1) x)) : $k |] )
                  [| [] |] nms ) |]


data ContainerOpts = ContainerOpts { objects_per_block :: Int
                                   , filetype_label :: B.ByteString }

-- Writing a container file.  This is an 'Enumeratee', we read a list of
-- suitable types, we write a header containing the generated schema,
-- and a series of blocks with serialized data.
writeAvroContainer :: (MonadIO m, Nullable s, ListLike s a, Avro a)
                   => ContainerOpts -> Enumeratee s B.ByteString m r
writeAvroContainer ContainerOpts{..} out = do
        ma <- peekStream
        sync_marker <- liftIO $ B.pack <$> replicateM 16 randomIO

        let schema = encode . runMkSchema . toSchema . fromJust $ ma

            meta :: H.HashMap T.Text B.ByteString
            meta = H.fromList [( "avro.schema", B.concat $ BL.toChunks schema )
                              ,( "avro.codec", "null" )
                              ,( "biohazard.filetype", filetype_label )]

            hdr = fromByteString "Obj\1" <> toBin meta <> fromByteString sync_marker

        let enc_blocks = iterLoop $ \out' -> do (num,code) <- joinI $ takeStream objects_per_block $
                                                                foldStream (\(!n,c) o -> (n+1, c <> toBin o)) (0::Int,mempty)

                                                let code1 = toLazyByteString code
                                                    block = toBin num <> toBin (BL.length code1) <>
                                                            fromLazyByteString code1 <> fromByteString sync_marker
                                                lift (enumList (BL.toChunks $ toLazyByteString block) out')

        lift (enumList (BL.toChunks $ toLazyByteString hdr) out) >>= enc_blocks

-- XXX Possible codecs: null, zlib, snappy, lzma; all missing
-- XXX Should check schema on reading.

readAvroContainer :: (Monad m, ListLike s a, Avro a) => Enumeratee B.ByteString s m r
readAvroContainer out = do
        4 <- heads "Obj\1"  -- enough magic?
        meta <- iterGet (fromBin :: Get (H.HashMap T.Text B.ByteString))
        sync_marker <- iGetString 16

        flip iterLoop out $ \o -> do num <- iterGet zagInt
                                     sz <- iterGet fromBin
                                     o' <- joinI $ takeStream sz $ -- codec goes here
                                              convStream (LL.singleton `liftM` iterGet fromBin) o
                                     16 <- heads sync_marker
                                     return o'

-- | Repeatedly apply an 'Iteratee' to a value until end of stream.
-- Returns the final value.
iterLoop :: (Nullable s, Monad m) => (a -> Iteratee s m a) -> a -> Iteratee s m a
iterLoop it a = do e <- isFinished
                   if e then return a
                        else it a >>= iterLoop it


iterGet :: Monad m => Get a -> Iteratee B.ByteString m a
iterGet = go . runGetIncremental
  where
    go (Fail  _ _ err) = throwErr (iterStrExc err)
    go (Done rest _ a) = idone a (Chunk rest)
    go (Partial   dec) = liftI $ \ck -> case ck of
        Chunk s -> go (dec $ Just s)
        EOF  mx -> case dec Nothing of
            Fail  _ _ err -> throwErr (iterStrExc err)
            Partial     _ -> throwErr (iterStrExc "<partial>")
            Done rest _ a | B.null rest -> idone a (EOF mx)
                          | otherwise   -> idone a (Chunk rest)