{-# LANGUAGE TupleSections #-}
module Language.Java.Paragon.TypeCheck.Monad.TcCodeM 
    (
     module Language.Java.Paragon.TypeCheck.Monad.TcDeclM,
     module Language.Java.Paragon.TypeCheck.Monad.CodeEnv,
     module Language.Java.Paragon.TypeCheck.Monad.CodeState,

     TcCodeM, (|||), runTcCodeM,
     getEnv, withEnv,                                  -- Reader
     getState, setState, updateState, mergeWithState,  -- State
     addConstraint,                                    -- Writer

     withTypeMapAlwaysC, 

     setupStartState, 

    ) where

import Language.Java.Paragon.Syntax
import Language.Java.Paragon.Interaction

import Language.Java.Paragon.TypeCheck.Monad.TcDeclM

import Language.Java.Paragon.TypeCheck.Monad.CodeEnv
import Language.Java.Paragon.TypeCheck.Monad.CodeState

import Language.Java.Paragon.TypeCheck.Constraints (ConstraintWMsg, Constraint)
import Language.Java.Paragon.TypeCheck.Actors
import Language.Java.Paragon.TypeCheck.Locks (noMods)
import Language.Java.Paragon.TypeCheck.TypeMap

import Control.Monad 
import Control.Applicative

import qualified Data.Map as Map

tcCodeMonadModule :: String
tcCodeMonadModule = typeCheckerBase ++ ".Monad.TcCodeM"

-------------------------------------------------
-- All the cool methods of this monad

setupStartState :: TcDeclM r CodeState
setupStartState = do
  tm <- getTypeMap
  let aMap = gatherActorInfo tm
  return $ CodeState aMap noMods Map.empty


gatherActorInfo :: TypeMap -> Map (Name ()) ActorInfo
gatherActorInfo = gatherActorInfo' Nothing

    where gatherActorInfo' mPre tm =
              let acts = Map.assocs $ actors tm -- :: [(Ident, ActorId)]
                  aMap = Map.fromList $ map (mkInfo mPre $ fields tm) acts
                  tMap = gatherActorInfoAux TName mPre 
                           (Map.assocs $ Map.map (tMembers . snd) $ types tm)
                  pMap = gatherActorInfoAux PName mPre (Map.assocs $ packages tm)
              in foldl1 Map.union [aMap, tMap, pMap]

          mkStab :: VarFieldSig -> Stability
          mkStab (VSig _ _ _ _ final) = 
              if final then Stable else FieldV Nothing

          mkInfo :: Maybe (Name ())
                 -> Map (Ident ()) VarFieldSig 
                 -> (Ident (), ActorId) 
                 -> (Name (), ActorInfo)
          mkInfo mPre fs (i,aid) = 
              case Map.lookup i fs of
                Just vti -> (Name () EName mPre i, AI aid (mkStab vti))
                _ -> panic (tcCodeMonadModule ++ ".gatherActorInfo") $ 
                     "No field for corresponding actor " ++ show i

          gatherActorInfoAux :: NameType 
                             -> Maybe (Name ()) 
                             -> [(Ident (), TypeMap)] 
                             -> Map (Name ()) ActorInfo
          gatherActorInfoAux nt mPre = foldl Map.union Map.empty . map aux
              where aux :: (Ident (), TypeMap) -> Map (Name ()) ActorInfo
                    aux (i,tm) = 
                        let pre = Name () nt mPre i
                        in gatherActorInfo' (Just pre) tm


-- Running in parallel
infix 1 |||
(|||) :: TcCodeM r a -> TcCodeM r b -> TcCodeM r (a,b)
(TcCodeM f1) ||| (TcCodeM f2) = 
    TcCodeM $ \te ts -> do
      (a, ts1, cs1) <- f1 te ts
      (b, ts2, cs2) <- f2 te ts
      ts' <- mergeStatesDecl ts1 ts2
      return $ ((a,b), ts', cs1 ++ cs2)

mergeStatesDecl :: CodeState -> CodeState -> TcDeclM r CodeState
mergeStatesDecl s1 s2 = do
  u <- getUniqRef
  liftIO $ mergeStates u s1 s2






--------------------------------------------------
-- The monad used for typechecking code snippets

newtype TcCodeM r a = 
    TcCodeM (CodeEnv -> CodeState -> TcDeclM r (a, CodeState, [ConstraintWMsg]) )

runTcCodeM :: CodeEnv -> CodeState -> TcCodeM r a -> TcDeclM r (a, [ConstraintWMsg])
runTcCodeM env state (TcCodeM f) = do
  (a,_,cs) <- f env state
  return (a, cs)

instance Monad (TcCodeM r) where
  return x = TcCodeM $ \_ s -> return (x, s, [])

  TcCodeM f >>= h = TcCodeM $ \e s0 -> do
                    (a, s1, cs1) <- f e s0
                    let TcCodeM g = h a
                    (b, s2, cs2) <- g e s1
                    return (b, s2, cs1 ++ cs2)

  fail err = TcCodeM $ \_ _ -> fail err

instance Functor (TcCodeM r) where
  fmap = liftM

instance Applicative (TcCodeM r) where
  (<*>) = ap
  pure = return

instance MonadIO (TcCodeM r) where
  liftIO = liftTcDeclM . liftIO

instance MonadBase (TcCodeM r) where
  liftBase = liftTcDeclM . liftBase
  withErrCtxt' ecf (TcCodeM f) = TcCodeM $ \e s -> withErrCtxt' ecf (f e s)

instance MonadPR (TcCodeM r) where
  liftPR = liftTcDeclM . liftPR

instance MonadTcBaseM (TcCodeM r) where
  liftTcBaseM = liftTcDeclM . liftTcBaseM
  withTypeMap tmf (TcCodeM f) = TcCodeM $ \e s -> withTypeMap tmf (f e s)

instance MonadTcDeclM TcCodeM where
  liftTcDeclM tdm = TcCodeM $ \_ s -> (,s,[]) <$> tdm

liftCallCC :: ((((a, CodeState, [ConstraintWMsg]) -> TcDeclM r (b, CodeState, [ConstraintWMsg]))
                 -> TcDeclM r (a, CodeState, [ConstraintWMsg]))
                  -> TcDeclM r (a, CodeState, [ConstraintWMsg]))
              -> ((a -> TcCodeM r b) -> TcCodeM r a) -> TcCodeM r a
liftCallCC ccc k = TcCodeM $ \e s -> ccc $ \c ->
                       let (TcCodeM f) = k (\a -> TcCodeM $ \_ _ -> c (a, s, []))
                       in f e s

withTypeMapAlwaysC :: (TypeMap -> TypeMap) -> TcCodeM r a -> TcCodeM r a
withTypeMapAlwaysC tmf tcm = liftCallCC callCC $ \cont -> do
                               withTypeMap tmf $ tcm >>= cont



-- The environment

getEnv :: TcCodeM r CodeEnv
getEnv = TcCodeM (\e s -> return (e,s,[]))

withEnv :: (CodeEnv -> CodeEnv) -> TcCodeM r a -> TcCodeM r a
withEnv k (TcCodeM f) = TcCodeM (f . k)


-- The state

getState :: TcCodeM r CodeState
getState = TcCodeM (\_ s -> return (s,s,[]))

setState :: CodeState -> TcCodeM r ()
setState s = TcCodeM (\_ _ -> return ((),s,[]))

updateState :: (CodeState -> CodeState) -> TcCodeM r ()
updateState f = getState >>= return . f >>= setState

mergeWithState :: CodeState -> TcCodeM r ()
mergeWithState s = do
  sOld <- getState
  sNew <- liftTcDeclM $ mergeStatesDecl sOld s
  setState sNew

-- Constraints

addConstraint :: Constraint -> String -> TcCodeM r ()
addConstraint c str = TcCodeM (\_ s -> return ((), s, [(c,str)]))