{-#LANGUAGE ScopedTypeVariables, ExistentialQuantification, MultiParamTypeClasses, FlexibleContexts, GADTs, FunctionalDependencies, FlexibleInstances, TemplateHaskell #-}

module Join.Monads (
    VAR, -- TODO: Make this unnecessary
    Sync, -- same
    Process,
    withRules,
    spawn,
    spawnN,
    Join.Monads.call,
    asyncChan,
    AsyncChan,
    syncChan,
    SyncChan,
    joinList,
    Rules,
    var,
    def,
    (&),
    (.\.),
    onlyIf,
    reply,
    to,
    apply,
    apply2,
    value
) where

import Join.Join
import MultiSetRewrite.Base
import qualified MultiSetRewrite.StoreRepresentation as S
import qualified MultiSetRewrite.RuleCompiler as C
import qualified MultiSetRewrite.RuleSyntax as R hiding (when)

import Control.Concurrent
import Control.Monad 
import Control.Monad.Trans
import Language.Haskell.TH
import Data.Char
import Data.Bits
import Data.Typeable


infixr 4 & 
-- Not associative to make patterns clearer: all propagated messages to the left, all normal messages to the right.
infix 3 .\.
-- Same thing: just use boolean operators to combine guards...
infix 2 `onlyIf`

{-
 - Channels
 -}

data AsyncChan = forall a . (CollectArgs a) => 
    AC { name :: String, arg :: a }


-- r is the return type (phantom type)
data SyncChan r where  
    SC :: forall a . (CollectArgs a, Typeable r, Assign (VAR (Sync r)) r) => 
                String -> a -> VAR (Sync r) -> SyncChan r

aVar = varT (mkName "a")
varConstr = conT (mkName "VAR")
syncConstr = conT (mkName "Sync")

asyncChan :: String -> Q Type -> Q [Dec]
asyncChan name argType = 
        chanSplice name argType [|AC|]
            (appT (appT arrowT aVar) (conT (mkName "AsyncChan")))

syncChan :: String -> Q Type -> Q [Dec]
syncChan name arrowType = do
    -- allows for nicer user syntax: just the arrow type instead of 2 separate argument and return types
    (AppT (AppT ArrowT _argType) _retType) <- arrowType
    let argType = return _argType
    let retType = return _retType
    let retSig = appT (conT ''SyncChan) retType
    chanSplice name argType [|SC|]
        (appT   (appT arrowT aVar)
                (appT (appT arrowT
                            (appT varConstr
                                  (appT syncConstr
                                        retType)))
                      retSig))
        
chanSplice name argType constr sig = do
    let className = mkName (capitalize name)
    let funcName = mkName name
    class_ <- classD (cxt []) className [mkName "a"] [] 
        [sigD funcName sig]
    let instanceDecl typeT =
            instanceD (cxt []) (appT (conT className) typeT)
                [funD funcName [clause [] 
                    (normalB (appE constr [|name|]) ) [] ]]
    t <- argType
    instances <- mapM instanceDecl (map return (typePermutations t))
    return (class_:instances)
    where capitalize (x:xs) = (toUpper x):xs
-- What the above generates, for a sync channel get :: Int -> String
-- It will become readable when splices in [d| |] are added to TH...
-- [d|
--    class Get a where
--        get :: a -> VAR (Sync String) -> SyncChan String
--    instance Get Int where
--        get = SC "get"
--    instance Get (VAR Int) where
--        get = SC "get"
-- |]

-- Tuples require the VAR to be on each tuple constituent, forcing
-- us to walk down the tree to figure out if we have a tuple or not.
-- Tuples are as follows in the syntax tree:
-- AppT (AppT (AppT (TupleT 3) (..Int)) (..Int)) (..Int)
-- Update: in fact we need instances for each permutation of the VARs, eg
-- instance (Int,Int)... instance (VAR Int, Int)... instance (Int, VAR Int)...
-- This is starting to feel like a hack
typePermutations :: Type -> [Type]
typePermutations t = 
    let (lst,wastuple) = tupleTypes t in 
        if wastuple then map makeTuple (permutations lst)
                    else [t, makeVar t]
    where
        permutations lst = map (zipWith (flip ($)) lst) $ sequence (lst >> return [makeVar, id])
        makeVar t = AppT (ConT (mkName "VAR")) t
        tupleTypes (TupleT _) = ([],True)
        tupleTypes t@(AppT t1 t2) = let (lst,wastuple) = tupleTypes t1 in
            if wastuple then (lst ++ [t2],True)
                        else ([t],False)
        tupleTypes t = ([t],False)
        makeTuple lst = makeTuple' (length lst) (reverse lst)
        makeTuple' n [] = TupleT n
        makeTuple' n (x:xs) = AppT (makeTuple' n xs) x
        

{- 
 - The Process monad: what to do when a pattern matches, or in the initial system
 - (name comes from the Chemical Abstract Machine: processes floating in a solution,
 - travelling to reaction sites (patterns)).
 -}

data Environment = Env { jStoreS :: Join }

newtype Process a = S { runProcess :: Environment -> IO (a, Environment) }

instance Monad Process where
    return x = S $ \e -> return (x, e)
    s >>= f = S $ \e -> do (val, e1) <- runProcess s e
                           runProcess (f val) e1
    fail str = S $ \_ -> fail str

instance MonadIO Process where
    liftIO io = S $ \e -> do val <- io
                             return (val, e)

withRules :: Rules () -> Process a -> IO a
withRules rules solution = do
    joinStore <- createJoinStore (execRules rules)
    execProcessInternal joinStore solution 

