{-# LANGUAGE ScopedTypeVariables #-}
module Main where

import           Control.Monad
import           Control.Monad.Loops
import           Data.Bifunctor
import           Data.Bits
import qualified Data.ByteString       as BS
import qualified Data.ByteString.Lazy  as BL
import           Data.Machine
import           Data.Machine.Seekable
import           Data.Word
import           Test.Tasty
import           Test.Tasty.HUnit
import qualified Test.Tasty.QuickCheck as QC
import           Test.Tasty.QuickCheck ((==>), (===))

seekableWord64 :: forall m i a
               . (Monad m, Num i) => MachineT m (Seekable i Word64) a -> MachineT m (Seekable i BS.ByteString) a
seekableWord64 = seeking mkAtPnt readFrom
  where
    readFrom ::      MachineT m (Seekable i BS.ByteString) Word64
    mkAtPnt  :: i -> MachineT m (Seekable i BS.ByteString) Word64
{- When we read past the end, this can recover
    readFrom = (construct . forever $ awaits Read >>= yield) `stack` stackParse
    stackParse = construct . forever $ do
      v <- iterateUntilM ((<=) 8 . BS.length) (\a -> (a<>) <$> pop) mempty
      let (h, t) = BS.splitAt 8 v
      push t
      yield $ bsToW64 h --}
--{- This one stops after reading past the end
    readFrom = construct $
      stackParse mempty
    stackParse p = do
      v <- iterateUntilM ((<=) 8 . BS.length) (\a -> (a<>) <$> awaits Read) p
      let (h, t) = BS.splitAt 8 v
      yield $ bsToW64 h
      stackParse t --}
    mkAtPnt i = encased $ Await (\() -> readFrom) (SeekPos (i*8)) stopped
    bsToW64 :: BS.ByteString -> Word64
    bsToW64 bs = foldl (\a w -> (fromIntegral w) .|. shiftL a 8) 0 . take 8 . BS.unpack $ bs
  
main :: IO ()
main = defaultMain tests

takeAll :: PlanT (Seekable i o) o m a
takeAll = forever $ awaits Read >>= yield

tests :: TestTree
tests = testGroup "machine-seekable"
  [ testCase "Empty pass through" $ do
    let bs = BS.pack []
    (BL.fromStrict bs) @=? (BL.fromChunks (run (seekByteString bs (construct takeAll))))
  , testCase "Simple pass through" $ do
    let bs = BS.pack [0]
    (BL.fromStrict bs) @=? (BL.fromChunks (run (seekByteString bs (construct takeAll))))
  , testCase "Longer pass through" $ do
    let bs = BS.pack [0,1,2,3,4,5,6,7,8,9,10]
    (BL.fromStrict bs) @=? (BL.fromChunks (run (seekByteString bs (construct takeAll))))
  , QC.testProperty "Pass through property" $
      \bs' ->
        let bs = BS.pack bs'
        in (BL.fromStrict bs) ===
           (BL.fromChunks (run (seekByteString bs (construct takeAll))))
  , QC.testProperty "Seek & Pass through property" $
      \d bs' ->
        let bs = BS.pack bs'
        in (d >= 0) ==>
           (BL.drop (fromIntegral d) $ BL.fromStrict bs) ===
           (BL.fromChunks (run (seekByteString bs (construct $ awaits (SeekPos d) >> takeAll))))
  , testCase "Composite Machine reads" $ do
    let bs = BS.pack [0,0,0,0,0,0,0,0x1
                     ,0,0,0,0,0,0,0,0x2
                     ,0,0,0,0,0,0,0,0x3
                     ,0,0,0,0,0,0,0,0x4
                     ,0,0,0,0,0,0,0,0x5
                     ,0,0,0,0,0,0,0,0x6
                     ,0,0,0,0,0,0,0,0x7
                     ,0,0,0,0,0,0,0,0x8
                     ,0,0,0,0,0,0,0,0x9
                     ,0,0,0,0,0,0,0xA,0]
    ([0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA00]) @=?
      (run (seekByteString bs $ seekableWord64 (construct takeAll)))
  , testCase "Composite Machine seeks & reads" $ do
    let bs = BS.pack [0,0,0,0,0,0,0,0x1
                     ,0,0,0,0,0,0,0,0x2
                     ,0,0,0,0,0,0,0,0x3
                     ,0,0,0,0,0,0,0,0x4
                     ,0,0,0,0,0,0,0,0x5
                     ,0,0,0,0,0,0,0,0x6
                     ,0,0,0,0,0,0,0,0x7
                     ,0,0,0,0,0,0,0,0x8
                     ,0,0,0,0,0,0,0,0x9
                     ,0,0,0,0,0,0,0xA,0]
    ([0x1, 0x4, 0x4, 0x7, 0xA00, 0x1, 0x2, 0x4]) @=?
      (run . seekByteString bs . seekableWord64 . construct $ do
          awaits Read >>= yield
          awaits (SeekPos 3) >> awaits Read >>= yield
          awaits (SeekPos 3) >> awaits Read >>= yield
          awaits (SeekPos 6) >> awaits Read >>= yield
          awaits (SeekPos 9) >> awaits Read >>= yield
          awaits (SeekPos 0) >> awaits Read >>= yield
          awaits Read >>= yield
          awaits (SeekPos 3) >> awaits Read >>= yield
      )
  , QC.testProperty "Composite Seek & Pass through property" $
      \(d'::Word8) ->
        let d = fromIntegral d'
            r = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA00]
            bs = BS.pack [0,0,0,0,0,0,0,0x1
                         ,0,0,0,0,0,0,0,0x2
                         ,0,0,0,0,0,0,0,0x3
                         ,0,0,0,0,0,0,0,0x4
                         ,0,0,0,0,0,0,0,0x5
                         ,0,0,0,0,0,0,0,0x6
                         ,0,0,0,0,0,0,0,0x7
                         ,0,0,0,0,0,0,0,0x8
                         ,0,0,0,0,0,0,0,0x9
                         ,0,0,0,0,0,0,0xA,0]
        in (d >= 0) ==>
           (drop (fromIntegral d) r) ===
           (run (seekByteString bs $ seekableWord64 (construct $ awaits (SeekPos d) >> takeAll)))
  , QC.testProperty "Seek reconstruction property" $
      \(sks'::[(Word8, Word8)]) ->
      let sks::[(Word64, Word64)] = map (bimap fromIntegral fromIntegral) sks'
          r = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA00]
          bs = BS.pack [0,0,0,0,0,0,0,0x1
                       ,0,0,0,0,0,0,0,0x2
                       ,0,0,0,0,0,0,0,0x3
                       ,0,0,0,0,0,0,0,0x4
                       ,0,0,0,0,0,0,0,0x5
                       ,0,0,0,0,0,0,0,0x6
                       ,0,0,0,0,0,0,0,0x7
                       ,0,0,0,0,0,0,0,0x8
                       ,0,0,0,0,0,0,0,0x9
                       ,0,0,0,0,0,0,0xA,0]
          expect = concatMap (\(p,l) -> take (fromIntegral l) $ drop (fromIntegral p) r) .
                     take (length (takeWhile (\(p, l) ->
                               (l<=0) || (p+l <= (fromIntegral $ length r))) sks)+1) $
                      sks
          result = run . seekByteString bs . seekableWord64 . construct . forM_ sks $ \(p, l) -> do
              awaits (SeekPos p)
              replicateM_ (fromIntegral l) $ awaits Read >>= yield
      in expect === result
  ]