module Language.Java.Paragon.TypeCheck.TcStmt where

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

import Language.Java.Paragon.TypeCheck.Policy
import Language.Java.Paragon.TypeCheck.Monad
import Language.Java.Paragon.TypeCheck.Locks
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 Language.Java.Paragon.TypeCheck.TcExp
import Language.Java.Paragon.TypeCheck.Evaluate

--import Control.Monad (liftM)
import Control.Applicative ( (<$>) )

-----------------------------------
--    Typechecking Statements    --
-----------------------------------

tcStmt :: Stmt -> Tc r ()

-- Rule EMPTY
tcStmt Empty = return ()

-- Rule EXPSTMT
tcStmt (ExpStmt e) = ignore $ tcExp e

-- Rule BLOCK
tcStmt (StmtBlock b) = tcBlock b

-- Rule IF
tcStmt (IfThenElse c s1 s2) = do
  (tyC, pC) <- tcExp c
  checkM (liftCont $ tyC <: booleanT) $ "Cannot convert type to boolean"
  extendBranchPC pC $ do
    ignore $ (maybeM (mLocks tyC) (\ls -> applyLockMods ([],ls)) >> tcStmt s1) ||| tcStmt s2

-- Rule IFTHEN
tcStmt (IfThen c s) = tcStmt (IfThenElse c s Empty)

-- Rule WHILE
tcStmt (While c sBody) = do
  s <- getState                  -- Starting state S
  (tyC, pC) <- tcExp c
  checkM (liftCont $ tyC <: booleanT) $ "Cannot convert type to boolean"
  extendBranchPC pC $
    addBranchPC breakE bottom $
    addBranchPC continueE bottom $ 
     do maybeM (mLocks tyC) (\ls -> applyLockMods ([],ls))
        -- First iteration of body
        tcStmt sBody

        -- Approximate state to the possible ones
        -- at the start of the loop
        sCont <- getExnState ExnContinue -- TODO smart constructor
        maybeM sCont $ mergeWithState
        mergeWithState s             -- S* = S <> Ss <> Ss[exns](continue)[state]
        deactivateExn ExnContinue    -- S** = S*[exns /= continue]

        -- Second iteration of condition
        _ <- tcExp c -- This is only to get the (potentially) more difficult constraints
        maybeM (mLocks tyC) (\ls -> applyLockMods ([],ls))
        
        -- Second iteration of body
        tcStmt sBody -- Again for the sake of the constraints, and outgoing state
        
        -- Fix outgoing state
        sBreak <- getExnState ExnBreak   -- Sb = Ss'[exns](break)[state]
        maybeM sBreak mergeWithState     -- S' = Se' <> Sb
        deactivateExn ExnBreak           -- S'' = S'[exns /= break]

-- Rule BREAK
tcStmt (Break Nothing) = do
  bpc <- getBranchPC breakE
  epc <- getExnPC
  s <- getState
  throwExn ExnBreak (bpc `join` epc)

-- Rule CONTINUE
tcStmt (Continue Nothing) = do
  bpc <- getBranchPC continueE
  epc <- getExnPC
  s <- getState
  throwExn ExnContinue (bpc `join` epc)

-- Rule RETURNVOID
tcStmt (Return Nothing) = do
  (tyR, pR) <- getReturn
  check (tyR == voidT) $ "Encountered unexpected empty return statement"
  check (pR == top)   $ "Internal error: tcStmt: " 
                          ++ "void return with non-top return policy should never happen"

  bpc <- getBranchPC returnE
  epc <- getExnPC
  s <- getState
  throwExn ExnReturn (bpc `join` epc)

