{-# 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 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 a -> TcCodeM b -> TcCodeM (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 CodeState
mergeStatesDecl s1 s2 = do
  u <- getUniqRef
  liftIO $ mergeStates u s1 s2






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

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

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

instance Monad TcCodeM 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 where
  fmap = liftM

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

instance MonadIO TcCodeM where
  liftIO = liftTcDeclM . liftIO

instance MonadBase TcCodeM where
  liftBase = liftTcDeclM . liftBase
  withErrCtxt' ecf (TcCodeM f) = TcCodeM $ \e s -> withErrCtxt' ecf (f e s)
  tryM (TcCodeM f) = TcCodeM $ \e s -> do
                             esa <- tryM $ f e s
                             case esa of
                               Right (a, s', cs) -> return (Right a, s', cs)
                               Left err -> return (Left err, s, [])

instance MonadPR TcCodeM where
  liftPR = liftTcDeclM . liftPR

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

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


-- The environment

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

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


-- The state

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

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

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

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

-- Constraints

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

-- Exceptions
{-
tryCatch :: TcCodeM a -> (String -> TcCodeM a) -> TcCodeM a
tryCatch (TcCodeM f) ctch = TcCodeM $ \ec e s -> do
                              esa <- f ec e s
                              case esa of
                                Right a -> return $ Right a
                                Left err -> do
                                         let TcCodeM g = ctch err
                                         g ec e s
-}