module Driver 
    ( Transition
    , DriverState
    , applyTransition
    , initDriverState
    , takeState
    ) where

import When

import Control.Concurrent
import Control.Concurrent.MVar
import Data.Maybe (catMaybes)
import Data.Time.Clock
import Data.Time.Format()
import qualified Data.Map as M

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

type Transition st a 
    = st -> ([a], [Timed a], UTCTime -> IO (), st)

newtype DriverState st a 
    = DriverState (MVar (M.Map a ThreadId, st))

-----------------
 
initDriverState :: st -> IO (DriverState st a)
initDriverState st = do
    state <- newMVar (M.empty, st)
    return $ DriverState state

takeState :: DriverState st a -> IO ([a], st)
takeState (DriverState dst) = do
    (ma, st) <- takeMVar dst
    return (M.keys ma, st)

applyTransition
    :: (Eq a, Ord a, Show a) 
    => DriverState st a
    -> (a -> Transition st a) 
    -> Transition st a
    -> IO ()

applyTransition (DriverState dst) action transition     = appTransition Nothing transition     
 where
    appTransition this_action tr = do
        time <- getCurrentTime
        res  <- modifyMVar dst $ \(am, st) -> do

            let (cs, as, res, new_st) = tr st

            let am' = maybe am (`M.delete` am) this_action

            mapM_ killThread $ catMaybes $ map (`M.lookup` am') $ cs ++ map snd as

            let am'' = foldl (flip M.delete) am' cs

            tids <- mapM (forkIO . initAction time) as

            let am''' = M.fromList (zip (map snd as) tids) `M.union` am'' 

            return ((am''', new_st), res)

        res time 

    initAction time (when, a) = do 
        case when of
            Now         -> return ()
            AfterEval x -> x `seq` return ()
            At t        -> waitTill t
            Later dt    -> waitTill (addUTCTime dt time)
        appTransition (Just a) (action a)


waitTill time = do
    time' <- getCurrentTime
    case round $ 1000000 * (time `diffUTCTime` time') of
        n | n > 0   -> threadDelay n
        _           -> return ()