{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GADTs #-}

import Data.Time.Clock (UTCTime, diffUTCTime, getCurrentTime)
import Control.Monad.ST
import Data.IORef

-- warning: clunky Free monad implementation
data Free f a = Pure a | Run (f (Free f a))
instance (Functor f) => Monad (Free f) where
  return = Pure
  (>>=) (Pure a) f = f a
  (>>=) (Run r) f = Run $ fmap (>>= f) r
instance (Functor f) => Functor (Free f) where
  fmap = undefined
instance (Functor f) => Applicative (Free f) where
  (<*>) = undefined
  pure = return
liftF :: (Functor f) => f a -> Free f a
liftF x = Run (fmap Pure x)


data SomePromise where
    SomePromise :: IORef (Promise a) -> SomePromise

data PromiseF a
    = PLift (IO a) -- run an effect and yield
    | PQueue SomePromise (() -> a) -- add a new anonymous thread to the event loop
    deriving Functor
type Promise = Free PromiseF

liftP :: IO a -> Promise a
liftP = liftF . PLift

data Handle a = Handle (Promise (Maybe a))
scan (Handle handle) = handle
fork :: Promise a -> Promise (Handle a)
fork p = do
    forkRef <- liftP $ newIORef p
    liftF $ PQueue (SomePromise forkRef) id
    return (Handle $ do
        p' <- liftP $ readIORef forkRef
        return $ case p' of
            Pure a -> Just a
            Run _ -> Nothing)

step :: IORef [SomePromise] -> IORef (Promise a) -> IO (Maybe a)
step threadsRef threadRef = do
    r <- readIORef threadRef
    case r of
        Pure a -> return (Just a) -- thread done!
        Run (PQueue t c) -> do
            modifyIORef threadsRef (t :)
            writeIORef threadRef (c ())
            return Nothing
        Run (PLift mr) -> (mr >>= writeIORef threadRef) >> return Nothing

runPromise :: Promise a -> IO a
runPromise root = do
    threadsRef <- newIORef []
    rootRef <- newIORef root
    let forever = do
            r <- step threadsRef rootRef
            case r of
                Just a -> return a
                Nothing -> do
                    -- update the other threads and keep going
                    threads <- readIORef threadsRef
                    foldr (\(SomePromise p) acc -> step threadsRef p >> acc >> return ())
                        (return ()) threads
                    forever
    forever

---------------------------------------------------------------------------------------------

await :: Handle a -> Promise a
await handle = do
    ma <- scan handle
    case ma of
        Just a -> return a -- done!
        Nothing -> await handle -- take a break and try again later

testAsync :: Float -> String -> Promise String
testAsync t n = do
    liftP $ putStrLn (n ++ " called")
    timeout t
    liftP $ putStrLn (n ++ " did some stuff, it's halfway through now")
    timeout t
    return (n ++ " has finished")

test :: Promise ()
test = do
    liftP $ putStrLn "Hello World"
    handle1 <- fork $ testAsync 2 "test2"
    handle2 <- fork $ testAsync 3 "test3"
    liftP $ putStrLn "Tests invoked"
    await handle1
    await handle2
    liftP $ putStrLn "All done, yay!"
    return ()

timeout :: Float -> Promise () -- empty promise; yields while waiting, resolves once finished
timeout length = do
    start <- liftP $ getCurrentTime
    let waitUntil = do
            curr <- liftP $ getCurrentTime
            if realToFrac (diffUTCTime curr start) > length then return ()
            else waitUntil -- not done yet; allow eventloop to work on other promises
    waitUntil

main = runPromise test