module Language.Java.Paragon.TypeCheck.Monad.TcMonad where

import Language.Java.Paragon.Syntax
--import Language.Java.Paragon.Pretty

import Language.Java.Paragon.TypeCheck.Monad.TcBase
import Language.Java.Paragon.TypeCheck.TcEnv
import Language.Java.Paragon.TypeCheck.TcState
import Language.Java.Paragon.TypeCheck.Constraints
import Language.Java.Paragon.TypeCheck.Actors
import Language.Java.Paragon.TypeCheck.Locks (noMods)

import Control.Monad (liftM)         -- liftM only to instantiate fmap

import qualified Data.Map as Map

import Debug.Trace

------------------------------------------------
--               The Tc monad                 --
------------------------------------------------
--
-- A monad on top of TcBase for typechecking
-- on the level of method declarations.
-- Has a method environment, a mergeable state, 
-- and a writer for constraints

newtype Tc a =
    Tc (TcEnv -> TcState -> TcBase (a, TcState, [Constraint]))

runTc :: TcEnv -> TcState -> Tc a -> TcBase (a, [Constraint])
runTc env state (Tc f) = do
  (a,_,cs) <- f env state
  return (a, cs)

-------------------------------------
setupStartState :: TcBase TcState
setupStartState = do
  tm <- getTypeMap
  let aMap = gatherActorInfo tm
  return $ TcState aMap noMods Map.empty


gatherActorInfo :: TypeMap -> Map Name ActorInfo
gatherActorInfo = Map.mapKeysMonotonic Name . gatherActorInfoAux
  where
    gatherActorInfoAux :: TypeMap -> Map [Ident] ActorInfo
    gatherActorInfoAux tm = 
        --trace ("TRACE: " ++ show tm) $
        let acts = Map.assocs $ actors tm -- :: [(Ident, ActorId)]
            aMap = Map.fromList $ map (mkInfo $ fields tm) acts
            tMap = gatherActorInfoT (Map.assocs $ pkgsAndTypes tm)
        in Map.union aMap tMap

            where mkStab :: VTypeInfo -> Stability
                  mkStab (VTI _ _ _ final) = 
                       if final then Stable else FieldV (Name []) -- UGLY
                
                  mkInfo :: Map Ident VTypeInfo -> (Ident,ActorId) -> ([Ident], ActorInfo)
                  mkInfo fs (i,aid) = 
                    case Map.lookup i fs of
                      Just vti -> ([i], AI aid (mkStab vti))
                      _ -> error $ "Internal error: no field for corresponding actor " ++ show i

    gatherActorInfoT :: [(Ident, TypeMap)] -> Map [Ident] ActorInfo
    gatherActorInfoT = foldl Map.union Map.empty . map aux
        where aux :: (Ident, TypeMap) -> Map [Ident] ActorInfo
              aux (i,tm) = let aMap = gatherActorInfoAux tm
                           in Map.mapKeysMonotonic (i:) $ Map.map (extendT i) aMap
              extendT :: Ident -> ActorInfo -> ActorInfo
              extendT i (AI aid (FieldV (Name is))) = AI aid (FieldV (Name $ i:is))
              extendT _ a = a
      
-------------------------------------


instance Monad Tc where
  return x = Tc $ \_ s -> return (x, s, [])
  Tc f >>= k = Tc $ \e s0 -> do
                     (a, s1, cs1) <- f e s0
                     let Tc g = k a
                     (b, s2, cs2) <- g e s1
                     return (b, s2, cs1 ++ cs2)
  fail err = Tc $ \_ _ -> fail err


instance Functor Tc where
  fmap = liftM

liftBase :: TcBase a -> Tc a
liftBase tcba = Tc $ \_ s -> do
                  a <- tcba
                  return (a, s, [])

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

mergeStatesBase :: TcState -> TcState -> TcBase TcState
mergeStatesBase s1 s2 = do
  u <- getUniqRef
  liftIO $ mergeStates u s1 s2

-- The environment

getEnv :: Tc TcEnv
getEnv = Tc (\e s -> return (e,s,[]))

withEnv :: (TcEnv -> TcEnv) -> Tc a -> Tc a
withEnv k (Tc f) = Tc (f . k)

-- The state

getState :: Tc TcState
getState = Tc (\_ s -> return (s,s,[]))

setState :: TcState -> Tc ()
setState s = Tc (\_ _ -> return ((),s,[]))

updateState :: (TcState -> TcState) -> Tc ()
updateState f = getState >>= return . f >>= setState

mergeWithState :: TcState -> Tc ()
mergeWithState s = do
  sOld <- getState
  sNew <- liftBase $ mergeStatesBase sOld s
  setState sNew

-- Constraints

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