{-# LANGUAGE ExistentialQuantification #-}
module OS where

-- TODO try avoiding double-matching for metamessages by closing over the callback
-- TODO try using PIDs instead of HIDs to avoid the search in runMetaHandler

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Control.Arrow ((>>>))
import Data.Maybe
import Control.Concurrent.Chan

data VM m mm = VM { hidCounter :: Int,
                    suspensions :: [Suspension m mm],
                    messages :: [m],
                    metaMessages :: [mm],
                    runnables :: [Callback m mm ()] }
type Callback m mm state = (state -> KernelModeTransition m mm)
data Subscription m mm = forall state. Subscription
                                        state
                                        (Maybe (Callback m mm state))
                                        [Handler m mm state]
                                        [MetaHandler m mm state]
data Suspension m mm = forall state. Suspension
                                        state
                                        (Maybe (Callback m mm state))
                                        [Handler m mm state]
                                        (IntMap (MetaHandler m mm state))
data KernelModeTransition m mm = KernelModeTransition (Subscription m mm)
                                                      [m]
                                                      [mm]
                                                      [Callback m mm ()]
type HID = Int
type Handler m mm state = (m -> Maybe (Callback m mm state))
type MetaHandler m mm state = (mm -> Maybe (Callback m mm state))

newVM boot = VM { hidCounter = 0,
                  suspensions = [],
                  messages = [],
                  metaMessages = [],
                  runnables = [boot] }

runVM :: VM m mm -> KernelModeTransition mm mm1
runVM = rebuildSuspensions enqueuePoller >>> runRunnables >>> dispatchMessages >>> \ vm ->
        let mms = reverse $ metaMessages vm
            mmhs = concatMap extractDown $ suspensions vm
            pollerK = if shouldPoll vm then Just runVM else Nothing
        in
          KernelModeTransition (Subscription (vm { metaMessages = [] }) pollerK mmhs []) mms [] []

rebuildSuspensions :: (VM m mm -> Suspension m mm -> VM m mm) -> (VM m mm) -> (VM m mm)
rebuildSuspensions f vm = foldl f (vm { suspensions = [] }) (suspensions vm)

enqueuePoller :: VM m mm -> Suspension m mm -> VM m mm
enqueuePoller vm susp@(Suspension state k _ _) =
    if isPolling susp
    then vm { runnables = (\ () -> (fromJust k) state) : runnables vm }
    else vm { suspensions = susp : suspensions vm }

runRunnables :: VM m mm -> VM m mm
runRunnables vm = foldl runRunnable (vm { runnables = [] }) (reverse (runnables vm))
runRunnable vm r = performTransition (r ()) vm

dispatchMessages :: VM m mm -> VM m mm
dispatchMessages vm = foldr dispatchMessage (vm { messages = [] }) (messages vm)

dispatchMessage m = rebuildSuspensions matchSuspension
    where matchSuspension vm susp@(Suspension state _ mhs _) =
              let searchHandlers [] = vm { suspensions = susp : suspensions vm }
                  searchHandlers (mh:mhs) =
                      case mh m of
                        Just h -> performTransition (h state) vm
                        Nothing -> searchHandlers mhs
              in searchHandlers mhs

extractDown (Suspension _ _ _ mmhs) =
    IntMap.foldWithKey (\ hid mmh xs -> dispatchMetaMessage hid mmh : xs) [] mmhs

dispatchMetaMessage :: HID -> MetaHandler m mm state -> Handler mm mm1 (VM m mm)
dispatchMetaMessage hid mmh mm =
    case mmh mm of
      Nothing -> Nothing
      Just _ -> Just (runMetaHandler hid mm)

isPolling (Suspension _ pollerK _ _) = isNothing pollerK
isBlocked = not . isPolling

shouldPoll vm@(VM { messages = [], runnables = []}) = not $ all isBlocked (suspensions vm)
shouldPoll _ = True

runMetaHandler :: HID -> mm -> (VM m mm) -> (KernelModeTransition mm mm1)
runMetaHandler hid mm = runVM . rebuildSuspensions runHid
    where runHid vm susp@(Suspension state _ _ mmhs) =
              case IntMap.lookup hid mmhs of
                Just h -> performTransition ((fromJust $ h mm) state) vm
                Nothing -> vm { suspensions = susp : suspensions vm }

addHids vm mmhs =
    foldr (\ mmh (vm, acc) -> let h = hidCounter vm in
                              (vm { hidCounter = h + 1 }, IntMap.insert h mmh acc))
          (vm, IntMap.empty)
          mmhs

performTransition :: KernelModeTransition m mm -> VM m mm -> VM m mm
performTransition (KernelModeTransition (Subscription state pollerK mhs mmhs) ms mms cbs) vm =
    let (vm, hidMap) = addHids vm mmhs in
    vm { suspensions = (Suspension state pollerK mhs hidMap) : suspensions vm,
         messages = reverse ms ++ (messages vm),
         metaMessages = reverse mms ++ (metaMessages vm),
         runnables = reverse cbs ++ (runnables vm) }

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

type LabelledMessage a b = (a, b)

groundVM :: Callback m (IO ()) () -> IO ()
groundVM boot = do inboundChannel <- newChan
                   mainloop inboundChannel (runVM (makeVM boot))
    where mainloop ch (KernelModeTransition (Subscription newState pollingK mhs []) ms [] []) = do
            runActions ch ms
            case (newState, pollingK, mhs) of
              (VM { messages = [], metaMessages = [], runnables = [] }, Nothing, []) ->
                  -- inert.
                  return ()
              _ ->
                  
                  
    (when (not (nested-vm-inert? (kernel-mode-transition-subscription transition)))
      (match transition
	[(kernel-mode-transition (subscription new-state
					       polling-k
					       message-handlers
					       '())
				 _
				 '()
				 '())
	 (define inbound-messages
	   (map (match-lambda [(message-handler e k) (wrap-evt e (lambda (v) (cons v k)))])
		message-handlers))
	 (match-define (cons inbound-value inbound-continuation)
	   (apply sync
		  (wrap-evt (if polling-k always-evt never-evt)
			    (lambda (v) (cons (void)
					      (lambda (dummy) polling-k))))
		  inbound-messages))
	 (loop ((inbound-continuation inbound-value) new-state))]
	[_
	 (error 'ground-vm
		"Outermost VM may not spawn new siblings or send or receive metamessages")]))))