{-
 -  ``Control/Monad/Event/EventT''
 -}
{-# LANGUAGE
    GeneralizedNewtypeDeriving,
    MultiParamTypeClasses,
    FlexibleContexts,
    FlexibleInstances,
    UndecidableInstances,
    TypeSynonymInstances,
    KindSignatures
  #-}

module Control.Monad.EventT
    ( 
      EventT
    , runEventT
    , runEventGraph
    , runEventGraphWithState
    
    , EventT_RState, EventT_RWState
    , newEventT_RState
    , newEventT_RWState
    
    , HandlerAccessor
    , 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.RWS
import Control.Monad.Loops

import Data.Typeable

import Text.PrettyPrint.HughesPJ

import IO
import Control.Monad.Trans
import Control.Concurrent (threadDelay)

-- imports just for classes to implement
import Control.Monad.Error.Class    (MonadError(..))
import Control.Monad                (MonadPlus(..))
import Control.Monad.Cont.Class     (MonadCont(..))
import Control.Monad.Fix            (MonadFix(..))

import qualified Data.Map as M

{- Time management -}
instance Monad m => MonadTime (EventT t m) t where
    getCurrentTime = EventT (gets currentTime)

{- Time Management (internal) -}
setCurrentTime :: Monad m => t -> EventT t m ()
setCurrentTime t1 = do
    rwState <- EventT get
    
    let t0 = currentTime rwState
    EventT (put (rwState {currentTime = t1}))
    
    invokeHandler clockChangedHandler (t0,t1)

{- simulation control -}
instance Monad m => MonadSimControl (EventT t m) where
    resumeSimulation = do
        rwState <- EventT get
        EventT (put (rwState {simRunning = True}))
    pauseSimulation = do
        rwState <- EventT get
        EventT (put (rwState {simRunning = False}))
    isSimulationRunning = EventT (gets simRunning)

{- MonadEvent instance for unadorned (EventT t m) actions -}
instance (Monad m, Typeable (EventT t m a)) => MonadEvent (EventT t m) (EventT t m a) where
    describeEvent e = return (text "Undocumented Event")
    runEvent e = e >> return ()

{- MonadEvent instance for Event Descriptors - this does most of 
    the real grunt work of running an event  -}
instance (Monad m, Show t, Typeable t, Typeable1 m) => MonadEvent (EventT t m) (EventDescriptor (EventT t m) 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 eventDispatchHandler event
        
        EventT (local (\env -> env {currentEvent = Just event}) (unWrapEventT (runEvent e)))
        
        invokeHandler eventCompleteHandler event

{- event scheduling, canceling, etc. -}
instance (Monad m, Ord t, MonadEvent (EventT t m) e) => ScheduleEvent (EventT t m) t e where
    scheduleEventAt t e = do
        eid <- getNextEventId
        let event = EventDescriptor 
                {eventId = eid, eventTime = t, event = e}
        
        enqueue eventQueueRef event
        invokeHandler eventScheduleHandler event
        
        return eid

instance (Monad m, Ord t) => CancelEvent (EventT t m) t where
    cancelEvent eid = do
        maybeEvent <- dequeueWhere eventQueueRef (\e -> eventId e == eid)
        
        invokeHandler eventCancelHandler (maybe (Left eid) Right maybeEvent)
        return maybeEvent

{- services for use in implementing events -}
instance Monad m => GetCurrentEvent (EventT t m) t where
    getCurrentEvent = EventT (asks currentEvent)

instance Monad m => RetryEvent (EventT t m) 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 (Monad m, Ord t) => MonadEventQueueInfo (EventT t m) t where
    eventQueueSize = withEventQueue queueSize
    eventQueueContents = withEventQueue peekQueue


{- The EventT type and related administrative bits (state management stuff) -}
-- |A monad transformer which adds an event queue and related operations to
-- an underlying monad.  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.
newtype EventT t m a = EventT { unWrapEventT :: RWST (EventT_RState t m) () (EventT_RWState t m) m a }
    deriving (Functor, Monad, MonadIO)

tyCon_EventT = mkTyCon "Control.Monad.EventT.EventT"
instance (Typeable t, Typeable1 m) => Typeable1 (EventT t m) where
    typeOf1 y = mkTyConApp tyCon_EventT 
        [ typeOf  ((undefined :: EventT t m a -> t)    y)
        , typeOf1 ((undefined :: EventT t m a -> m ()) y)
        ]

instance MonadTrans (EventT t) where
    lift = EventT . lift

{- Monad type-class pass-through implementations -}
instance MonadReader r m => MonadReader r (EventT t m) where
    ask = EventT (lift ask)
    local f (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        (a, rwState, w) <- lift (local f (runRWST x rState rwState))
        put rwState
        tell w
        return a

instance MonadWriter w m => MonadWriter w (EventT t m) where
    tell w = EventT (lift (tell w))
    listen (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        ((a, rwState, w), w2) <- lift (listen (runRWST x rState rwState))
        put rwState
        tell w
        return (a, w2)
    pass (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        let knot = do
            ((a, f), rwState, w) <- runRWST x rState rwState 
            return ((a, rwState, w), f)
        (a, rwState, w) <- lift (pass knot)
        put rwState
        tell w
        return a

instance MonadState s m => MonadState s (EventT t m) where
    get = EventT (lift get)
    put s = EventT (lift (put s))

instance MonadError e m => MonadError e (EventT t m) where
    throwError e = EventT (lift (throwError e))
    catchError (EventT x) onErr = EventT $ do
        rState  <- ask
        rwState <- get
        let onErr' err = runRWST (unWrapEventT (onErr err)) rState rwState
        (a, rwState, w) <- lift (catchError (runRWST x rState rwState) onErr')
        put rwState
        tell w
        return a

-- is this valid?
instance MonadPlus m => MonadPlus (EventT t m) where
    mzero = EventT (lift mzero)
    mplus (EventT a) (EventT b) = EventT $ do
        rState  <- ask
        rwState <- get
        (c, rwState, w) <- lift (runRWST a rState rwState `mplus` runRWST b rState rwState)
        put rwState
        tell w
        return c

-- what semantics does this inherit?  is it sensible?  is it unique?
instance MonadCont m => MonadCont (EventT t m) where
    callCC f = EventT $
        callCC $ \cont ->
            unWrapEventT (f (EventT . cont))

instance MonadFix m => MonadFix (EventT t m) where
    mfix f = EventT $
        mfix $ \x ->
            unWrapEventT (f x)

{- Running (EventT t m) 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.
runEventT :: Monad m => EventT t m a -> EventT_RState t m -> EventT_RWState t m -> m (a, EventT_RWState t m)
runEventT (EventT x) rState rwState = do
    (a, rwState, _) <- runRWST x rState rwState
    return (a, rwState)

-- |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 :: 
    ( Monad m, Ord t, Show t
    , Typeable t, Typeable1 m
    ) => EventT_RState t m -> EventT_RWState t m -> m (EventT_RWState t m)
runEventGraphWithState rState rwState = do
    (_, rwState) <- runEventT (whileJust_ dequeueNextEvent runEvent) rState rwState
    return rwState

-- |Initialize the event queue and other stuff, enqueue the provided \"start 
-- event\", and run the queue until it's empty or until the simulation is
-- paused.
runEventGraph ::
    ( Monad m, MonadEvent (EventT t m) e
    , Ord t, Num t, Show t
    , Typeable t, Typeable1 m
    ) => e -> m (EventT_RState t m, EventT_RWState t m)
runEventGraph e = do
    let rState  = newEventT_RState
        rwState = newEventT_RWState 0
    
    (_, rwState) <- runEventT (scheduleEventIn 0 e) rState rwState
    rwState <- runEventGraphWithState rState rwState
    
    return (rState, rwState)


{- the main state vectors -}
data EventT_RState t m = EventT_RState
    { currentEvent  :: Maybe (EventDescriptor (EventT t m) t)
    }

data EventT_RWState t m = EventT_RWState
    { currentTime   :: t
    , simRunning    :: Bool
    , nextEventId   :: EventID
    , eventQueue    :: EventQueue t m
    , handlers      :: EventTHandlers t m
    }

-- |A new instance of the read/write portion of the 'EventT' internal
-- state.  The parameter is the initial time value.
newEventT_RWState :: (Monad m, Ord t) => t -> EventT_RWState t m
newEventT_RWState t = EventT_RWState
    { currentTime   = t
    , simRunning    = True
    , nextEventId   = EventID 1
    , eventQueue    = newEventQueue
    , handlers      = newEventTHandlers
    }

-- |A new instance of the read-only portion of the 'EventT' internal
-- state.
newEventT_RState :: EventT_RState t m
newEventT_RState = EventT_RState
    { currentEvent  = Nothing
    }

{- Miscellaneous small state management functions -}
getNextEventId :: Monad m => EventT t m EventID
getNextEventId = do
    rwState <- EventT get
    let eID = nextEventId rwState
    EventT (put (rwState {nextEventId = succ eID}))
    return eID

{- Support for debugging event handlers -}
data EventTHandlers t m = EventTHandlers
    { eventScheduleHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , eventCancelHandler    :: HandlerSet (EventT t m) (Either EventID (EventDescriptor (EventT t m) t)) ()
    , eventDispatchHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , eventCompleteHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , clockChangedHandler   :: HandlerSet (EventT t m) (t, t) () -- (old time, new time)
    }

type HandlerAccessor t m a b = 
    ( EventTHandlers t m -> HandlerSet (EventT t m) a b
    , HandlerSet (EventT t m) a b -> EventTHandlers t m -> EventTHandlers t m
    )

-- |Fires after an event is scheduled.  Is passed an 'EventDescriptor' for
-- the event.
onEventSchedule :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventSchedule = (eventScheduleHandler, \h hs -> hs {eventScheduleHandler = h})

-- |Fires after an event is canceled.  Is passed either an 'EventID' (if the
-- cancellation failed) or an 'EventDescriptor' for the event that was canceled.
onEventCancel   :: HandlerAccessor t m (Either EventID (EventDescriptor (EventT t m) t)) ()
onEventCancel = (eventCancelHandler, \h hs -> hs {eventCancelHandler = h})

-- |Fires just before an event is dispatched.  Is passed an 'EventDescriptor'
-- describing the event about to be run.
onEventDispatch :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventDispatch = (eventDispatchHandler, \h hs -> hs {eventDispatchHandler = h})

-- |Fires after an event returns.  Is passed an 'EventDescriptor' for the
-- event that just finished.
onEventComplete :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventComplete = (eventCompleteHandler, \h hs -> hs {eventCompleteHandler = h})

-- |Fires whenever the clock changes, and is passed a tuple containing
-- (old time, new time)
onClockChanged  :: HandlerAccessor t m (t, t) ()
onClockChanged = (clockChangedHandler, \h hs -> hs {clockChangedHandler = h})

newEventTHandlers :: Monad m => EventTHandlers t m
newEventTHandlers = EventTHandlers
    { eventScheduleHandler  = emptyHandlerSet
    , eventCancelHandler    = emptyHandlerSet
    , eventDispatchHandler  = emptyHandlerSet
    , eventCompleteHandler  = emptyHandlerSet
    , clockChangedHandler   = emptyHandlerSet
    }

-- |Add an event handler to be called when the specified event happens.
addHandler :: Monad m => HandlerAccessor t m a b -> (a -> EventT t m b) -> EventT t m HandlerID
addHandler (getter, setter) h = do
    hSet <- EventT (gets (getter.handlers))
    
    let (newHSet, hId) = addHandlerToSet h hSet
    
    EventT (modify (\s -> s {handlers = setter newHSet (handlers s)}))
    return hId

-- |Remove an event handler given its ID, and return it if it was in the set.
removeHandler :: Monad m => HandlerAccessor t m a b -> HandlerID -> EventT t m (Maybe (a -> EventT t m b))
removeHandler (getter, setter) hId = do
    hSet <- EventT (gets (getter.handlers))
    
    let (newHSet, h) = removeHandlerFromSet hId hSet
    
    EventT (modify (\s -> s {handlers = setter newHSet (handlers s)}))
    return h

invokeHandler :: Monad m => (EventTHandlers t m -> HandlerSet (EventT t m) a b) -> a -> EventT t m b
invokeHandler h args = do
    hSet <- EventT (gets (h.handlers))
    invokeHandlers hSet args


{- internal support relating to the event queue -}
type EventQueue t m = PQ (EventDescriptor (EventT t m) t)
type EventQueueRef t m = PriorityQueue (EventT t m) (EventDescriptor (EventT t m) t)

newEventQueue :: Ord t => EventQueue t m
newEventQueue = emptyPQ eventTime

withEventQueue :: (Monad m, Ord t) => (EventQueueRef t m -> EventT t m a) -> EventT t m a
withEventQueue f = f eventQueueRef

eventQueueRef :: (Monad m, Ord t) => EventQueueRef t m
eventQueueRef = mkPriorityQueue (UnsafeModifyRef (Accessor (Getter getEventQueue, Setter setEventQueue)))
    where 
        getEventQueue = EventT (gets eventQueue)
        setEventQueue eq = EventT $ do
            modify (\rwState -> rwState {eventQueue = eq})

dequeueNextEvent :: (Monad m, Ord t) => EventT t m (Maybe (EventDescriptor (EventT t m) t))
dequeueNextEvent = do
    running <- isSimulationRunning
    if running
        then withEventQueue dequeue
        else return Nothing

{- Some handy event handlers for logging and speed regulation -}
{- at some point these may be extracted to a type class -}
-- |(warning: very likely to change or disappear)
-- 
-- A collection of 'Bool's telling 'SetDebugHandlers' which handlers to install.
data DebugHandlerOptions = DebugHandlerOptions
    { -- |Run roughly in sync with the wall clock (using seconds as the simulation time unit)
      sync              :: Bool
    , -- |Print event dispatches to 'stderr' as they occur
      showDispatches    :: Bool 
    , -- |Print event cancellations to 'stderr' as they occur
      showCancels       :: Bool
    , -- |Print event scheduling actions to 'stderr' as they occur
      showSchedules     :: Bool
    , -- |Dump the entire contents of the event queue to
      -- 'stderr' every time an event is dispatched
      dumpQueue         :: Bool
    } deriving Show

defaultHandlerOptions = DebugHandlerOptions
    { sync              = False
    , showDispatches    = False
    , showCancels       = False
    , showSchedules     = False
    , dumpQueue         = False
    }