-- Rule RETURN
tcStmt (Return (Just e)) = do
  (tyR, pR) <- getReturn
  (tyE, pE) <- tcExp e
  checkM (liftCont $ tyE <: tyR) $ "Returned value has wrong type: " ++ prettyPrint tyE
  
  -- Check pE <=[L] pR
  l <- getCurrentLockState
  constraint l pE pR
  -- Check E[branchPC](return) <= pR
  bpc <- getBranchPC returnE
  constraint_ bpc pR
  -- Check exnPC(S) <= pR
  epc <- getExnPC
  constraint_ epc pR

  throwExn ExnReturn pR

-- Rule THROW
tcStmt (Throw eX) = do
  (tyX, pX) <- tcExp eX
  -- TODO: check (tyX <: "Throwable")
  (rX, wX) <- lookupExn tyX
  -- Check E[branchPC](X) <= E[exns](X)[write]
  bpc <- getBranchPC (exnE tyX)
  constraint_ bpc wX
  -- Check exnPC(S) <= E[exns](X)[write]
  epc <- getExnPC
  constraint_ epc wX
  -- Check pX <=[L] E[exns](X)[read]
  l <- getCurrentLockState
  constraint l pX rX

  throwExn (ExnType tyX) wX

-- Rule TRYCATCH
tcStmt (Try block [catch] Nothing) = do
  let Catch (FormalParam ms t isVar (VarId i)) cBlock = catch
  tyP <- liftCont $ evalSrcType t
  -- TODO check tyP <: "Throwable"
  pR <- liftCont $ getReadPolicy ms -- getParamPolicy i ms
  pW <- newMetaPolVar             -- \pi, where \pi is fresh
  addBranchPC (exnE tyP) bottom $ -- E' = E[branchPC{tyP +-> bottom},
    registerExn tyP pR pW $ do    --        exns{tyP +-> (pR, \pi)}]
      tcBlock block
      extendVarEnv i (VSig tyP pR False (isFinal ms)) $ do -- E* = E[vars{x +-> (tyP, pR)}]
        msX <- getExnState (ExnType tyP)
        maybeM msX $ mergeWithState  -- S* = St <> St[exns](X)[state]
        tcBlock cBlock
        deactivateExn (ExnType tyP)
 
-- Rule TRYFINALLY
tcStmt (Try block [] (Just finBlock)) = do
  sS <- getState     -- Must store starting state for later
  tcBlock block
  -- TODO: we actually merge with "too many" states here -- unfortunate, but conservative
  mergeActiveExnStates  -- S* = St <> <>{ Sx | X <- dom(St[exns]), St[exns](X)[state] = Sx }
  sStar <- getState
  useExnState sS        -- S** = S*[exns = S[exns]]
  tcBlock finBlock
  sF <- getState
  useExnState sStar     -- S' = Sf[exns = S*[exns]]
  mergeWithState sF     --       <> Sf

-- Pre-process TRYCATCH/TRYFINALLY
tcStmt (Try block catches mFin) = do
    let catchChain = linearise catches mFin block
    tcStmt catchChain
        where linearise :: [Catch] -> Maybe Block -> Block -> Stmt
              linearise [] justFin  block = Try block [] justFin
              linearise [c] Nothing block = Try block [c] Nothing
              linearise (c:cs) mFin block = linearise cs mFin $ bl (Try block [c] Nothing)
              bl :: Stmt -> Block
              bl = Block . return . BlockStmt

-- Rule OPEN
-- TODO change the list of actor names to a list of expressions (parser, AST, here)
tcStmt (Open (Lock n as)) = do
  arL <- lookupLockArity n
  (_,pL) <- lookupVar n
  check (length as == arL) $ 
            "Lock " ++ prettyPrint n ++ " expects " ++ show arL 
                    ++ " arguments but has been given " ++ show (length as)
  -- Check pI <=[L] pL
  l <- getCurrentLockState
  psAs <- map snd <$> mapM lookupActorName as
  mapM_ (flip (constraint l) pL) psAs
  -- Check E[branchPC](L) <= pL
  bpc <- getBranchPC (lockE n)
  constraint_ bpc pL
  -- Check exnPC(S) <= pL
  epc <- getExnPC
  constraint_ epc pL

  aids <- mapM tcActorName as
  openLock (TcLock n aids)

