{-# LANGUAGE OverloadedStrings #-}

module System.IO.Streams.Tests.ByteString (tests) where

------------------------------------------------------------------------------
import           Control.Concurrent
import           Control.Monad
import           Data.ByteString.Char8                (ByteString)
import qualified Data.ByteString.Char8                as S
import qualified Data.ByteString.Lazy.Char8           as L
import           Data.List                            hiding (lines,
                                                       takeWhile, unlines,
                                                       unwords, words)
import           Data.Maybe                           (isJust)
import           Data.Monoid
import           Prelude                              hiding (lines, read,
                                                       takeWhile, unlines,
                                                       unwords, unwords,
                                                       words)
import qualified Prelude
import           System.IO.Streams                    hiding (filter,
                                                       intersperse, mapM_)
import           System.IO.Streams.Tests.Common
import           System.Timeout
import           Test.Framework
import           Test.Framework.Providers.HUnit
import           Test.Framework.Providers.QuickCheck2
import           Test.HUnit                           hiding (Test)
import           Test.QuickCheck                      hiding (output)
import           Test.QuickCheck.Monadic
------------------------------------------------------------------------------

tests :: [Test]
tests = [ testBoyerMoore
        , testBoyerMoore2
        , testCountInput
        , testCountInput2
        , testCountOutput
        , testThrowIfTooSlow
        , testReadExactly
        , testTakeWhile
        , testTakeBytes
        , testTakeBytes2
        , testTakeBytes3
        , testTakeBytes4
        , testTakeExactly
        , testTakeExactly2
        , testTakeExactly3
        , testThrowIfProducesMoreThan
        , testThrowIfProducesMoreThan2
        , testThrowIfProducesMoreThan3
        , testThrowIfConsumesMoreThan
        , testThrowIfConsumesMoreThan2
        , testTrivials
        , testWriteLazyByteString
        , testGiveBytes
        , testGiveExactly
        , testLines
        , testWords
        ]


------------------------------------------------------------------------------
testCountInput :: Test
testCountInput = testProperty "bytestring/countInput" $
                 monadicIO $
                 forAllM arbitrary prop
  where
    prop :: [ByteString] -> PropertyM IO ()
    prop l = liftQ $ do
        is          <- fromList l
        (is', grab) <- countInput is

        x <- toList is'
        n <- grab

        assertEqual "countInput1" (L.length $ L.fromChunks l) n
        assertEqual "countInput2" (L.length $ L.fromChunks x) n

        read is' >>= assertEqual "eof" Nothing
        unRead "ok" is'
        peek is >>= assertEqual "peek" (Just "ok")
        read is' >>= assertEqual "read" (Just "ok")


------------------------------------------------------------------------------
testCountInput2 :: Test
testCountInput2 = testCase "bytestring/countInput2" $ do
    is              <- fromList txt
    (is', getCount) <- countInput is
    (Just x)        <- read is'

    unRead "0, " is'
    getCount >>= assertEqual "count1" 5
    peek is  >>= assertEqual "pushback propagates" (Just "0, ")

    (liftM (L.fromChunks . (x:)) $ toList is') >>=
        assertEqual "output" expectedOutput

    getCount >>= assertEqual "count2" (L.length $ L.fromChunks txt)

  where
    txt            = ["testing ", "1, ", "2, ", "3"]
    expectedOutput = "testing 0, 1, 2, 3"


------------------------------------------------------------------------------
testCountOutput :: Test
testCountOutput = testProperty "bytestring/countOutput" $
                  monadicIO $
                  forAllM arbitrary prop
  where
    prop :: [ByteString] -> PropertyM IO ()
    prop l = liftQ $ do
        is            <- fromList l
        (os0, grab)   <- listOutputStream
        (os, grabLen) <- countOutput os0

        connect is os

        xs <- grab
        n  <- grabLen

        assertEqual "countOutput1" l xs
        assertEqual "countOutput2" (L.length $ L.fromChunks l) n


------------------------------------------------------------------------------
testTakeBytes :: Test
testTakeBytes = testProperty "bytestring/takeBytes" $
                monadicIO $
                forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = pre (L.length l > 5) >> liftQ (do
        let (a,b) = L.splitAt 4 l

        is  <- fromList (L.toChunks l)
        is' <- takeBytes 4 is

        x   <- liftM L.fromChunks $ toList is'
        y   <- liftM L.fromChunks $ toList is

        assertEqual "take1" a x
        assertEqual "take2" b y
        )

------------------------------------------------------------------------------
testTakeBytes2 :: Test
testTakeBytes2 = testProperty "bytestring/takeBytes2" $
                 monadicIO $
                 forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = liftQ $ do
        is  <- fromList (L.toChunks l)
        is2 <- takeBytes 0 is

        x   <- toList is2
        y   <- liftM L.fromChunks $ toList is

        assertEqual "takeBytes3" [] x
        assertEqual "takeBytes4" l y

        -- Test that pushback makes it back to the source inputstream
        is3 <- takeBytes 20 is
        void $ toList is3
        unRead "ok2" is3
        unRead "ok1" is3

        z   <- toList is
        assertEqual "takeBytes5" ["ok1", "ok2"] z


------------------------------------------------------------------------------
testTakeBytes3 :: Test
testTakeBytes3 = testCase "bytestring/takeBytes3" $ do
    is <- fromLazyByteString (L.fromChunks ["The", "quick", "brown", "fox"]) >>=
          takeBytes 100
    _  <- toList is
    m  <- read is

    assertEqual "takeBytes3" Nothing m


------------------------------------------------------------------------------
testTakeBytes4 :: Test
testTakeBytes4 = testCase "bytestring/takeBytes4" $ do
    is <- makeInputStream (threadDelay 20000000 >> return Nothing)
          >>= takeBytes 0
    mb  <- timeout 100000 $ toList is
    assertBool "takeBytes4" $ isJust mb


------------------------------------------------------------------------------
testTakeExactly :: Test
testTakeExactly = testProperty "bytestring/takeExactly" $
                monadicIO $
                forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = pre (L.length l > 5) >> liftQ (do
        let (a,b) = L.splitAt 4 l

        is  <- fromList (L.toChunks l)
        is' <- takeExactly 4 is

        x   <- liftM L.fromChunks $ toList is'
        y   <- liftM L.fromChunks $ toList is

        assertEqual "take1" a x
        assertEqual "take2" b y
        )

------------------------------------------------------------------------------
testTakeExactly2 :: Test
testTakeExactly2 = testProperty "bytestring/takeExactly2" $
                monadicIO $
                forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = pre (L.length l < 10) >> liftQ (do
        is  <- fromList (L.toChunks l)
        is' <- takeExactly 10 is

        expectExceptionH $ toList is'
        )

------------------------------------------------------------------------------
testTakeExactly3 :: Test
testTakeExactly3 = testCase "bytestring/takeExactly3" $ do
    is  <- fromList ["one", "two"]
    is' <- takeExactly 7 is
    expectExceptionH $ toList is'

------------------------------------------------------------------------------
testThrowIfProducesMoreThan :: Test
testThrowIfProducesMoreThan =
    testProperty "bytestring/throwIfProducesMoreThan" $
    monadicIO $ forAllM arbitrary prop

  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = do
        pre (L.length l > 5)

        liftQ $ do
            is  <- fromList $ L.toChunks l
            is' <- throwIfProducesMoreThan 4 is
            expectExceptionH $ toList is'


------------------------------------------------------------------------------
testThrowIfProducesMoreThan2 :: Test
testThrowIfProducesMoreThan2 =
    testProperty "bytestring/throwIfProducesMoreThan2" $
    monadicIO $ forAllM arbitrary prop

  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = do
        let n = L.length l

        liftQ $ do
            is  <- fromList $ L.toChunks l
            is' <- throwIfProducesMoreThan (n + 1) is
            l'  <- liftM L.fromChunks $ toList is'
            assertEqual "throwIfProducesMoreThan2" l l'

            m   <- read is'
            assertEqual "throwIfProducesMoreThan2-2" Nothing m

            unRead "ok2" is'
            unRead "ok1" is'
            z   <- toList is
            assertEqual "throwIfProducesMoreThan2-3" ["ok1", "ok2"] z


------------------------------------------------------------------------------
testThrowIfProducesMoreThan3 :: Test
testThrowIfProducesMoreThan3 =
    testCase "bytestring/throwIfProducesMoreThan3" $ do
        is <- fromList ["lo", "ngstring"] >>= throwIfProducesMoreThan 4
        s  <- readExactly 4 is
        assertEqual "throwIfProducesMoreThan split" "long" s

        l <- fromList ["ok", "", "", "", ""] >>= throwIfProducesMoreThan 2 >>=
             toList
        assertEqual "throwIfProducesMoreThan3" ["ok", "", "", "", ""] l


------------------------------------------------------------------------------
testThrowIfConsumesMoreThan :: Test
testThrowIfConsumesMoreThan =
    testProperty "bytestring/throwIfConsumesMoreThan" $
    monadicIO $
    forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = do
        let n = L.length l
        pre (n > 0)
        liftQ $ do
             is      <- fromList (L.toChunks l)
             (os, _) <- listOutputStream
             os'     <- throwIfConsumesMoreThan (n-1) os

             expectExceptionH $ connect is os'


------------------------------------------------------------------------------
testThrowIfConsumesMoreThan2 :: Test
testThrowIfConsumesMoreThan2 =
    testProperty "bytestring/throwIfConsumesMoreThan2" $
    monadicIO $
    forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = do
        let n = L.length l

        liftQ $ do
             is         <- fromList (L.toChunks l)
             (os, grab) <- listOutputStream
             os'        <- throwIfConsumesMoreThan n os

             connect is os'

             l' <- liftM L.fromChunks grab
             assertEqual "throwIfConsumesMoreThan" l l'


------------------------------------------------------------------------------
testGiveExactly :: Test
testGiveExactly = testCase "bytestring/giveExactly" $ do
    f 2 >>= assertEqual "ok" ["ok"]
    expectExceptionH $ f 1
    expectExceptionH $ f 3

  where
    f n = do
      is <- fromByteString "ok"
      outputToList (giveExactly n >=> connect is)


------------------------------------------------------------------------------
testGiveBytes :: Test
testGiveBytes = testProperty "bytestring/giveBytes" $
                monadicIO $
                forAllM arbitrary prop
  where
    prop :: L.ByteString -> PropertyM IO ()
    prop l = do
        pre (L.length l > 5)
        let a = L.take 4 l

        liftQ $ do
            is         <- fromList (L.toChunks l)
            (os, grab) <- listOutputStream
            os'        <- giveBytes 4 os

            connect is os'
            write Nothing os'

            x <- liftM L.fromChunks grab

            assertEqual "giveBytes1" a x

        liftQ $ do
            is <- fromList $ L.toChunks a
            (os, grab) <- listOutputStream
            os'        <- giveBytes 10 os

            connect is os'
            write Nothing os'
            x <- liftM L.fromChunks grab
            assertEqual "giveBytes2" a x


------------------------------------------------------------------------------
testThrowIfTooSlow :: Test
testThrowIfTooSlow = testCase "bytestring/throwIfTooSlow" $ do
    is <- mkList
    expectExceptionH $ trickleFrom is

    is' <- mkList
    void $ toList is'
    x <- read is'

    assertEqual "throwIfTooSlow" Nothing x

    src  <- mkSrc
    src' <- throwIfTooSlow (return ()) 10 2 src

    void $ toList src'
    unRead "ok2" src'
    unRead "ok1" src'
    l <- toList src

    assertEqual "throwIfTooSlow/pushback" ["ok1", "ok2"] l


  where
    mkSrc = fromList $ Prelude.take 100 $ cycle $
            intersperse " " ["the", "quick", "brown", "fox"]

    mkList = mkSrc >>= throwIfTooSlow (return ()) 10 2

    trickleFrom is = go
      where
        go = read is >>= maybe (return ())
                               (\x -> x `seq` (threadDelay 2000000 >> go))


------------------------------------------------------------------------------
testBoyerMoore :: Test
testBoyerMoore = testProperty "bytestring/boyerMoore" $
                 monadicIO $ forAllM gen prop
  where
    genBS range = liftM S.pack $ listOf $ choose range

    gen :: Gen (ByteString, [ByteString])
    gen = do
        needle <- genBS ('a', 'z')
        n      <- choose (0, 10)
        hay    <- replicateM n $ genBS ('A', 'Z')
        return (needle, hay)

    prop :: (ByteString, [ByteString]) -> PropertyM IO ()
    prop (needle, haystack') = do
        let lneedle   = L.fromChunks [needle]
        let lhaystack = L.fromChunks haystack'

        pre ((not $ S.null needle) &&
             (not $ L.null lhaystack) &&
             (not $ S.isInfixOf needle $ S.concat haystack'))


        (lhay, toklist0) <- insertNeedle lneedle lhaystack
        let stream  = L.toChunks $ L.concat [lneedle, lhay]
        let toklist = (Match needle) : toklist0

        -- there should be exactly three matches
        out <- liftQ (fromList stream >>= search needle >>= toList)

        let nMatches = length $ filter isMatch out

        let out' = concatAdj Nothing id out

        when (nMatches /= 3 || out' /= toklist) $ liftQ $ do
            putStrLn "got wrong output!!"
            putStrLn "needle:\n"
            putStrLn $ show lneedle
            putStrLn "\nhaystack:\n"
            mapM_ (putStrLn . show) stream
            putStrLn "\noutput stream:"
            mapM_ (putStrLn . show) out
            putStrLn "\noutput stream (minified):"
            mapM_ (putStrLn . show) out'
            putStrLn "\nexpected output:"
            mapM_ (putStrLn . show) toklist
            putStrLn ""

        liftQ $ do
            assertEqual "boyer-moore matches" 3 nMatches
            assertEqual "boyer-moore output" toklist out'


    isMatch (Match _) = True
    isMatch _         = False

    concatAdj :: Maybe MatchInfo
              -> ([MatchInfo] -> [MatchInfo])
              -> [MatchInfo]
              -> [MatchInfo]
    concatAdj prefix dl []     = dl $ maybe [] (:[]) prefix
    concatAdj prefix dl (x:xs) =
        maybe (concatAdj (Just x) dl xs)
              (\p -> maybe (concatAdj (Just x) (dl . (p:)) xs)
                           (\x' -> concatAdj (Just x') dl xs)
                           (merge p x))
              prefix

      where
        merge (NoMatch z) y
            | S.null z  = Just y
            | otherwise = case y of
                            NoMatch x' -> Just $ NoMatch $ z `mappend` x'
                            _          -> Nothing

        merge (Match _) _ = Nothing

    insertNeedle lneedle lhaystack = do
        idxL  <- pick $ choose (0, lenL-1)
        idxN  <- pick $ choose (0, lenN-1)
        idxN2 <- pick $ choose (0, lenN-1)
        let (l1, l2) = L.splitAt (toEnum idxL) lhaystack
        let (n1, n2) = L.splitAt (toEnum idxN) lneedle
        let (n3, n4) = L.splitAt (toEnum idxN2) lneedle

        let out1 = L.concat [ l1, n1, n2, l2, n3, n4 ]

        let res = concatAdj Nothing id
                      [ NoMatch $ strict l1
                      , Match   $ strict lneedle
                      , NoMatch $ strict l2
                      , Match   $ strict lneedle
                      ]

        return (out1, res)

      where
        strict = S.concat . L.toChunks

        lenN = fromEnum $ L.length lneedle
        lenL = fromEnum $ L.length lhaystack


------------------------------------------------------------------------------
testBoyerMoore2 :: Test
testBoyerMoore2 = testCase "bytestring/boyerMoore2" $ do
    fromList ["bork", "no", "bork", "bor"]
        >>= search "bork"
        >>= toList
        >>= assertEqual "bork!" [ Match "bork", NoMatch "no", Match "bork"
                                , NoMatch "bor" ]

    fromList [] >>= search "bork" >>= toList >>= assertEqual "nothing" []

    fromList ["borkbo", "r"] >>= search "bork" >>= toList
        >>= assertEqual "borkbo" [Match "bork", NoMatch "bor"]

    fromList ["borkborkborkb", "o", "r", "k", "b", "o"]
        >>= search "borkborkbork"
        >>= toList
        >>= assertEqual "boooooork" [Match "borkborkbork", NoMatch "borkbo"]

    fromList ["bbbbb", "o", "r", "k", "bork"]
        >>= search "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
        >>= toList
        >>= assertEqual "bbbbbbbbb" [NoMatch "bbbbborkbork"]

    fromList ["bbbbbbbbb", "o", "r", "k"]
        >>= search "bbbbbbbb"
        >>= toList
        >>= assertEqual "bbb2" [Match "bbbbbbbb", NoMatch "bork"]

    fromList ["bababa", "bo", "rk", "bz", "z", "z", "z"] >>= search "babababork"
        >>= toList
        >>= assertEqual "zzz" [Match "babababork", NoMatch "bzzzz"]

    fromList ["bab", "a", "b"] >>= search "bababa"
        >>= toList
        >>= assertEqual "bab" [NoMatch "babab"]

    fromList ["xxx", "xxxzx", "xx"] >>= search "xxxx" >>= toList
        >>= assertEqual "xxxx" [Match "xxxx", NoMatch "xxzxxx"]


------------------------------------------------------------------------------
testWriteLazyByteString :: Test
testWriteLazyByteString = testProperty "bytestring/writeLazy" $
                          monadicIO $
                          forAllM arbitrary prop
  where
    prop :: [ByteString] -> PropertyM IO ()
    prop l0 = liftQ $ do
        let l = filter (not . S.null) l0
        let s = L.fromChunks l
        (os, grab) <- listOutputStream
        writeLazyByteString s os

        l' <- grab
        assertEqual "writeLazy" l l'


------------------------------------------------------------------------------
testReadExactly :: Test
testReadExactly = testProperty "bytestring/readExactly" $
                  monadicIO $
                  forAllM arbitrary prop
  where
    prop l0 = liftQ $ do
        let l = filter (not . S.null) l0
        is <- fromList l

        let s = L.fromChunks l
        let n = fromEnum $ L.length s

        t <- readExactly n is
        assertEqual "eq" s $ L.fromChunks [t]

        unRead t is
        expectExceptionH $ readExactly (n+1) is

        when (n > 0) $ do
            is' <- fromList l
            u   <- readExactly (n-1) is'
            assertEqual "eq2" (L.take (toEnum $ n-1) s) (L.fromChunks [u])
            v   <- readExactly 1 is'
            assertEqual "eq3" (L.drop (toEnum $ n-1) s) (L.fromChunks [v])


------------------------------------------------------------------------------
testTakeWhile :: Test
testTakeWhile = testCase "bytestring/takeBytesWhile" $ do
    is <- fromList ["test", "ing\n", "1-2-3\n1-2-3"]

    takeBytesWhile (/= '\n') is >>=
        assertEqual "takeBytesWhile1" (Just "testing")
    takeBytesWhile (/= '\n') is >>=
        assertEqual "takeBytesWhile2" (Just "")
    readExactly 1 is >>= assertEqual "readExactly" "\n"
    takeBytesWhile (/= '\n') is >>=
        assertEqual "takeBytesWhile3" (Just "1-2-3")
    readExactly 1 is >>= assertEqual "readExactly" "\n"
    takeBytesWhile (/= '\n') is >>=
        assertEqual "takeBytesWhile4" (Just "1-2-3")
    takeBytesWhile (/= '\n') is >>= assertEqual "takeBytesWhile4" Nothing


------------------------------------------------------------------------------
testLines :: Test
testLines = testCase "bytestring/testLines" $ do
    fromList ["th", "e\nquick\nbrown", "\n", "", "fox"] >>= lines >>=
             toList >>= assertEqual "lines" ["the", "quick", "brown", "fox"]
    fromList [] >>= lines >>= toList >>= assertEqual "empty lines" []

    fromList ["ok", "cool"] >>=
      \is -> outputToList (\os -> unlines os >>= connect is) >>=
      assertEqual "unlines" ["ok", "\n", "cool", "\n"]


------------------------------------------------------------------------------
testWords :: Test
testWords = testCase "bytestring/testWords" $ do
    fromList ["the quick brown    \n\tfox"] >>= words >>=
             toList >>= assertEqual "words" ["the", "quick", "brown", "fox"]

    fromList ["ok", "cool"] >>=
      \is -> outputToList (\os -> unwords os >>= connect is) >>=
      assertEqual "unlines" ["ok", " ", "cool"]


------------------------------------------------------------------------------
testTrivials :: Test
testTrivials = testCase "bytestring/testTrivials" $ do
    coverTypeableInstance (undefined :: TooManyBytesReadException)
    coverShowInstance     (undefined :: TooManyBytesReadException)
    coverTypeableInstance (undefined :: TooFewBytesWrittenException)
    coverShowInstance     (undefined :: TooFewBytesWrittenException)
    coverTypeableInstance (undefined :: TooManyBytesWrittenException)
    coverShowInstance     (undefined :: TooManyBytesWrittenException)
    coverTypeableInstance (undefined :: RateTooSlowException)
    coverShowInstance     (undefined :: RateTooSlowException)
    coverTypeableInstance (undefined :: ReadTooShortException)

    coverEqInstance $ Match ""

    coverShowInstance $ Match ""
    coverShowInstance $ NoMatch ""