execProcessInternal :: Join -> Process a -> IO a
execProcessInternal joinStore solution = do
    (value, env) <- runProcess solution (Env joinStore)
    return value

runS sol defs = runProcess sol (Env (jStoreDef defs))

execS sol defs = do
    (action,_env) <- runS sol defs
    return action

class Spawn a where
    spawn :: a -> Process ()

instance Spawn AsyncChan where
    spawn (AC name arg) = S $ \e -> do 
        Join.Join.call (jStoreS e) name arg
        return ((), e)
instance Spawn [AsyncChan] where
    spawn [] = return ()
    -- ensures that processes are spawned in text order, as 
    -- the Process monad ends up reversing stuff (same issue as def)
    spawn (x:xs) = spawn xs >> spawn x

spawnN n p = mapM_ (\_ -> spawn p) [1..n]

call :: (VAR (Sync r) -> SyncChan r) -> Process r
call chan = call' (chan undefined)
call' :: SyncChan r -> Process r
call' (SC name arg _) = S $ \e -> do 
    value <- callSync (jStoreS e) name arg
    return (value, e)
    
to = ()
reply :: (Assign (VAR (Sync r)) v) => v -> () -> VAR (Sync r) -> Process ()
reply value _to resvar = S $ \e -> 
    resvar .=. value >> return ((),e)

value :: VAR a -> Process a
value x = liftIO (readVar x)

-- TODO: class Apply ?
apply f x = do v <- value x
               return (f v)
apply2 f x y = do vx <- value x
                  vy <- value y
                  return (f vx vy)
{-
 - The Rules monad: defining join patterns
 -}

data Definitions = Def { jStoreDef :: Join, 
                         defList :: [([R.MatchTask Method], C.Code_RHS ())]
                       }
newtype Rules a = R { runRules :: Definitions -> IO (a, Definitions) }

instance Monad Rules where
    r >>= f = R $ \defs -> do (val, defs') <- runRules r defs
                              runRules (f val) defs'
    return x = R $ \d -> return (x, d)
    fail str = R $ \d -> fail str

execRules rules jStore = do
    (val, defs) <- runRules rules (Def jStore [])
    return (defList defs)

var :: (CollectArgs (VAR z), Assign (VAR (Sync z)) (VAR z)) => Rules (VAR z)
var = R $ \d -> do v <- newVar
                   return (v, d)

def :: Pattern a h => a -> Process () -> Rules ()
def pattern solution = R $ \defs -> do
    def <- newdef pattern solution defs
    -- append, not prepend, to preserve the ordering of rules in
    -- user code, as rules listed first by user have precedence
    -- (same issue as spawn [])
    return ((), defs {defList = (defList defs)++[def]})
    where
    newdef pattern solution defs = do
        (convPat, env) <- runS (convertPattern pattern) defs
        return (convPat R..->. execS solution defs)

{-
 -  Patterns
 -}

class (R.ConvertHead h Method) => Pattern t h | t -> h where
    convertPattern :: t -> Process h
    
instance Pattern AsyncChan [Method] where
    convertPattern (AC name arg) = return [method name arg]

instance Pattern (SyncChan r) [Method] where
    convertPattern (SC name arg resvar) = S $ \env ->
        return ([method name (resvar, arg)], env)

data Guarded = forall p . (Pattern p [Method]) => Guarded { pattern :: p, guard :: Process Bool }
pattern `onlyIf` guard = Guarded pattern guard

instance Pattern Guarded (R.WITHGUARD [Method]) where
    convertPattern (Guarded p g) = S $ \e@(Env join) -> do
        (convPat,e') <- runProcess (convertPattern p) e
        return (convPat `R.WITHGUARD` (execProcessInternal join g), e')

class (Pattern a [Method]) => Channel a
instance Channel (SyncChan r)
instance Channel AsyncChan
data AnyChan = forall a . (Channel a) => AnyChan a

instance Pattern [AnyChan] [Method] where
    convertPattern [] = return []
    convertPattern ((AnyChan c):xs) = do
        l1 <- convertPattern c
        l2 <- convertPattern xs
        return (l1 ++ l2)

class (Pattern g [Method]) => ChannelGroup g where
    (&) :: forall c . (Channel c) => c -> g -> [AnyChan]
instance ChannelGroup (SyncChan r) where
    c & g = (AnyChan c):[AnyChan g]
instance ChannelGroup AsyncChan where
    c & g = (AnyChan c):[AnyChan g]
instance ChannelGroup [AnyChan] where
    c & g = (AnyChan c):g

joinList :: Channel a => [a] -> [AnyChan]
joinList xs = (map AnyChan xs)


(.\.) :: (ChannelGroup l, ChannelGroup r) => l -> r -> (l,r)
(.\.) x y = (x,y)
instance (ChannelGroup l, ChannelGroup r) => Pattern (l,r) ([Method],[Method]) where
    convertPattern (l,r) = do
        patL <- convertPattern l
        patR <- convertPattern r
        return (patL, patR)

{-
    Various utilities and wrappers around the original library
-}

callSync :: (CollectArgs (Sync r), CollectArgs a) => 
    Join -> String -> a -> IO r
callSync join name arg = do
    x <- newSync
    Join.Join.call join name (x, arg)
    val <- waitSync x
    return val

-- Join store initialization
createJoinStore :: (Join -> IO [([R.MatchTask Method], C.Code_RHS ())]) -> IO Join
createJoinStore rules = do
    jStore <- newJoinStore
    return Join {store = jStore, rules = wrapRules rules}
    where
    wrapRules rules join activeMethod = do
        definitions <- rules join
        let patterns = translateJoinDefinitions definitions
        res <- runJoinOnGoal (store join) activeMethod patterns
        maybe (return ()) id res