{-# LANGUAGE TupleSections, BangPatterns #-}
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, 

     getPrefix, touchPrefix

    ) 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.Policy
--import Language.Java.Paragon.TypeCheck.TypeMap
--import Language.Java.Paragon.TypeCheck.Types

import Control.Monad 
import Control.Applicative
--import Control.Arrow (second)

import qualified Data.Map.Strict as Map
--import qualified Data.ByteString.Char8 as B

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

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


setupStartState :: TcDeclM CodeState
setupStartState = return emptyCodeState
{-
setupStartState = do
  tm <- getTypeMap
  let aMap = gatherActorInfo tm
      pMap = undefined -- gatherPolicyBounds tm
      iMap = undefined
      tMap = undefined
  return $ CodeState (VarMap aMap pMap iMap tMap) noMods Map.empty
-}
{-
gatherActorInfo :: TypeMap -> ActorMap -- Map Ident (AI { Bool ActorId } )
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 . (\(_,_,x) -> x)) $ 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

-- TODO: non-final policies should have bounds bottom/top
gatherPolicyBounds :: TypeMap -> Map (Name ()) ActorPolicyBounds
gatherPolicyBounds = gatherPolicyBounds' Nothing

    where gatherPolicyBounds' mPre tm =
              let pols = Map.assocs $ policies tm -- :: [(Ident, PrgPolicy)]
                  aMap = Map.fromList $ map (mkPols mPre $ fields tm) 
                         $ map (second RealPolicy) pols
                  tMap = gatherPolicyBoundsAux TName mPre
                           (Map.assocs $ Map.map (tMembers . (\(_,_,x) -> x)) $ types tm)
                  pMap = gatherPolicyBoundsAux PName mPre (Map.assocs $ packages tm)
              in foldl1 Map.union [aMap, tMap, pMap]

          mkPols :: Maybe (Name ())
                 -> Map (Ident ()) VarFieldSig 
                 -> (Ident (), ActorPolicy) 
                 -> (Name (), ActorPolicyBounds)
          mkPols mPre fs (i,p) = 
              case Map.lookup i fs of
                Just _ -> (Name () EName mPre i, KnownPolicy p)
                _ -> panic (tcCodeMonadModule ++ ".gatherActorInfo") $ 
                     "No field for corresponding actor " ++ show i

          gatherPolicyBoundsAux :: NameType 
                              -> Maybe (Name ()) 
                              -> [(Ident (), TypeMap)] 
                              -> Map (Name ()) ActorPolicyBounds
          gatherPolicyBoundsAux nt mPre = foldl Map.union Map.empty . map aux
              where aux :: (Ident (), TypeMap) -> Map (Name ()) ActorPolicyBounds
                    aux (i,tm) = 
                        let pre = Name () nt mPre i
                        in gatherPolicyBounds' (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' <- mergeStatesDeclM ts1 ts2
      return ((a,b), ts', cs1 ++ cs2)


mergeStatesDeclM :: Maybe CodeState -> Maybe CodeState -> TcDeclM (Maybe CodeState)
mergeStatesDeclM (Just s1) (Just s2) = Just <$> mergeStatesDecl s1 s2
mergeStatesDeclM Nothing a = return a
mergeStatesDeclM a Nothing = return a

mergeStatesDecl :: CodeState -> CodeState -> TcDeclM CodeState
mergeStatesDecl s1 s2 = liftBase $ mergeStates s1 s2



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

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

runTcCodeM :: CodeEnv -> CodeState -> TcCodeM a -> TcDeclM (a, [ConstraintWMsg])
runTcCodeM env state (TcCodeM f) = do
  (a,_,cs) <- f env (Just 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 = do
  ms <- getStateM
  case ms of
    Just s -> return s
    Nothing -> panic (tcCodeMonadModule ++ ".getState")
               $ "Calling getState in dead code analysis"

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

setState :: CodeState -> TcCodeM ()
setState s = TcCodeM (\_ _ -> return ((),Just 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)]))




-------------------------------------
-- Working with the VarMap
-------------------------------------


-- Using a zipper technique
touchPrefix :: Maybe (Name ()) -> TcCodeM (VarMap, VarMap -> CodeState)
touchPrefix mn = do
  case mn of
    Nothing -> do
            st <- getState
            return (varMapSt st, \vm -> st { varMapSt = vm })
    Just n@(Name _ nt mPre i) -> do
        (vm, vmf) <- touchPrefix mPre
        case nt of
          _ | nt `elem` [TName, PName] -> do
              let upd newVm = vmf $ vm { typesSt = Map.insert (unIdent i) newVm $ typesSt vm }
              case Map.lookup (unIdent i) $ typesSt vm of
                Just tvm -> return (tvm, upd)
                Nothing -> do setState $ upd emptyVM
                              return (emptyVM, upd)
          EName -> do
              let ist = instanceSt vm
              case Map.lookup (unIdent i) ist of
                Just ii -> 
                    let upd newVm = 
                            let newII = ii { iMembers = newVm }
                            in vmf $ vm { instanceSt = Map.insert (unIdent i) newII ist }
                    in return $ (iMembers ii, upd)
                Nothing -> panic (tcCodeMonadModule ++ ".touchPrefix")
                           $ "Prefix not in state: " ++ show n
          _ -> panic (tcCodeMonadModule ++ ".touchPrefix")
               $ "Unexpected name: " ++ show n
    _ -> panic (tcCodeMonadModule ++ ".touchPrefix")
         $ show mn

-- Using a zipper technique
getPrefix :: Maybe (Name ()) -> TcCodeM VarMap
getPrefix mn = do
  case mn of
    Nothing -> varMapSt <$> getState
    Just n@(Name _ nt mPre i) -> do
        !vm <- getPrefix mPre
--        (vm, vmf) <- touchPrefix mPre
        case nt of
          _ | nt `elem` [TName, PName] -> do
              case Map.lookup (unIdent i) $ typesSt vm of
                Just tvm -> return tvm -- , upd)
                Nothing -> do (_, vmf) <- touchPrefix mPre
                              let upd newVm = vmf $ vm { typesSt = Map.insert (unIdent i) newVm $ typesSt vm }
                              setState $ upd emptyVM
                              return emptyVM -- , upd)
          EName -> do
              let ist = instanceSt vm
              case Map.lookup (unIdent i) ist of
                Just ii -> 
--                    let upd newVm = 
--                            let newII = ii { iMembers = newVm }
--                            in vmf $ vm { instanceSt = Map.insert (unIdent i) newII ist }
--                    in 
                      return $ iMembers ii --, upd)
                Nothing -> panic (tcCodeMonadModule ++ ".getPrefix")
                           $ "Prefix not in state: " ++ show n
          _ -> panic (tcCodeMonadModule ++ ".getPrefix")
               $ "Unexpected name: " ++ show n
    _ -> panic (tcCodeMonadModule ++ ".getPrefix")
         $ show mn