{-# 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
     tryCatch,                                         -- Error

--     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 $ \ec te ts -> do
      esa <- f1 ec te ts
      esb <- f2 ec te ts
      case (esa, esb) of
        (Right (a, ts1, cs1), Right (b, ts2, cs2)) -> do
                          ts' <- mergeStatesDecl ts1 ts2
                          return $ Right ((a,b), ts', cs1 ++ cs2)
        (Left err, _) -> return $ Left err
        (_, Left err) -> return $ Left err

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 (ErrCtxt -> CodeEnv -> CodeState -> TcDeclM r (Either String (a, CodeState, [ConstraintWMsg])))

runTcCodeM :: CodeEnv -> CodeState -> TcCodeM r a -> TcDeclM r (a, [ConstraintWMsg])
runTcCodeM env state (TcCodeM f) = do
  esa <- f id env state
  case esa of
    Left err -> fail err
    Right (a,_,cs) -> return (a, cs)

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

  TcCodeM f >>= h = TcCodeM $ \ec e s0 -> do
                      esa <- f ec e s0
                      case esa of
                        Left err -> return $ Left err
                        Right (a, s1, cs1) -> do
                            let TcCodeM g = h a
                            esb <- g ec e s1
                            case esb of
                              Left err -> return $ Left err
                              Right (b, s2, cs2) -> return $ Right (b, s2, cs1 ++ cs2)

  fail err = TcCodeM $ \ec _ _ -> return (Left $ ec 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 $ \ec e s -> (f (ecf ec) e s)

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

instance MonadTcBaseM (TcCodeM r) 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 -> Right . (,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 $ Right (e,s,[]))

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


-- The state

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

setState :: CodeState -> TcCodeM r ()
setState s = TcCodeM (\_ _ _ -> return $ Right ((),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 $ Right ((), s, [(c,str)]))

-- Exceptions

tryCatch :: TcCodeM r a -> (String -> TcCodeM r a) -> TcCodeM r 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