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 = liftCont $ liftIO $ putStrLn $ "DEBUG: Tc: " ++ str

tcExp :: Exp -> Tc r (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
                (VSig 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
            (VSig 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: " ++ prettyPrint lhs

  (tyRhs, pRhs) <- tcExp rhs

  -- Check all the necessary constraints:
  checkM (liftCont $ tyRhs <: tyV) $ 
             "Type mismatch: " ++ prettyPrint tyRhs ++ " <=> " ++ prettyPrint 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 <- liftCont $ evalSrcClsType ct
  -- DEBUG
  --debugTc $ "Type: " ++ show tyT
  --tm <- liftCont 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 <- liftCont $ mapM evalSrcTypeArg tas
 -- tm <- liftCont getTypeMap
  let cti = instantiate (zip tps tArgs) genCti
  let (CSig 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
  checkM (liftCont $ 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) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $ 
            "Post-increment operator used at non-numeric type " ++ prettyPrint tyE
  return (tyE, pE)
tcExp (PostDecrement e) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $ 
            "Post-decrement operator used at non-numeric type " ++ prettyPrint tyE
  return (tyE, pE)
tcExp (PreIncrement e) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $ 
            "Pre-increment operator used at non-numeric type " ++ prettyPrint tyE
  return (tyE, pE)
tcExp (PreDecrement e) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $ 
            "Pre-decrement operator used at non-numeric type " ++ prettyPrint tyE
  return (tyE, pE)

-- Unary promotion prefix operator

-- Rule PREP
tcExp (PrePlus e) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $
            "Pre-plus operator used at non-numeric type " ++ prettyPrint tyE
  
  return (unaryNumPromote_ tyE, pE)

tcExp (PreMinus e) = do
  (tyE, pE) <- tcExp e
  check (isNumConvertible tyE) $
            "Pre-minus operator used at non-numeric type " ++ prettyPrint tyE
  return (unaryNumPromote_ tyE, pE)

tcExp (PreBitCompl e) = do
  (tyE, pE) <- tcExp e
  check (isIntConvertible tyE) $
        "Pre-complement bit operator used at non-integral type " ++ prettyPrint tyE
  return (unaryNumPromote_ tyE, pE)

tcExp (PreNot e) = do
  (tyE, pE) <- tcExp e
  check (isBoolConvertible tyE) $
        "Pre-complement boolean operator used at non-boolean type " ++ prettyPrint tyE
  return (booleanT, pE)


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

tcExp (FieldAccess fa) = tcFieldAccess fa

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

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

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

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

tcMethodInv :: MethodInvocation -> Tc r (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 <- liftCont $ mapM (evalSrcTypeArg . ActualArg) tas
--            tm <- liftCont getTypeMap
            let mti = instantiate (zip tps tArgs) genMti
            return (mti, psArgs, pE)
        _ -> fail $ "tcMethodInv: Unsupported method call"
  let (MSig 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 r 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 r (TcClause TcActor)
tcClause (Clause h ats) = do
  tcH <- tcActor h
  tcAs <- mapM tcAtom ats
  return (TcClause tcH tcAs)

tcAtom :: Atom -> Tc r 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 r TcActor
tcActor (Var i) = return $ TcVar i
tcActor (Actor n) = TcActor <$> tcActorName n

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

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

opType :: Op -> TcType -> TcType -> Tc r TcType
-- First the special cases: policy operators, and String conversion
opType Mult (TcPolicyPolT p1) (TcPolicyPolT p2) = return (TcPolicyPolT (p1 `join` p2))
opType Add  (TcPolicyPolT p1) (TcPolicyPolT p2) = return (TcPolicyPolT (p1 `meet` p2))
opType Add t1 t2 | t1 == stringT || t2 == stringT = return stringT

opType op t1 t2
    -- Numeric operators
    | op `elem` [Mult, Div, Rem, Add, Sub] = do
        check (isNumConvertible t1) $
              "Numeric operator " ++ prettyPrint op ++ 
                " used with non-numeric operand of type " ++ prettyPrint t1
        check (isNumConvertible t2) $
              "Numeric operator " ++ prettyPrint op ++ 
                " used with non-numeric operand of type " ++ prettyPrint t2
        return $ binaryNumPromote_ t1 t2

    -- Shift operators
    | op `elem` [LShift, RShift, RRShift] = do
        check (isIntConvertible t1) $
              "Shift operator " ++ prettyPrint op ++
                " used with non-integral operand of type " ++ prettyPrint t1
        check (isIntConvertible t2) $
              "Shift operator " ++ prettyPrint op ++
                " used with non-integral operand of type " ++ prettyPrint t2
        return $ unaryNumPromote_ t1

    -- Relational operators
    | op `elem` [LThan, GThan, LThanE, GThanE] = do
        check (isNumConvertible t1) $
              "Numerical comparison operator " ++ prettyPrint op ++ 
                " used with non-numeric operand of type " ++ prettyPrint t1
        check (isNumConvertible t2) $
              "Numerical comparison operator " ++ prettyPrint op ++ 
                " used with non-numeric operand of type " ++ prettyPrint t2
        return booleanT

    | op `elem` [Equal, NotEq] = do
        case binaryNumPromote t1 t2 of
          Just _ -> return ()
          _ | isBoolConvertible t1 && isBoolConvertible t2 -> return ()
            | isRefType t1 && isRefType t2 -> return ()
          _ -> fail $ "Equality operator " ++ prettyPrint op ++
                       " used with incomparable operands of types " ++
                       prettyPrint t1 ++ " and " ++ prettyPrint t2
        return booleanT

    | op `elem` [And, Or, Xor] =
        if isBoolConvertible t1 
         then do
           check (isBoolConvertible t2) $ 
                     "Logical operator " ++ prettyPrint op ++
                      " used with non-boolean operand of type " ++ prettyPrint t2
           return booleanT
         else if isIntConvertible t1
               then do
                 check (isIntConvertible t2) $
                       "Bitwise operator " ++ prettyPrint op ++
                        " used with non-integral operand of type " ++ prettyPrint t2
                 return $ binaryNumPromote_ t1 t2
               else fail $ "Bitwise/logical operator " ++ prettyPrint op ++
                             " used with non-integral and non-boolean operand of type " ++
                             prettyPrint t1
    | op `elem` [CAnd, COr] = do
        check (isBoolConvertible t1) $
                  "Logical operator " ++ prettyPrint op ++
                    " used with non-boolean operand of type " ++ prettyPrint t1
        check (isBoolConvertible t2) $
                  "Logical operator " ++ prettyPrint op ++
                    " used with non-boolean operand of type " ++ prettyPrint t2
        return booleanT


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])