{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Yield where

-- base
import Control.Monad (forever)
import Data.List.NonEmpty (NonEmpty, reverse)

-- transformers
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer (Writer, tell, runWriter, execWriter)

-- test-framework
import Test.Framework

-- test-framework-hunit
import Test.Framework.Providers.HUnit

-- HUnit
import Test.HUnit hiding (Test)

-- monad-schedule
import Control.Monad.Schedule.Class (scheduleAndFinish, schedule)
import Control.Monad.Schedule.Yield
import Control.Monad.Schedule.Trans (runScheduleT, runScheduleIO)
import Control.Monad.Trans.Reader
import Data.Foldable (forM_)
import Test.QuickCheck (NonEmptyList(NonEmpty), (===), (==>), counterexample)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (fromJust, isJust, maybeToList)
import Test.Framework.Providers.QuickCheck2 (testProperty)

sampleActions :: NonEmpty (MySchedule ())
sampleActions = [yield, yield]

tests = testGroup "Trans"
  [ testCase "Only leftover time is waited"
    $ assertRunsLike sampleActions [Yielded]
  , testCase "Scheduling two waits"
    $ assertRunsEqual sampleActions (Data.List.NonEmpty.reverse sampleActions)
  , testCase "Different number of waits"
    $ assertRunsLike
      [ myLog "Thread 1 starts" >> yield >> myLog "Thread 1 action" >> yield >> myLog "Thread 1 done"
      , myLog "Thread 2 starts" >> yield >> myLog "Thread 2 done"
      ]
      [ Log "Thread 1 starts"
      , Log "Thread 2 starts"
      , Yielded
      , Log "Thread 1 action"
      , Log "Thread 2 done"
      , Yielded
      , Log "Thread 1 done"
      ]
  , testCase "Blocking thread doesn't starve other thread"
    $ assertRunContains
      [ forever $ myLog "Busy loop starts" >> yield >> myLog "Busy loop ends"
      , myLog "One off thread starts" >> yield >> myLog "One off thread does a thing" >> yield >> myLog "One off thread done"
      ]
      $ Log "One off thread done"
  , testCase "Programs with continuations can be scheduled"
    $ assertProgramsInitiallyRunsLike
      [[Log "Thread 1 active",Log "Thread 2 active",Yielded],[Log "Thread 1 active",Log "Thread 2 active",Yielded],[Log "Thread 1 active",Log "Thread 2 active",Yielded],[Log "Thread 1 active",Log "Thread 2 active",Yielded],[Log "Thread 1 active",Log "Thread 2 active",Yielded]]
      [foreverP (const ["Thread 1 active"]), foreverP (const ["Thread 2 active"])]
      $ repeat True
  , testCase "Two programs that tick alternately can be scheduled"
    $ assertProgramsInitiallyRunsLike
      [[Log "1 Nope",Log "2 Yes",Yielded,Log "1 Nope"],[Log "2 Nope",Yielded,Log "2 Nope",Log "1 Yes",Yielded,Log "2 Nope"],[Log "1 Yes",Yielded,Log "2 Nope"],[Log "1 Nope",Yielded,Log "1 Nope",Log "2 Yes",Yielded,Log "1 Nope"],[Log "2 Yes",Yielded,Log "1 Nope"],[Log "2 Yes",Yielded,Log "1 Nope"],[Log "2 Yes",Yielded,Log "1 Nope"],[Log "2 Nope",Yielded,Log "2 Nope",Log "1 Yes",Yielded,Log "2 Nope"]]
      twoPrograms
      [True, False, False, True, True, True, True, False, False]
  , testProperty "Two programs that tick alternately can be scheduled with arbitrary input"
    $ \(inputs :: [Bool]) skip ->
    let log = take (20 * length inputs) $ drop skip $ concat $ runProgramWith inputs $ schedulePrograms twoPrograms
        isContained expectedEntry = expectedEntry `elem` log
    in counterexample (show log)
    $ all (`elem` drop skip inputs) ([True, False] :: [Bool]) ==> all isContained ([Log "1 Yes", Log "2 Yes"] :: [Event])
  ]

assertRunsEqual :: NonEmpty (MySchedule a1) -> NonEmpty (MySchedule a2) -> Assertion
assertRunsEqual actions1 actions2 = assertEqual "Should run the same under scheduling" (runMySchedule actions1) (runMySchedule actions2)

assertRunsLike :: NonEmpty (MySchedule a) -> [Event] -> Assertion
assertRunsLike actions events = assertEqual "Should run like the following under scheduling" events (runMySchedule actions)

assertRunContains :: NonEmpty (MySchedule a) -> Event -> Assertion
assertRunContains actions event = assertBool ("The run should contain the event " ++ show event) $ event `elem` runMySchedule actions

assertInitiallyRunsLike :: NonEmpty (MySchedule a) -> [Event] -> Assertion
assertInitiallyRunsLike actions events = assertEqual "Should, at the beginning, run like the following under scheduling" events $ take (length events) $ runMySchedule actions


data Event
  = Log String
  | Yielded
  deriving (Eq, Show)

type MySchedule a = YieldT (Writer [Event]) a

myLog :: String -> MySchedule ()
myLog = lift . tell . pure . Log

runMySchedule :: NonEmpty (MySchedule a) -> [Event]
runMySchedule = execWriter . runScheduleT (tell . pure . const Yielded) . scheduleAndFinish

type MyReaderSchedule a = YieldT (ReaderT Bool (Writer [Event])) a

-- Ok this is basically ListT
newtype Program = Program { unProgram :: MyReaderSchedule (Maybe Program) }

foreverP :: (Bool -> [String]) -> Program
foreverP action = go
  where
    go = Program $ do
      input <- lift ask
      lift $ lift $ tell $ Log <$> action input
      yield
      return $ Just go

-- Returning Left means looping further, Right means going to the next continuation
wait :: (Bool -> Either String String) -> Program
wait f = go
  where
    go = Program loop
    loop = do
      input <- lift ask
      case f input of
        Left msg -> do
          myReaderLog msg
          yield
          loop
        Right msg -> do
          myReaderLog msg
          yield
          return $ Just go

twoPrograms =
  [ wait (\b -> if b then Left "1 Nope" else Right "1 Yes")
  , wait (\b -> if b then Right "2 Yes" else Left "2 Nope")
  ]

runProgram :: Program -> MyReaderSchedule ()
runProgram Program { .. } = do
  tick <- unProgram
  forM_ tick runProgram

myReaderLog :: String -> MyReaderSchedule ()
myReaderLog = lift . lift . tell . pure . Log

schedulePrograms :: NonEmpty Program -> Program
schedulePrograms programs = Program $ do
  (done, running) <- schedule $ unProgram <$> programs
  return $ fmap schedulePrograms $ NonEmpty.nonEmpty $ (fromJust <$> NonEmpty.filter isJust done) ++ (Program <$> running)

-- Should be possible with some recursion scheme
runProgramWith :: [Bool] -> Program -> [[Event]]
runProgramWith [] _ = []
runProgramWith (input : inputs) Program { .. }
  = let (cont, events) = runWriter $ flip runReaderT input $ runScheduleT (const $ lift $ tell [Yielded]) unProgram
  in events : (runProgramWith inputs =<< maybeToList cont)

assertProgramsInitiallyRunsLike :: [[Event]] -> NonEmpty Program -> [Bool] -> Assertion
assertProgramsInitiallyRunsLike events programs inputs = assertEqual "The programs, when scheduled, should run like" events
  $ take (length events) $ runProgramWith inputs $ schedulePrograms programs