module Language.Java.Paragon.TypeCheck.TcExp where

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

import Language.Java.Paragon.TypeCheck.Policy
import Language.Java.Paragon.TypeCheck.Actors
import Language.Java.Paragon.TypeCheck.Monad
import Language.Java.Paragon.TypeCheck.Types

import Language.Java.Paragon.TypeCheck.TcEnv     -- The relevant things should really 
import Language.Java.Paragon.TypeCheck.TcState   -- be re-exported by .Monad

--import Control.Monad (liftM)
import Control.Applicative ( (<$>) )
import Data.List ( (\\) )
import Control.Arrow ( first )

import qualified Data.Map as Map
-----------------------------------
--    Types of literal values    --
-----------------------------------

litType :: Literal -> TcType
litType (Int _) = intT
litType (Word _) = longT
litType (Float _) = floatT
litType (Double _) = doubleT
litType (Boolean _) = booleanT
litType (Char _) = charT
litType (String _) = stringT
litType (Null) = nullT
litType l = error $ "litType: " ++ show l

-----------------------------------
--    Checking expressions       --
-----------------------------------

debugTc str = liftBase $ liftIO $ putStrLn $ "DEBUG: Tc: " ++ str

tcExp :: Exp -> Tc (TcType, TcPolicy)

-- Rule PAREN
tcExp (Paren e) = tcExp e

-- Rule LIT
tcExp (Lit l) = return (litType l, bottom)

-- Rule THIS
tcExp This = do
  ty <- getThisType
  --debugTc $ "getThisType: " ++ show ty
  return (ty, bottom)

-- Rule BINOP
tcExp (BinOp e1 op e2) = do
  (ty1, p1) <- tcExp e1
  (ty2, p2) <- tcExp e2
  tyRes <- opType op ty1 ty2
  return (tyRes, p1 `join` p2)

-- Rule VAR/FIELD
tcExp (ExpName n) = do
  --debugTc $ "tcExp ExpName: " ++ show n
  -- lookupVar handles all the complexity of 
  -- finding the correct actor id, if applicable
  (ty,pol) <- lookupVar n
  --debugTc $ show (ty,pol)
  return (ty, pol)
  
-- Rule VARASS/FIELDASS
tcExp (Assign lhs op rhs) = do

  (tyV, pV, mtyO, mpO, mEnt, mN) <- 
    case lhs of
      NameLhs n ->  do
          let (mo,f) = splitName n -- TODO: This is bogus      
          case mo of
            Nothing -> do -- VARASS
                (tyV, pV) <- lookupVar f
                return (tyV, pV, Nothing, Nothing, Just (varE n), Just n)
            Just o -> do -- FIELDASS
                (tyO,pO) <- lookupVar o
                let (Name [i]) = f
                (VTI tyF pF _ _) <- lookupFieldT tyO i
                return (tyF, pF, Just tyO, Just pO, Just (varE n), Just n)
      FieldLhs (PrimaryFieldAccess e fi) -> do
            (tyE, pE) <- tcExp e
            (VTI tyF pF _ _) <- lookupFieldT tyE fi
            let eEnt = case e of
                         This -> Just $ thisFE fi
                         _ -> Nothing
            return (tyF, pF, Just tyE, Just pE, eEnt, Nothing)
      _ -> fail $ "Lhs not supported: " ++ show lhs

  (tyRhs, pRhs) <- tcExp rhs

  -- Check all the necessary constraints:
  check (tyRhs <: tyV) $ "Type mismatch: " ++ show tyRhs ++ " <=> " ++ show tyV
  -- Check: E[branchPC](n) <= pV
  bpc <- maybe getBranchPC_ getBranchPC mEnt

  --debugTc $ "BPC: " ++ show bpc
  constraint_ bpc pV
  -- Check: exnPC(S) <= pV
  epc <- getExnPC
  constraint_ epc pV
  -- Check: pRhs <= pV modulo L
  l <- getCurrentLockState
  constraint l pRhs pV
  -- Check: pO <= pV, if pO exists
  maybeM mpO (\pO -> constraint_ pO pV)

  -- Update actor tracker if applicable
  maybeM (mActorId tyRhs) $ \aid -> do
         maybeM mtyO $ \tyO -> scrambleActors (Just tyO)
         maybeM mN $ \n -> setActorId n aid
  
  return (tyV, pV)

-- Rule CALL
tcExp (MethodInv mi) = tcMethodInv mi

