{-# LANGUAGE CPP, DeriveDataTypeable, PatternGuards, RelaxedPolyRec #-}
module Language.Java.Paragon.TypeCheck.TcEnv where

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

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

import qualified Data.Map as Map
import Control.Monad ( msum )

import Data.Generics.Uniplate.Data

#ifdef BASE4
import Data.Data
#else
import Data.Generics (Data(..),Typeable(..))
#endif

import Debug.Trace

type Map = Map.Map

data TypeMap = TypeMap {
      -- signatures
      fields      :: Map (Ident ()) VarFieldSig,
      methods     :: Map (Ident (), [TcType]) ([TypeParam ()], MethodSig),
      constrs     :: Map [TcType] ([TypeParam ()], ConstrSig),
      locks       :: Map (Ident ()) LockSig,
      -- known policy-level entities
      policies    :: Map (Ident ()) TcPolicy,
      actors      :: Map (Ident ()) ActorId,
      -- typemethod eval info
      typemethods :: Map (Ident ()) ([Ident ()], Block ()),
      -- types and packages
      types       :: Map (Ident ()) ([TypeParam ()], TypeSig),
      packages    :: Map (Ident ()) TypeMap
    }
  deriving (Show, Data, Typeable)

emptyTM :: TypeMap
emptyTM = TypeMap {
                fields      = Map.empty,
                methods     = Map.empty,
                constrs     = Map.empty,
                locks       = Map.empty,
                policies    = Map.empty,
                actors      = Map.empty,
                typemethods = Map.empty,
                types       = Map.empty,
                packages    = Map.empty
              }

hardCodedArrayTM :: TcType -> TcPolicy -> TypeMap
hardCodedArrayTM _ _ = 
    emptyTM {
         fields = Map.fromList [(Ident () "length", VSig intT thisP False True)]
       , methods = Map.fromList [] -- TODO
       }

clearToPkgs :: TypeMap -> TypeMap
clearToPkgs tm = emptyTM { packages = packages tm }
                   
clearToTypes :: TypeMap -> TypeMap
clearToTypes tm = emptyTM { packages = packages tm, types = types tm }

data TcEnv = TcEnv  {
      -- lockProps :: ??,
      vars      :: Map (Ident ()) VarFieldSig,
      lockstate :: [TcLock],
      returnI   :: Maybe (TcType, TcPolicy),
      exnsE     :: Map TcType (TcPolicy, TcPolicy),
      branchPCE :: (Map Entity [(TcPolicy, String)], [(TcPolicy, String)])
    }
  deriving (Show, Data, Typeable)

-- Env to use when typechecking expressions not inside method
-- bodies, e.g. in field initializers and policy modifiers
simpleEnv :: TcPolicy -> String -> TcEnv
simpleEnv brPol str = 
    TcEnv {
      vars = Map.empty,
      lockstate = [],
      returnI = Nothing, -- error "No returns in simple env",
      exnsE = Map.empty,
      branchPCE = (Map.empty, [(brPol,str)])
    }

data Entity = VarEntity (Name ())
            | ThisFieldEntity (Ident ())
            | ExnEntity TcType
            | LockEntity (Name ())
            | BreakE | ContinueE | ReturnE
  deriving (Show, Eq, Ord, Data, Typeable)

varE, lockE :: Name () -> Entity
varE = VarEntity
lockE = LockEntity

exnE :: TcType -> Entity
exnE = ExnEntity

thisFE :: Ident () -> Entity
thisFE = ThisFieldEntity 

breakE, continueE, returnE :: Entity
breakE = BreakE
continueE = ContinueE
returnE = ReturnE

data VarFieldSig = VSig { 
      varType   :: TcType,
      varPol    :: TcPolicy,
      varStatic :: Bool,
      varFinal  :: Bool
    }
  deriving (Show, Data, Typeable)

data MethodSig = MSig {
      mRetType :: TcType,
      mRetPol  :: TcPolicy,
      mPars    :: [TcPolicy],
      mWrites  :: TcPolicy,
      mExpects :: [TcLock],
      mLMods   :: ([TcLock],[TcLock]),
      mExns    :: [(TcType, ExnSig)]
    }
  deriving (Show, Data, Typeable)

data ExnSig = ExnSig { 
      exnReads :: TcPolicy, 
      exnWrites :: TcPolicy, 
      exnMods :: ([TcLock],[TcLock]) 
    }
  deriving (Show, Data, Typeable)

data ConstrSig = CSig {
      cPars    :: [TcPolicy],
      cWrites  :: TcPolicy,
      cExpects :: [TcLock],
      cLMods   :: ([TcLock],[TcLock]),
      cExns    :: [(TcType, ExnSig)]
    }
  deriving (Show, Data, Typeable)

data LockSig = LSig {
      lPol     :: TcPolicy,
      lArity   :: Int
    }
  deriving (Show, Data, Typeable)

data TypeSig = TSig {
      tIsClass    :: Bool,
      tIsFinal    :: Bool,
      tSupers     :: [TcType],
      tImpls      :: [TcType],
      tMembers    :: TypeMap
    }
  deriving (Show, Data, Typeable)

--------------------------------------
--   Type argument instantiation    --
--------------------------------------

instantiate :: Data a => [(TypeParam (),TcTypeArg)] -> a -> a
instantiate pas = transformBi instT
                     . transformBi instA
                     . transformBi instP
                     . transformBi instLs
    where instT :: TcRefType -> TcRefType
          instT tv@(TcTypeVar i) =
              case lookup i typs of
                Just rt -> rt
                Nothing -> tv
          instT rt = rt

          instA :: ActorId -> ActorId
          instA av@(ActorTPVar i) =
              case lookup i as of
                Just a -> a
                Nothing -> av
          instA a = a

          instP :: TcPolicy -> TcPolicy
          instP pv@(TcRigidVar i) =
              case lookup i ps of
                Just p -> p
                Nothing -> pv
          instP p = p

          instLs :: [TcLock] -> [TcLock]
          instLs = concatMap instL

          instL :: TcLock -> [TcLock]
          instL lv@(TcLockVar i) =
              case lookup i locks of
                Just le -> le
                Nothing -> [lv]
          instL l = [l]

          typs  = [ (i, rt) | (TypeParam   _ i _, TcActualType  rt) <- pas ] 
          as    = [ (i, n ) | (ActorParam  _ i  , TcActualActor  n) <- pas ]
          ps    = [ (i, p ) | (PolicyParam _ i  , TcActualPolicy p) <- pas ]
          locks = [ (i, ls) | (LockStateParam _ i, TcActualLockState ls) <- pas ]



--------------------------------------
--    Working with the branchPC     --
--------------------------------------

branchPC :: Maybe Entity -> TcEnv -> [(TcPolicy, String)]
branchPC men (TcEnv { branchPCE = (bm, def) }) = 
    flip (maybe def) men $ \en ->
        maybe def id (Map.lookup en bm)

joinBranchPC :: TcPolicy -> String -> TcEnv -> TcEnv
joinBranchPC p str env = let (bm, def) = branchPCE env
                         in env { branchPCE = (Map.map ((p, str):) bm, (p,str):def) }

--------------------------------------
--    Working with the lookups      --
--------------------------------------

lookupNamed :: (TypeMap -> Map (Ident ()) a) -> Name () -> TypeMap -> Maybe a
lookupNamed recf nam@(Name _ is) tm =
   case lookupTypeOfN (Name () $ init is) tm of
     Right actualTm -> Map.lookup (last is) (recf actualTm)
     Left err -> Nothing

lookupNamedMethod :: Name () -> [TcType] -> TypeMap -> Maybe ([TypeParam ()],MethodSig)
lookupNamedMethod (Name _ is) ts tm = 
    case lookupTypeOfN (Name () $ init is) tm of
      Right actualTm -> Map.lookup (last is, ts) (methods actualTm)
      Left err -> Nothing

pkgsAndTypes :: TypeMap -> Map (Ident ()) TypeMap
pkgsAndTypes tm = Map.union (packages tm)
                    -- disregard type parameters
                    (Map.map (tMembers . snd) $ types tm)

lookupTypeOfN :: Name () -> TypeMap -> Either (Maybe String) TypeMap
lookupTypeOfN nam@(Name _ is) tm = aux is (Right tm) tm
    where aux :: [Ident ()] 
              -> Either (Maybe String) TypeMap 
              -> TypeMap 
              -> Either (Maybe String) TypeMap
          aux _ err@(Left _) _ = err
          aux [] etm _ = etm
          aux (i:is) (Right tm) baseTm =
            let mNewTm = 
                 case Map.lookup i (fields tm) of
                   Just (VSig typ _ _ _) -> lookupTypeOfT typ baseTm
                   Nothing -> case Map.lookup i (pkgsAndTypes tm) of
                                Just aTm -> Right aTm
                                Nothing -> Left Nothing -- $ "Unknown name: " ++ prettyPrint nam
            in aux is mNewTm baseTm

lookupTypeOfT :: TcType -> TypeMap -> Either (Maybe String) TypeMap
lookupTypeOfT (TcRefT refT) = lookupTypeOfT' refT
lookupTypeOfT t = const $ Left Nothing

lookupTypeOfT' :: TcRefType -> TypeMap -> Either (Maybe String) TypeMap
lookupTypeOfT' (TcClsRefT (TcClassT iargs)) tm = aux iargs (Right tm)
    where aux :: [(Ident (), [TcTypeArg])] 
              -> Either (Maybe String) TypeMap 
              -> Either (Maybe String) TypeMap
          aux _ err@(Left _) = err
          aux [] etm = etm
          aux ((i, args):iargs) (Right tm) = -- traceShow args $
              let eNewTm =
                   case Map.lookup i (types tm) of
                     Just (pars, tsig) -> 
                         if length pars == length args
                         then Right $ instantiate (zip pars args) (tMembers tsig)
                         else Left $ Just $ 
                                  "Wrong number of type arguments in class type.\n" ++
                                  "Type " ++ prettyPrint i ++ " expects " ++ show (length pars) ++
                                  " arguments but has been given " ++ show (length args)
                     Nothing -> 
                         case Map.lookup i (packages tm) of
                           Just tm -> Right tm
                           Nothing -> Left Nothing -- $ "Unknown type : " ++ prettyPrint i
              in aux iargs eNewTm

lookupTypeOfT' (TcArrayT ty pol) _ = Right $ hardCodedArrayTM ty pol
-- TODO: Insert array support