{-
 -      ``Control/Monad/EventM''
 -      (c) 2009 Cook, J. MR  SSD, Inc.
 -}
{-# LANGUAGE
    GeneralizedNewtypeDeriving,
    MultiParamTypeClasses,
    FlexibleContexts,
    FlexibleInstances,
    UndecidableInstances,
    TypeSynonymInstances,
    KindSignatures
  #-}


module Control.Monad.EventM
    ( 
      EventM
    , EventIO
    , runEventIO
    , runEventGraph
    , runEventGraphWithState
    
    , newEventIOState
    
    , onClockChanged
    , onEventDispatch
    , onEventComplete
    , onEventSchedule
    , onEventCancel
    
    , addHandler
    , removeHandler
    ) where


import Control.Monad.Event.Classes
import Control.Monad.Event.Internal.Types

import Data.Handler
import Data.PriorityQueue
import Data.StateRef
import Control.Monad.Reader
import Control.Monad.Loops

import Text.PrettyPrint.HughesPJ

{- Time management -}
instance MonadTime (EventIO t) t where
    getCurrentTime = EventIO (asks currentTime) >>= readReference

{- Time Management (internal) -}
setCurrentTime :: t -> EventIO t ()
setCurrentTime t1 = do
    state <- EventIO ask
    t0 <- atomicModifyReference (currentTime state) (\t0 -> (t1, t0))
    
    invokeHandler onClockChanged (t0,t1)

{- simulation control -}
instance MonadSimControl (EventIO t) where
    resumeSimulation = do
        state <- EventIO ask
        writeReference (simRunning state) True
    pauseSimulation = do
        state <- EventIO ask
        writeReference (simRunning state) False
    isSimulationRunning = EventIO (asks simRunning) >>= readReference

{- MonadEvent instance for unadorned (EventIO t) actions -}
instance MonadEvent (EventIO t) (EventIO t a) where
    describeEvent e = return (text "Undocumented Event")
    runEvent e = e >> return ()

instance MonadEvent (EventIO t) (IO a) where
    describeEvent e = EventIO . lift . return $ text "Undocumented Event"
    runEvent e = EventIO . lift $ (e >> return ())

{- MonadEvent instance for Event Descriptors - this does most of 
    the real grunt work of running an event  -}
instance Show t => MonadEvent (EventIO t) (EventDescriptor (EventIO t) t) where
    describeEvent (EventDescriptor {eventId = eid, eventTime = t, event = e}) = do
        eventDescription <- describeEvent e
        return (
            brackets (
                text (fill 5 (show eid))
                <> text "|"
                <> text (fill 10 (show t))
            ) <> colon
            <+> eventDescription
         )
         
         where  fill    n     []  = replicate n ' '
                fill    0     xs  = xs
                fill (n+1) (x:xs) = x : fill n xs
    runEvent event@(EventDescriptor {event = e}) = do
        setCurrentTime (eventTime event)
        
        invokeHandler onEventDispatch event
        
        EventIO (local (\env -> env {currentEvent = Just event}) (unWrapEventIO (runEvent e)))
        
        invokeHandler onEventComplete event

{- event scheduling, canceling, etc. -}
instance (MonadEvent (EventIO t) e) => ScheduleEvent (EventIO t) t e where
    scheduleEventAt t e = do
        eid <- getNextEventId
        let event = EventDescriptor 
                {eventId = eid, eventTime = t, event = e}
        
        q <- getEventQueue
        enqueue q event
        invokeHandler onEventSchedule event
        
        return eid

instance CancelEvent (EventIO t) t where
    cancelEvent eid = do
        q <- getEventQueue
        maybeEvent <- dequeueWhere q (\e -> eventId e == eid)
        
        invokeHandler onEventCancel (maybe (Left eid) Right maybeEvent)
        return maybeEvent

{- services for use in implementing events -}
instance GetCurrentEvent (EventIO t) t where
    getCurrentEvent = EventIO (asks currentEvent)

instance RetryEvent (EventIO t) t where
    retryEventAt t = do
        maybeCe <- getCurrentEvent
        case maybeCe of
            Nothing -> fail "retry called outside an event"
            Just (EventDescriptor {event = e}) ->
                scheduleEventAt t e

instance MonadEventQueueInfo (EventIO t) t where
    eventQueueSize     = getEventQueue >>= queueSize
    eventQueueContents = getEventQueue >>= peekQueue


{- The EventT type and related administrative bits (state management stuff) -}
-- |A monad which extends IO with an event queue and related operations.
-- The \"t\" parameter specifies the type of the simulation time.
--
-- Several hooks are provided to allow special handling of various events,
-- such as the progression of time, the scheduling or canceling or dispatch
-- of an event, etc.
instance HasRef (EventIO t) where
    newRef x = fmap Ref ((newReference :: a -> EventIO t (IORef a)) x)

instance NewRef    (Ref IO a) (EventIO t) a where
    newReference = liftIO . newRef
instance ReadRef   (Ref IO a) (EventIO t) a where
    readReference = liftIO . readReference
instance WriteRef  (Ref IO a) (EventIO t) a where
    writeReference r = liftIO . writeReference r
instance ModifyRef (Ref IO a) (EventIO t) a where
    modifyReference r = liftIO . modifyReference r
    atomicModifyReference r = liftIO . atomicModifyReference r

instance NewRef    (Ref (EventIO t) a) IO a where
    newReference x = do
        r <- newReference x `asTypeOf` (undefined :: IO (IORef a))
        return (Ref r)

type EventM = EventIO Double
newtype EventIO t a = EventIO { unWrapEventIO :: ReaderT (EventIOState t) IO a }
    deriving (Functor, Monad, MonadIO, MonadFix)

{- Running (EventIO t) actions and whole event graphs -}
-- |Run an 'EventT' wrapped action.  This is a \"raw\" action - there is no 
-- interaction with the state (including event graph) except whatever the
-- action itself does.
runEventIO :: EventIO t a -> EventIOState t -> IO a
runEventIO (EventIO x) state = runReaderT x state

-- |Repeatedly pull and run the next event in the queue until it's 
-- empty or until the simulation is paused using 'pauseSimulation'
-- or something equivalent.
runEventGraphWithState :: (Ord t, Show t) => EventIOState t -> IO ()
runEventGraphWithState state = runEventIO (whileJust_ dequeueNextEvent runEvent) state

-- |Initialize the event queue and other stuff, run the provided \"start 
-- event\", and run the queue until it's empty or until the simulation is
-- paused.
runEventGraph ::
    ( MonadEvent (EventIO t) e
    , Ord t, Num t, Show t
    ) => e -> IO (EventIOState t)
runEventGraph e = do
    state <- newEventIOState 0
    
    runEventIO (scheduleEventIn 0 e) state
    runEventGraphWithState state
    
    return state

{- the main state vectors -}
data EventIOState t = EventIOState
    { currentEvent  :: Maybe (EventDescriptor (EventIO t) t)
    , currentTime   :: Ref IO t
    , simRunning    :: Ref IO Bool
    , nextEventId   :: Ref IO EventID
    , eventQueue    :: EventQueue t
    , handlers      :: EventIOHandlers t
    }

newEventIOState :: Ord t => t -> IO (EventIOState t)
newEventIOState t = do
    currentTime   <- newRef t
    simRunning    <- newRef True
    nextEventId   <- newRef (EventID 1)
    eventQueue    <- newPriorityQueue eventTime
    handlers      <- newEventIOHandlers
    
    return $ EventIOState
        { currentEvent  = Nothing
        , currentTime   = currentTime
        , simRunning    = simRunning
        , nextEventId   = nextEventId
        , eventQueue    = eventQueue
        , handlers      = handlers
        }

{- Miscellaneous small state management functions -}
getNextEventId :: EventIO t EventID
getNextEventId = do
    state <- EventIO ask
    atomicModifyReference (nextEventId state) (\i -> (succ i, i))

{- Support for debugging event handlers -}
data EventIOHandlers t = EventIOHandlers
    { onEventSchedule  :: Ref IO (HandlerSet (EventIO t) (EventDescriptor (EventIO t) t) ())
    , onEventCancel    :: Ref IO (HandlerSet (EventIO t) (Either EventID (EventDescriptor (EventIO t) t)) ())
    , onEventDispatch  :: Ref IO (HandlerSet (EventIO t) (EventDescriptor (EventIO t) t) ())
    , onEventComplete  :: Ref IO (HandlerSet (EventIO t) (EventDescriptor (EventIO t) t) ())
    , onClockChanged   :: Ref IO (HandlerSet (EventIO t) (t, t) ()) -- (old time, new time)
    }

type HandlerAccessor t a b = EventIOHandlers t -> Ref IO (HandlerSet (EventIO t) a b)

newEventIOHandlers :: IO (EventIOHandlers t)
newEventIOHandlers = do
    onEventSchedule     <- newRef emptyHandlerSet
    onEventCancel       <- newRef emptyHandlerSet
    onEventDispatch     <- newRef emptyHandlerSet
    onEventComplete     <- newRef emptyHandlerSet
    onClockChanged      <- newRef emptyHandlerSet
    return $ EventIOHandlers
        { onEventSchedule   = onEventSchedule
        , onEventCancel     = onEventCancel
        , onEventDispatch   = onEventDispatch
        , onEventComplete   = onEventComplete
        , onClockChanged    = onClockChanged
        }

-- |Add an event handler to be called when the specified event happens.
addHandler :: HandlerAccessor t a b -> (a -> EventIO t b) -> EventIO t HandlerID
addHandler hSel h = do
    hSet <- EventIO (asks (hSel.handlers))
    atomicModifyReference hSet (addHandlerToSet h)

-- |Remove an event handler given its ID, and return it if it was in the set.
removeHandler :: HandlerAccessor t a b -> HandlerID -> EventIO t (Maybe (a -> EventIO t b))
removeHandler hSel hId = do
    hSet <- EventIO (asks (hSel.handlers))
    atomicModifyReference hSet (removeHandlerFromSet hId)

invokeHandler :: HandlerAccessor t a b -> a -> EventIO t b
invokeHandler h args = do
    hSet <- EventIO (asks (h.handlers)) >>= readReference
    invokeHandlers hSet args


{- internal support relating to the event queue -}
type EventQueue t = PriorityQueue (EventIO t) (EventDescriptor (EventIO t) t)

getEventQueue :: EventIO t (EventQueue t)
getEventQueue = EventIO (asks eventQueue)

dequeueNextEvent :: EventIO t (Maybe (EventDescriptor (EventIO t) t))
dequeueNextEvent = do
    running <- isSimulationRunning
    if running
        then getEventQueue >>= dequeue
        else return Nothing