{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}

import Control.Monad
import Control.Monad.IO.Class
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import Data.ByteString as BS
import Data.Word
import Streamly hiding (parallel)
import Streamly.Binary
import qualified Streamly.Data.Fold as FL
import qualified Streamly.Prelude as S
import Test.Hspec
import Test.QuickCheck
import qualified Test.QuickCheck.Gen as Gen
import Test.QuickCheck.Monadic

-- One element of stream is one encoded Object
prop_normal :: [Object] -> Property
prop_normal objs = monadicIO do
  rt <-
    run $
      S.toList $
        decodeStream $ encodeStream $ S.fromList objs
  return $ rt === objs

-- One element of stream is one byte (represented as bytestring)
prop_single_bytes :: [Object] -> Property
prop_single_bytes objs = monadicIO do
  rt <-
    run $
      S.toList $
        decodeStream
          (encodeStream (S.fromList objs) >>= S.fromList . fmap BS.singleton . BS.unpack)
  return $ rt === objs

-- One element of stream is all encoded objects concatenated as one big bytestring
prop_one_bytestring :: [Object] -> Property
prop_one_bytestring objs = monadicIO do
  rt <-
    run $
      S.toList $
        decodeStream $
          S.yieldM (S.fold FL.mconcat $ encodeStream $ S.fromList objs)
  return $ rt === objs

main :: IO ()
main =
  hspec $
    parallel do
      describe "Binary instance" $
        it "decode . encode === id" $ property \obj ->
          obj === decode (encode (obj :: Object))
      describe "decodeStream . encodeStream === id" do
        it "stream of bytestrings" $ property prop_normal
        it "stream of bytes" $ property prop_single_bytes
        it "stream of one large bytestring" $ property prop_one_bytestring

data Object
  = ObjectNil
  | ObjectNum Word32
  | ObjectStr String
  | ObjectBin ByteString
  | ObjectMap [(Object, Object)]
  deriving (Show, Eq)

instance Arbitrary Object where
  arbitrary = Gen.sized \n ->
    Gen.oneof
      [ pure ObjectNil,
        ObjectNum <$> arbitrary,
        ObjectStr <$> arbitrary,
        ObjectBin <$> (pack <$> arbitrary),
        ObjectMap <$> Gen.resize (n `div` 4) arbitrary
      ]

instance Binary Object where
  put ObjectNil = putWord8 0
  put (ObjectNum w) = putWord8 1 >> putWord32be w
  put (ObjectStr s) = putWord8 2 >> put s
  put (ObjectBin b) = putWord8 3 >> put b
  put (ObjectMap m) = putWord8 4 >> putList m
  get =
    getWord8 >>= \case
      0 -> pure ObjectNil
      1 -> ObjectNum <$> getWord32be
      2 -> ObjectStr <$> get
      3 -> ObjectBin <$> get
      4 -> ObjectMap <$> get
      _ -> fail "unknown tag"