-- Rule NEW
tcExp e@(InstanceCreation tas ct args mBody) = do
  --debugTc $ "tcExp: " ++ show e
  tyT <- liftBase $ evalSrcClsType ct
  -- DEBUG
  --debugTc $ "Type: " ++ show tyT
  --tm <- liftBase getTypeMap
  --let tyTm = lookupTypeOfT (clsTypeToType tyT) tm
  --debugTc $ "TypeMap: " ++ show (Map.keys $ constrs tyTm)
  -- END DEBUG
  (tysArgs, psArgs) <- unzip <$> mapM tcExp args
  (tps,genCti) <- lookupConstr tyT tysArgs
  tArgs <- liftBase $ mapM evalSrcTypeArg tas
 -- tm <- liftBase getTypeMap
  let cti = instantiate (zip tps tArgs) genCti
  let (CTI psPars pW lExp lMods exns) = cti

  -- Check lockstates
  l <- getCurrentLockState
  check (null (lExp \\ l)) $ "Lockstate too weak"
  -- Check argument constraints
  mapM_ (uncurry (constraint l)) (zip psArgs psPars)
  -- Check E[branchPC](*) <= pW
  bpc <- getBranchPC_
  constraint_ bpc pW
  -- Check exnPC(S) <= pW
  epc <- getExnPC
  constraint_ epc pW
  -- Check Exns(X)[write] <= E[exns](X)[write] AND
  -- Check Exns(X)[read]  >= E[exns](X)[read]  
  mapM (uncurry exnConsistent) exns

  -- Fix outgoing state
  let exns' = map (first ExnType) exns
  activateExns exns'     -- ==> S' = Sn[exns{X +-> (Sx, exns(X)[write])}]
  applyLockMods lMods    -- ==> S'' = S'[lockMods ||>>= lMods, 
  scrambleActors Nothing -- ==>          actors scrambled]

  return (clsTypeToType tyT, bottom)

-- Rule COND
tcExp (Cond c e1 e2) = do
  (tyC, pC) <- tcExp c
  check (tyC <: booleanT) $ "Cannot convert type to boolean"
  extendBranchPC pC $ do
    ((ty1, p1), (ty2, p2)) <- 
        (maybeM (mLocks tyC) (\ls -> applyLockMods ([], ls)) >> tcExp e1) ||| tcExp e2
    check (ty1 == ty2) $ "Types of branches don't match"
    
    return (ty1, pC `join` p1 `join` p2)  

tcExp (PolicyExp pl) = do
  pRep <- tcPolicyExp pl
  return (policyPolT pRep, bottom)

--tcExp (PostIncrement e) = error "tcExp: PostIncrement"

tcExp (PreMinus e) = do
  (tyE, pE) <- tcExp e
  check (tyE <: intT) $ "Wrong use of unary minus"
  return (tyE, pE)


tcExp (Cast t e) = do
  (tyE, pE) <- tcExp e
  tyC <- liftBase $ evalSrcType t
  check (tyE <: tyC) $ "Wrong type at cast"
  return (tyC, pE)

tcExp (FieldAccess fa) = tcFieldAccess fa

tcExp e = fail $ "tcExp: Unsupported expression: " ++ show e

--------------------------
-- Field Access

tcFieldAccess :: FieldAccess -> Tc (TcType, TcPolicy)
tcFieldAccess (PrimaryFieldAccess e fi) = do
  (tyE,pE) <- tcExp e
  VTI tyF pFi _ _ <- lookupFieldT tyE fi
  return (tyF, pE `join` pFi)

--------------------------
-- Method invocations

tcMethodInv :: MethodInvocation -> Tc (TcType, TcPolicy)
tcMethodInv mi = do
  --debugTc $ "tcMethodInv: " ++ show mi
  (mti, psArgs, pE) <- 
      case mi of 
        MethodCall n args -> do
            (tysArgs, psArgs) <- unzip <$> mapM tcExp args
            (pPath,_tps,mti) <- lookupMethod n tysArgs
            return (mti, psArgs, pPath)
        PrimaryMethodCall e tas i args -> do
            (tyE, pE) <- tcExp e
            (tysArgs, psArgs) <- unzip <$> mapM tcExp args
            (tps,genMti) <- lookupMethodT tyE i tysArgs
            tArgs <- liftBase $ mapM (evalSrcTypeArg . ActualArg) tas