-- Rule CLOSE
tcStmt (Close (Lock n as)) = do
  arL <- lookupLockArity n
  (_,pL) <- lookupVar n
--  LTI arL pL <- lookupLock n
  check (length as == arL) $ 
            "Lock " ++ prettyPrint n ++ " expects " ++ show arL 
                    ++ " arguments but has been given " ++ show (length as)
  -- Check pI <=[L] pL
  l <- getCurrentLockState
  psAs <- map snd <$> mapM lookupActorName as
  mapM_ (flip (constraint l) pL) psAs
  -- Check E[branchPC](L) <= pL
  bpc <- getBranchPC (lockE n)
  constraint_ bpc pL
  -- Check exnPC(S) <= pL
  epc <- getExnPC
  constraint_ epc pL

  aids <- mapM tcActorName as
  closeLock (TcLock n aids)

-- Rule OPENIN
tcStmt (OpenBlock (Lock n as) block) = do
  arL <- lookupLockArity n
  (_,pL) <- lookupVar n
  --debugTc $ "pL: " ++ prettyPrint n ++ ": " ++ show pL
--  LTI arL pL <- lookupLock n
  check (length as == arL) $ 
            "Lock " ++ prettyPrint n ++ " expects " ++ show arL 
                    ++ " arguments but has been given " ++ show (length as)
  -- Check pI <=[L] pL
  l <- getCurrentLockState
  psAs <- map snd <$> mapM lookupActorName as
  --debugTc $ "psAs: " ++ show psAs
  mapM_ (flip (constraint l) pL) psAs

  aids <- mapM tcActorName as
  extendLockEnv [TcLock n aids] $
    tcBlock block

tcStmt s = fail $ "Unsupported statement: " ++ prettyPrint s

----------------------------------------
-- Typechecking Blocks and BlockStmts --
----------------------------------------

tcBlock :: Block -> Tc r ()
tcBlock (Block bss) = tcBlockStmts bss

tcBlockStmts :: [BlockStmt] -> Tc r ()
-- Rule EMPTYBLOCK
tcBlockStmts [] = return ()

-- Rule BLOCKSTMT
tcBlockStmts (BlockStmt stmt : bss) = tcStmt stmt >> tcBlockStmts bss

-- Rule LOCALVARINIT
tcBlockStmts (LocalVars ms t [VarDecl (VarId i) (Just (InitExp e))] : bss) = do
  (tyE, pE) <- tcExp e
  tyV <- liftCont $ evalSrcType t
  checkM (liftCont $ tyE <: tyV) $ 
             "Type mismatch: " ++ prettyPrint tyE ++ " <=> " ++ prettyPrint tyV
  pV <- localVarPol ms
  tyV' <- case mActorId tyE of
            Nothing -> return tyV
            Just aid -> do
              newActorIdWith i aid
              return $ actorIdT aid
  extendVarEnv i (VSig tyV' pV False (isFinal ms)) $ do
    tcBlockStmts bss

-- Rule LOCALVARDECL
tcBlockStmts (LocalVars ms t [VarDecl (VarId i) Nothing] : bss) = do
  tyV <- liftCont $ evalSrcType t
  pV <- localVarPol ms
  tyV' <- if tyV == actorT
           then actorIdT <$> newActorId i
           else return tyV
  extendVarEnv i (VSig tyV' pV False (isFinal ms)) $ do
    tcBlockStmts bss
    
tcBlockStmts (b:bss) = fail $ "Unsupported block statement: " ++ prettyPrint b

localVarPol :: [Modifier] -> Tc r TcPolicy
localVarPol ms = 
    case [ p | Reads p <- ms ] of
      [] -> return bottom --newMetaPolVar
      [p] -> liftCont $ evalPolicy p