{-# LANGUAGE ScopedTypeVariables, GADTs, KindSignatures #-}
module Distributed.Model.Simulate (
   simulateMachines
 ) where

import           Control.Monad.Trans.Class
import           Data.Foldable
import           Data.Machine.Is
import           Data.Machine.Plan
import           Data.Map (Map)
import qualified Data.Map as Map
import           Data.Random
import           Data.Random.List
import           Data.Random.RVar
import           Distributed.Model.Semantics
import           Safe

data SimStep addr p (m :: * -> *) a =
    SimCrashed
  | SimStopped a
  | SimYield (Msg addr p) (m (SimStep addr p m a))
  | SimAwait (Msg addr p -> (m (SimStep addr p m a))) (m (SimStep addr p m a))

simulateMachines :: forall m s id addr p a
                  . (Ord id, RandomSource m s)
                 => Map id (Msg addr p -> (id, Msg addr p), Model addr p m a)
                 -> s
                 -> m (Map id (Maybe a))
simulateMachines machines s = do
    simInitialState <- mapM startSim machines
    runRVarT (step simInitialState) s
  where
    startSim :: (Msg addr p -> (id, Msg addr p), Model addr p m a)
             -> m (Msg addr p -> (id, Msg addr p)
                  ,SimStep addr p m a
                  ,[Msg addr p])
    startSim (addrmp, m) = do
      s <- runPlanT m
             (return . SimStopped)
             (\o n -> return $ SimYield o n)
             (\f Refl g -> return $ SimAwait f g)
             (return SimCrashed)
      return (addrmp, s, mempty)
    step :: Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
         -> RVarT m (Map id (Maybe a))
    step sim =
      case simIsStopped sim of
        True -> return . fmap (\(_, s, _) -> case s of
                                               SimStopped a -> Just a
                                               SimCrashed -> Nothing)
                       $ sim
        False -> stepSim sim >>= step

simIsStopped :: Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
             -> Bool
simIsStopped = all isStopped

isStopped :: (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p]) -> Bool
isStopped (_, SimCrashed, _) = True
isStopped (_, SimStopped _, _) = True
isStopped _ = False

queueMsg :: Ord id
         => (id, Msg addr p)
         -> Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
         -> Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
queueMsg (destAddr, msg) =
  Map.adjust (\(msgmap, ss, msgbuff) -> (msgmap, ss, msg:msgbuff)) destAddr


stepMachine :: (Ord id, Monad m)
            => id -> m (SimStep addr p m a)
            -> Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
            -> m (Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p]))
stepMachine mid mss sim = do
 ss <- mss
 return . Map.adjust (\(msgmap, _, msgbuff) -> (msgmap, ss, msgbuff)) mid $ sim

couldSend :: (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p]) -> Bool
couldSend (_, SimCrashed, _) = False
couldSend (_, SimStopped _, _) = False
couldSend (_, SimYield _ _, _) = True
couldSend (_, SimAwait _ _, []) = False
couldSend (_, SimAwait _ _, _) = True

-- | Takes a specific machine in the Sim and steps it forward one step.
stepSim :: (Ord id, Monad m)
        => Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p])
        -> RVarT m (Map id (Msg addr p -> (id, Msg addr p), SimStep addr p m a, [Msg addr p]))
stepSim sim = do
  (mid, (msgmap, ss, msgbuffer)) <- randomElementT (Map.assocs sim)
  case ss of
    SimCrashed -> return sim
    SimStopped _ -> return sim
    SimYield msg mss -> do
      lift $ queueMsg (msgmap msg) <$> stepMachine mid mss sim
    SimAwait onRecv onNeverMsgAgain -> do
      randMsgBuff <- shuffleT msgbuffer
      case headMay randMsgBuff of
        -- If there are no messages waiting to be delivered, and all
        -- machines are stopped or awaiting, we are jammed.
        Nothing | any couldSend sim -> return sim
                | otherwise -> lift $ stepMachine mid onNeverMsgAgain sim
        Just msg -> do
          lift $ stepMachine mid (onRecv msg) (Map.adjust (\(msgmap, ss', _) -> (msgmap, ss', tail randMsgBuff)) mid sim)