--            tm <- liftBase getTypeMap
            let mti = instantiate (zip tps tArgs) genMti
            return (mti, psArgs, pE)
        _ -> fail $ "tcMethodInv: Unsupported method call"
  let (MTI tyR pR psPars pW lExp lMods exns) = mti

  -- Check lockstates
  l <- getCurrentLockState
  check (null (lExp \\ l)) $ "Lockstate too weak: " ++ "(" ++ show lExp ++ ", " ++ show l ++ ")"
  -- Check argument constraints
  mapM_ (uncurry (constraint l)) (zip psArgs psPars)
  -- Check E[branchPC](*) <= pW
  bpc <- getBranchPC_
  constraint_ bpc pW
  -- Check exnPC(S) <= pW
  epc <- getExnPC
  constraint_ epc pW
  -- Check Exns(X)[write] <= E[exns](X)[write] AND
  -- Check Exns(X)[read]  >= E[exns](X)[read]  
  mapM (uncurry exnConsistent) exns

  -- Fix outgoing state
  let exns' = map (first ExnType) exns
  activateExns exns'     -- ==> S' = Sn[exns{X +-> (Sx, exns(X)[write])}]
  applyLockMods lMods    -- ==> S'' = S'[lockMods ||>>= lMods, 
  scrambleActors Nothing -- ==>          actors scrambled]

  return (tyR, pE `join` pR)

-----------------------------------
-- Policy expressions
-- Treat them as pure compile-time for now

tcPolicyExp :: PolicyExp -> Tc TcPolicy
tcPolicyExp (PolicyLit cs) = do
  tcCs <- mapM tcClause cs
  return $ TcPolicy tcCs
tcPolicyExp (PolicyOf i) = do
  _ <- lookupVar (Name [i])
  return $ TcRigidVar i
tcPolicyExp (PolicyTypeVar i) = do
  _ <- lookupVar (Name [i])
  return $ TcRigidVar i
 
tcClause :: Clause Actor -> Tc (TcClause TcActor)
tcClause (Clause h ats) = do
  tcH <- tcActor h
  tcAs <- mapM tcAtom ats
  return (TcClause tcH tcAs)

tcAtom :: Atom -> Tc TcAtom
tcAtom (Atom n as) = do
  ar <- lookupLockArity n
  check (length as == ar) $ "Arity mismatch in policy"
  tcAs <- mapM tcActor as
  return (TcAtom n tcAs)

tcActor :: Actor -> Tc TcActor
tcActor (Var i) = return $ TcVar i
tcActor (Actor n) = TcActor <$> tcActorName n

tcActorName :: ActorName -> Tc ActorId
tcActorName (ActorName n) = getActorId n
tcActorName (ActorTypeVar i) = return $ ActorTPVar i

-----------------------------------
--    Types of operators         --
-----------------------------------

opType :: Op -> TcType -> TcType -> Tc TcType
opType Add t1 t2 =
    case (t1, t2) of
      (TcPolicyPolT p1, TcPolicyPolT p2) -> return (TcPolicyPolT (p1 `meet` p2))
      (TcPrimT IntT, TcPrimT IntT) -> return t1
      _ -> fail "Operator applied at wrong types"
opType Mult t1 t2 =
    case (t1, t2) of
      (TcPolicyPolT p1, TcPolicyPolT p2) -> return (TcPolicyPolT (p1 `join` p2))
      (TcPrimT IntT, TcPrimT IntT) -> return t1
      _ -> fail "Operator applied at wrong types"
opType Div t1 t2 =
    case (t1, t2) of
      (TcPrimT IntT, TcPrimT IntT) -> return t1
      _ -> fail "Div applied at wrong types"
opType And t1 t2 =
    case (t1, t2) of
      (TcLockT ls1, TcLockT ls2) -> return (TcLockT (ls1 ++ ls2))
      _ | t1 <: booleanT && t2 <: booleanT -> return booleanT
      _ -> fail "Operator applied at wrong types"
opType op t1 t2 | op `elem` [Equal, NotEq, LThan, GThan, LThanE, GThanE] = do
    case (t1,t2) of
      (TcPrimT pt1, TcPrimT pt2) | pt1 == pt2 -> return ()
      (t1, t2) 
          | isRefType t1 && t2 == nullT || isRefType t2 && t1 == nullT 
              -> return ()
      _ -> fail "Operator NotEq applied at wrong types"
    return booleanT
opType op _ _ = fail $ "opType: Operator not supported: " ++ show op




(<:) :: TcType -> TcType -> Bool
t1 <: t2 
   | t1 == t2 = True
   | t1 == nullT && isRefType t2 = True
   | isLockType t1 && t2 == booleanT = True
   -- Should surely be more cases here
   | isPolicyType t1 && t2 == policyT = True
   -- When we add proper subtyping, it must also be in Tc
   | otherwise = False



splitName :: Name -> (Maybe Name, Name)
splitName (Name is) = 
    let (o,f) = (init is, last is)
        mo = if null o then Nothing else Just (Name o)
    in (mo, Name [f])