content: {-
  Copyright   : (c) Jean-Christophe Mincke, 2016-2017
  
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
  -}
  
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module Control.Distributed.Blast.Distributed.Master
  (
    runLocal
  )
  where
  
  
  import            Control.Concurrent.Async
  import            Control.Monad
  import            Control.Monad.IO.Class
  import            Control.Monad.Logger
  import            Control.Monad.Trans.State
  
  import qualified  Data.ByteString as BS
  import qualified  Data.List as L
  import qualified  Data.Map as M
  import            Data.Maybe (catMaybes)
  import qualified  Data.Serialize as S
  import qualified  Data.Set as S
  import qualified  Data.Vault.Strict as V
  import qualified  Data.Vector as Vc
  
  import            Control.Distributed.Blast.Types
  import            Control.Distributed.Blast.Distributed.Types
  import            Control.Distributed.Blast.Common.Analyser
  import            Control.Distributed.Blast.Master.Analyser
  
  toData :: Maybe a -> Data a
  toData (Just a) = Data a
  toData Nothing = NoData
  
  getLocalVaultKey :: MExp 'Local a -> LocalKey a
  getLocalVaultKey (MLConst _ k _) = k
  getLocalVaultKey (MCollect _ k _ _) = k
  getLocalVaultKey (MLApply _ k _ _) = k
  
  getRemote :: (Monad m) => StateT (s, V.Vault, InfoMap) m s
  getRemote = do
    (s, _, _) <- get
    return s
  
  getVault :: (Monad m) => StateT (s, V.Vault, InfoMap) m (V.Vault)
  getVault = do
    (_, vault, _) <- get
    return vault
  
  setVault :: (Monad m) => V.Vault -> StateT (s, V.Vault, InfoMap) m ()
  setVault vault= do
    (s, _, m) <- get
    put (s, vault, m)
  
  dereference :: (Monad m) => Int -> Int -> StateT (s, V.Vault, InfoMap) m Bool
  dereference parent child = do
    (s, vault, m) <- get
    let (GenericInfo crefs i) = m M.! child
    if S.member parent crefs
      then do
        let crefs' = S.delete parent crefs
        let m' = M.insert child (GenericInfo crefs' i) m
        put (s, vault, m')
        return $ S.null crefs'
      else error ("Remove parent reference twice for parent "++show parent ++" and child "++show child)
  
  handleRpcError :: Either String t -> t
  handleRpcError (Right a) = a
  handleRpcError (Left err) = error ("Error in RPC: "++err)
  
  handleRpcErrorM :: (Monad m) => Either String t -> m t
  handleRpcErrorM a = return $ handleRpcError a
  
  
  handleSlaveErrorM :: (Monad m) => SlaveResponse -> m SlaveResponse
  handleSlaveErrorM (LsRespError err) = error ("Error in Slave: "++err)
  handleSlaveErrorM r = return r
  
  
  
  runAndFetchRemoteOneSlaveStatefull ::(CommandClass s x, MonadIO m, S.Serialize a) => Int -> MExp 'Remote a -> StateT (s x , V.Vault, InfoMap) m (Maybe a)
  runAndFetchRemoteOneSlaveStatefull slaveId oe@(MRApply n (ExpClosure ce _) e) = do
    runRemoteOneSlaveStatefull slaveId oe
    fetchRes <- fetchOneSlaveResults slaveId oe
    case fetchRes of
      Just (Data a) -> return $ Just a
      Just (NoData) -> return $ Nothing
      Nothing ->
        -- We a fectch miss. Redo the remote computation
        -- Risk of infinite recursion: not really because a fetch miss
        -- is due to (temporary) slave failure.
        -- If the slave cannot never be reached gain, the backend will raise an specific error
        runAndFetchRemoteOneSlaveStatefull slaveId oe
  
  
  runRemoteOneSlaveStatefull ::(CommandClass s x, MonadIO m) => Int -> MExp 'Remote a -> StateT (s x , V.Vault, InfoMap) m ()
  runRemoteOneSlaveStatefull slaveId oe@(MRApply n (ExpClosure ce _) e) = do
    s <- getRemote
    let req = LsReqExecute n
    r <- liftIO $ send s slaveId req
    case handleRpcError r of
      LsRespExecute (RcRespCacheMiss CachedFreeVar) -> do
        vault <- getVault
        let keyce = getLocalVaultKey ce
        let cem = V.lookup keyce vault
        case cem of
          Nothing -> error "local value not cached while executing remote on one slave"
          Just (_, Nothing) -> error "local value not cached (BS) while executing remote on one slave"
          Just (_, Just p) -> do
            let ceId = getLocalIndex ce
            let csBs = toData $ getPart slaveId p
            let req' = LsReqCache ceId csBs
            r' <- liftIO $ send s slaveId req'
            sr <- handleRpcErrorM r'
            _ <- handleSlaveErrorM sr
            runRemoteOneSlaveStatefull slaveId oe
      LsRespExecute (RcRespCacheMiss CachedArg) -> do
        runRemoteOneSlaveStatefull slaveId e
        runRemoteOneSlaveStatefull slaveId oe
        -- todo uncache e if should not be cached
      LsRespExecute RcRespOk -> return ()
      LsRespExecute (RcRespError err) -> error ( "remote call: " ++ err)
      LsRespError err -> error ( "remote call: " ++ err)
      _ -> error ( "Should not reach here")
  
  runRemoteOneSlaveStatefull slaveId (MRConst n key _ _) = do
    s <- getRemote
    vault <- getVault
    case V.lookup key vault of
      Just partition -> do
        let bs = toData $ getPart slaveId partition
        let req = LsReqCache n bs
        r <- liftIO $ send s slaveId req
        -- TODO :verify next line.
        sr <- handleRpcErrorM r
        _ <- handleSlaveErrorM sr
        return ()
      Nothing ->  error ("MRConst value not cached"::String)
  
  
  
  runRemoteOneSlaveStateless ::(CommandClass s x, MonadIO m) => Int -> [SlaveRequest] -> MExp 'Remote a -> StateT (s x , V.Vault, InfoMap) m [SlaveRequest]
  runRemoteOneSlaveStateless slaveId requests (MRApply n (ExpClosure ce _) e) = do
    requests' <- runRemoteOneSlaveStateless slaveId requests e
    -- caching value of ce
    vault <- getVault
    let keyce = getLocalVaultKey ce
    let cem = V.lookup keyce vault
    case cem of
      Nothing -> error "local value not cached while executing remote on one slave"
      Just (_, Nothing) -> error "local value not cached (BS) while executing remote on one slave"
      Just (_, Just p) -> do
        let ceId = getLocalIndex ce
        let csBs = toData $ getPart slaveId p
        return (LsReqExecute n : LsReqCache ceId csBs : requests')
  
  runRemoteOneSlaveStateless slaveId requests (MRConst n key _ _) = do
    vault <- getVault
    case V.lookup key vault of
      Just partition -> do
        let bs = toData $ getPart slaveId partition
        return (LsReqCache n bs : requests)
      Nothing ->  error "MRConst value not cached"
  
  
  fetchOneSlaveResults :: forall a s x m.
    (S.Serialize a, CommandClass s x, MonadIO m)
    => Int -> MExp 'Remote a -> StateT (s x, V.Vault, InfoMap) m (Maybe (Data a))
  fetchOneSlaveResults slaveId e = do
    s <- getRemote
    let n = getRemoteIndex e
    let req = LsReqFetch n
    rE <- liftIO $ send s slaveId req
    case handleRpcError rE of
      LsRespFetch (Data bs) ->
        case S.decode bs of
          Right v -> return $ Just (Data v)
          Left err -> error ("Cannot decode fetched value: " ++ err)
      LsRespFetch NoData -> return $ Just NoData
      LsRespFetchMiss -> return Nothing
      LsRespError err -> error ("Cannot fetch results for node: "++ err)
      _ -> error ( "Should not reach here")
  
  --fetchOneSlaveResults :: forall a s x m.
  --  (S.Serialize a, CommandClass s x, MonadIO m)
  --  => Int -> MExp 'Remote a -> StateT (s x, V.Vault, InfoMap) m (Maybe a)
  --fetchOneSlaveResults slaveId e = do
  --  s <- getRemote
  --  let n = getRemoteIndex e
  --  let req = LsReqFetch n
  --  rE <- liftIO $ send s slaveId req
  --  case handleRpcError rE of
  --    LsRespFetch (Data bs) ->
  --      case S.decode bs of
  --        Right v -> return $ Just v
  --        Left err -> error ("Cannot decode fetched value: " ++ err)
  --    LsRespFetch NoData -> return Nothing
  --    LsRespError err -> error ("Cannot fetch results for node: "++ err)
  --    _ -> error ( "Should not reach here")
  
  
  --fetchResults :: (S.Serialize b, MonadIO m, CommandClass s x) =>
  --  UnChunkFun b a ->  MExp 'Remote b -> StateT (s x, V.Vault, InfoMap) m a
  --fetchResults unChunkFun e = do
  --  s <- getRemote
  --  let nbSlaves = getNbSlaves s
  --  let slaveIds = [0 .. nbSlaves - 1]
  --  st <- get
  --  r <- liftIO $ mapConcurrently (\slaveId -> evalStateT (fetchOneSlaveResults slaveId e) st) slaveIds
  --  return $ unChunkFun $ catMaybes r
  
  
  unCacheRemoteOneSlave :: (CommandClass s x, MonadIO m) => Int -> MExp 'Remote a -> StateT (s x , V.Vault, InfoMap) m ()
  unCacheRemoteOneSlave slaveId e = do
    s <- getRemote
    let req = LsReqUncache (getRemoteIndex e)
    rE <- liftIO $ send s slaveId req
    case handleRpcError rE of
      LsRespVoid -> return ()
      LsRespError err -> error ("Error, uncaching: "++err)
      _ -> error ( "Should not reach here")
  
  unCacheRemote :: (CommandClass s x, MonadIO m) => MExp 'Remote a -> StateT (s x , V.Vault, InfoMap) m ()
  unCacheRemote e = do
    s <- getRemote
    let nbSlaves = getNbSlaves s
    let slaveIds = [0 .. nbSlaves - 1]
    st <- get
    _ <- liftIO $ mapConcurrently (\slaveId -> evalStateT (unCacheRemoteOneSlave slaveId e) st) slaveIds
    return ()
  
  unCacheLocalOneSlave :: (CommandClass s x, MonadIO m) => Int -> MExp 'Local a -> StateT (s x , V.Vault, InfoMap) m ()
  unCacheLocalOneSlave slaveId e = do
    s <- getRemote
    let req = LsReqUncache (getLocalIndex e)
    rE <- liftIO $ send s slaveId req
    case handleRpcError rE of
      LsRespVoid -> return ()
      LsRespError err -> error ("Error, uncaching: "++err)
      _ -> error ("Should not reach here")
  
  unCacheLocal :: (CommandClass s x, MonadIO m) => MExp 'Local a -> StateT (s x , V.Vault, InfoMap) m ()
  unCacheLocal e = do
    s <- getRemote
    let nbSlaves = getNbSlaves s
    let slaveIds = [0 .. nbSlaves - 1]
    st <- get
    _ <- liftIO $ mapConcurrently (\slaveId -> evalStateT (unCacheLocalOneSlave slaveId e) st) slaveIds
    return ()
  
  runRemote ::(CommandClass s x, MonadLoggerIO m, S.Serialize b) =>
    UnChunkFun b a -> MExp 'Remote b -> StateT (s x, V.Vault, InfoMap) m a
  runRemote unChunkFun e = do
    prepareRunRemote e
    s <- getRemote
    case isStatefullSlave s of
      True -> do
        doRunAndFetchRemoteStatefull unChunkFun e
      False -> doRunRemoteStateless unChunkFun e
  
  doRunAndFetchRemoteStatefull ::(CommandClass s x, MonadLoggerIO m, S.Serialize b) =>
    UnChunkFun b a
    -> MExp 'Remote b
    -> StateT (s x, V.Vault, InfoMap) m a
  doRunAndFetchRemoteStatefull unChunkFun oe@(MRApply n (ExpClosure ce _) e) = do
    s <- getRemote
    doRunRemoteStatefull e
    cp <- runLocal ce
    case cp of
      (_, Just _) -> return ()
      (c, Nothing) -> do
          let nbSlaves = getNbSlaves s
          let partition = fmap S.encode $ chunk' nbSlaves c
          vault <- getVault
          let key = getLocalVaultKey ce
          setVault (V.insert key (c, Just partition) vault)
    let nbSlaves = getNbSlaves s
    let slaveIds = [0 .. nbSlaves - 1]
    st <- get
    r <- liftIO $ mapConcurrently (\slaveId -> evalStateT (runAndFetchRemoteOneSlaveStatefull slaveId oe) st) slaveIds
  
    -- dereference children and cleanup remote cache if necessary
    cleanupCacheE <- dereference n (getRemoteIndex e)
    when cleanupCacheE $ do unCacheRemote e
  
    cleanupCacheCe <- dereference n (getLocalIndex ce)
    when cleanupCacheCe $ do unCacheLocal ce
    return $ unChunkFun $ catMaybes r
  
  doRunRemoteStatefull ::(CommandClass s x, MonadLoggerIO m) =>
    MExp 'Remote a
    -> StateT (s x, V.Vault, InfoMap) m ()
  doRunRemoteStatefull e@(MRConst _ key chunkFun aio) = do
    a <- liftIO aio
    s <- getRemote
    vault <- getVault
    let nbSlaves = getNbSlaves s
    let slaveIds = [0 .. nbSlaves - 1]
    let partition = fmap S.encode $ chunkFun nbSlaves a
    let vault' = V.insert key partition vault
    setVault vault'
    st <- get
    _ <- liftIO $ mapConcurrently (\slaveId -> evalStateT (runRemoteOneSlaveStatefull slaveId e) st) slaveIds
    return ()
  
  doRunRemoteStateless :: forall a b m s x. (CommandClass s x, MonadLoggerIO m, S.Serialize b)
    => UnChunkFun b a -> MExp 'Remote b -> StateT (s x, V.Vault, InfoMap) m a
  doRunRemoteStateless unChunkFun oe@(MRApply n _ _) = do
    s <- getRemote
    let nbSlaves = getNbSlaves s
    let slaveIds = [0 .. nbSlaves - 1]
    st <- get
    rs <- liftIO $ mapConcurrently (\slaveId -> evalStateT (proc slaveId) st) slaveIds
    return $ unChunkFun $ catMaybes rs
    where
    proc slaveId = do
      requests <- runRemoteOneSlaveStateless slaveId [] oe
      -- fetch the results
      let requests' = L.reverse requests
      s <- getRemote
      let req = LsReqBatch n requests'
      aE <- liftIO $ send s slaveId req
      case handleRpcError aE of
        LsRespBatch (Data bs) ->
          case S.decode bs of
            Right a -> return $ Just (a::b)
            Left err -> error ("Cannot decode value from a batch execution: "++err)
        LsRespBatch NoData -> return Nothing
        LsRespError err -> error ("Batch error: "++err)
        _ -> error ( "Should not reach here")
  
  -- TODO uncomment
  doRunRemoteStateless unChunkFun (MRConst _ _ chunkFun aio) = do
    a <- liftIO aio
    return $ unChunkFun [(chunkFun 1 a Vc.! 0)]
  
  
  
  -- evaluate all local values (captured in closures)
  -- partition and cache MRConst's
  prepareRunRemote ::(CommandClass s x, MonadLoggerIO m) => MExp 'Remote a -> StateT (s x, V.Vault, InfoMap) m ()
  prepareRunRemote (MRApply _ (ExpClosure ce _) e) = do
    s <- getRemote
    prepareRunRemote e
    cp <- runLocal ce
    case cp of
      (_, Just _) -> return ()
      (c, Nothing) -> do
          let nbSlaves = getNbSlaves s
          let partition = fmap S.encode $ chunk' nbSlaves c
          vault <- getVault
          let key = getLocalVaultKey ce
          setVault (V.insert key (c, Just partition) vault)
    return ()
  
  
  prepareRunRemote (MRConst _ key chunkFun aio) = do
    a <- liftIO aio
    s <- getRemote
    vault <- getVault
    let nbSlaves = getNbSlaves s
    let partition = fmap S.encode $ chunkFun nbSlaves a
    let vault' = V.insert key partition vault
    setVault vault'
    return ()
  
  runLocal ::(CommandClass s x, MonadLoggerIO m) => MExp 'Local a -> StateT (s x, V.Vault, InfoMap) m (a, Maybe (Partition BS.ByteString))
  runLocal (MLConst _ key aio) = do
    vault <- getVault
    let cvm = V.lookup key vault
    case cvm of
      Just cv -> return cv
      Nothing -> do
        a <- liftIO aio
        setVault (V.insert key (a, Nothing) vault)
        return (a, Nothing)
  
  runLocal (MCollect n key unChunkFun e) = do
    vault <- getVault
    let cvm = V.lookup key vault
    case cvm of
      Just a -> return a
      Nothing -> do
        a <- runRemote unChunkFun e
        vault' <- getVault
        setVault (V.insert key (a, Nothing) vault')
  
          -- dereference children and clenup remote cache if necessary
        cleanupCacheE <- dereference n (getRemoteIndex e)
        when cleanupCacheE $ do unCacheRemote e
        return (a, Nothing)
  
  runLocal (MLApply _ key f e) = do
    vault <- getVault
    let cvm = V.lookup key vault
    case cvm of
      Just a -> return a
      Nothing -> do
        (a, _) <- runLocal e
        (f', _) <-  runLocal f
        let r = f' a
        vault' <- getVault
        setVault (V.insert key (r, Nothing) vault')
        return (r, Nothing)
  
  
  
  