{-# LANGUAGE GeneralizedNewtypeDeriving, NoMonomorphismRestriction,
             DataKinds, TypeFamilies, TemplateHaskell, ScopedTypeVariables,
             MagicHash, FlexibleContexts #-}
import Test.Tasty
import Test.Tasty.HUnit
import Control.Monad.Trans.Class
import qualified Data.Functor.Identity as I
import qualified Control.Monad.Trans.Reader as R
import qualified Control.Monad.Trans.Writer as W
import Control.Monad.Classes
import Control.Monad.Classes.Run
import Control.Applicative
import Control.Exception hiding (throw)
import Data.Lens.Light
import Data.Proxy
import GHC.Prim (Proxy#, proxy#)

-- for IO tests
import qualified Foreign.Storable as Foreign
import qualified Foreign.Marshal.Alloc as Foreign

-- for monad-control tests
import qualified Data.Conduit as C
import Control.Monad.Morph

-- for zoom tests
data Record = Record
  { _listL :: [Int]
  , _intL :: Int
  }
  deriving (Show, Eq)

makeLens ''Record

main = defaultMain tests

tests :: TestTree
tests = testGroup "Tests"
  [ readerTests
  , simpleStateTests
  , twoStatesTests
  , liftingTest
  , localState
  , exceptTests
  , execTests
  , zoomTests
  , liftNTests
  , liftConduitTest
  , mapWriterTest
  , readStateTest
  , polymorphicTests
  ]

readerTests = testGroup "Reader Tests"
  [ testCase "ask" $
      let base = 5 :: Integer
          power = 3 :: Int
          expected = 125 :: Integer
      in (runReader power action) base @?= expected
  , testCase "local ask" $
      let base = 5 :: Integer
          power = 2 :: Int
          altBase = 7 :: Integer
          altPower = 3 :: Int
          expected = 174 :: Integer
          action' = do
            x <- local (const altBase) action
            y <- local (const altPower) action
            pure (x + y)
      in (runReader power action') base @?= expected
  ]
  where
    f = (^) :: Integer -> Int -> Integer
    action = f <$> ask <*> ask

simpleStateTests = testGroup "Simple State"
  [ testCase "get" $
      (run $ runStateLazy (0 :: Int) get) @?= (0 :: Int, 0 :: Int)
  , testCase "put" $
      (run $ runStateLazy (0 :: Int) (put (1 :: Int))) @?= ((), 1 :: Int)
  , testCase "put-get-put" $
      (run $ runStateLazy (0 :: Int) (put (1 :: Int) *> get <* put (2 :: Int))) @?= (1 :: Int, 2 :: Int)
  ]

twoStatesComp = put 'b' >> put True >> put 'c'

twoStatesTests = testCase "Two States" $
  (run $ runStateLazy 'a' $ runStateLazy False twoStatesComp) @?= (((), True), 'c')

newtype Foo m a = Foo { runFoo :: m a }
  deriving (Functor, Applicative, Monad)
instance MonadTrans Foo where
  lift = Foo
type instance CanDo (Foo m) eff = False

liftingTest = testCase "Lifting through an unknown transformer" $
  (run $ runStateLazy 'a' $ runFoo $ runStateLazy False twoStatesComp) @?= (((), True), 'c')

localState = testCase "MonadLocal StateT" $
  (run $ evalStateStrict 'a' $
    do
      s1 <- get
      (s2,s3) <- local (toEnum . (+1) . fromEnum :: Char -> Char) $ do
        s2 <- get
        put 'x'
        s3 <- get
        return (s2,s3)
      s4 <- get
      return [s1,s2,s3,s4]) @?= "abxa"

exceptTests = testGroup "Except"
  [ testCase "Catch before IO" $ do
      r <- runExcept $ runStateStrict False $ throw $ ErrorCall "foo"
      (r :: Either ErrorCall ((), Bool)) @?= Left (ErrorCall "foo")
  , testCase "Let escape to IO" $ do
      r <- try $ runExcept $ runStateStrict False $ throw UserInterrupt
      (r :: Either AsyncException (Either ErrorCall ((), Bool))) @?= Left UserInterrupt
  ]

execTests = testCase "Exec" $ do
  r <- runWriterStrict $ exec $
    Foreign.alloca $ \ptr -> do
      Foreign.poke ptr True
      Foreign.peek ptr
  r @?= (True, ())

zoomTests = testCase "Zoom" $ do
  ((4, [2,5], 6), Record [2,5,10] 6) @?=
    (run $ runStateStrict (Record [2] 4) $ runZoom (vanLaarhoven intL) $ runZoom (vanLaarhoven listL) $ do
      (s0 :: Int) <- get
      tell [5 :: Int]
      (s1 :: [Int]) <- ask
      put (6 :: Int)
      (s2 :: Int) <- ask
      tell [10 :: Int]
      return (s0, s1, s2)
    )

liftNTests = testCase "liftN" $ do
  (run $ runReader 'a' $ runReader 'b' $ runReader 'c' $
    liftN (proxy# :: Proxy# (Succ Zero)) R.ask)
  @?= 'b'


liftConduit
  :: forall m n effM eff i o r .
     ( n ~ Find eff m
     , MonadLiftN n m
     , effM ~ Down n m
     , Monad effM
     )
  => Proxy# eff
  -> C.ConduitM i o effM r
  -> C.ConduitM i o m    r
liftConduit _ = C.transPipe (liftN (proxy# :: Proxy# n))

liftConduitTest = testCase "lift conduit" $
  (let
    src :: C.Source I.Identity Int
    src = C.yield 1 >> C.yield 2

    sink :: C.Sink Int (W.Writer [Int]) ()
    sink =
      C.await >>=
        maybe (return ()) (\x -> do lift $ tell [x::Int]; sink)
   in
    W.execWriter $ C.transPipe (liftN (proxy# :: Proxy# (Succ Zero))) src C.$$ sink
  ) @?= [1,2]
  {-

  execWriterStrict $ runReader (3 :: Int) $
    liftConduit (proxy# :: Proxy# (EffReader Int))
      (do
        x <- ask
        liftConduit (C.yield x)
        liftConduit (C.yield (x :: Int)))
    C.$$
      (proxy# :: Proxy# (EffWriter String)) (do C.awaitForever $ \y -> tell (show (y :: Int) ++ "\n")))
  @?= ""-}

mapWriterTest = testCase "mapWriter" $ do
  run (execWriterStrict $ mapWriter (\(w :: Char) -> [w]) $ do { tell 'a'; tell 'b'; tell 'c' }) @?= "abc"

readStateTest = testCase "ReadState" $ do
  let
    a1 :: MonadReader Char m => m Char
    a1 = ask

    a2 :: MonadState Char m => m Char
    a2 = runReadState (Proxy :: Proxy Char) a1

  run (evalStateStrict 'w' a2) @?= 'w'

polymorphicTests = testGroup "Polymorphic monadic values"
  [ testCase "MonadReader WriterT" $ do
      run (runReader 'c' (W.runWriterT polyReader1)) @?= ('c', ())
  , testCase "MonadReader ReaderT" $ do
      run (runReader 'c' (runReader False polyReader2)) @?= 'c'
  ]
  where
    polyReader1 :: MonadReader Char m => W.WriterT () m Char
    polyReader1 = ask
    polyReader2 :: MonadReader Char m => R.ReaderT Bool m Char
    polyReader2 = ask