content: {-# LANGUAGE TemplateHaskell, CPP, ScopedTypeVariables, OverloadedStrings #-}
  module IdeSession.RPC.Client (
      RpcServer
    , RpcConversation(..)
    , forkRpcServer
    , connectToRpcServer
    , rpc
    , rpcConversation
    , shutdown
    , forceShutdown
    , ExternalException(..)
    , illscopedConversationException
    , serverKilledException
    , getRpcExitCode
    , findProgram
    ) where
  
  import Control.Concurrent.MVar (MVar, newMVar, tryTakeMVar)
  import Control.Monad (void, unless)
  import Data.Binary (Binary, encode, decode)
  import Data.IORef (writeIORef, readIORef, newIORef)
  import Data.List (intercalate)
  import Data.Maybe (catMaybes)
  import Data.Monoid ((<>))
  import qualified Data.Text as Text
  import Data.Typeable (Typeable)
  import Prelude hiding (take)
  import System.Environment (lookupEnv)
  import System.Exit (ExitCode)
  import System.IO (Handle, hClose)
  import System.IO.Temp (openTempFile)
  
  import System.Process
    ( createProcess
    , proc
    , ProcessHandle
    , waitForProcess
    , CreateProcess(cwd, env)
    , getProcessExitCode
    )
  import IdeSession.Util.PortableProcess
  import qualified Control.Exception as Ex
  import qualified System.Directory  as Dir
  import Distribution.Verbosity (normal)
  import Distribution.Simple.Program.Find (
      findProgramOnSearchPath
    , ProgramSearchPath
    , ProgramSearchPathEntry(..)
    )
  
  import IdeSession.Util.BlockingOps
  import IdeSession.Util.Logger
  import IdeSession.RPC.API
  import IdeSession.RPC.Stream
  import IdeSession.RPC.Sockets
  
  import Network
  
  --------------------------------------------------------------------------------
  -- Client-side API                                                            --
  --------------------------------------------------------------------------------
  
  -- | Abstract data type representing RPC servers
  data RpcServer = RpcServer {
      -- | Handle to write requests to
      rpcRequestW  :: Handle
      -- | Temporary file the server will write uncaught exceptions to
    , rpcErrorLog :: FilePath
      -- | Handle on the server process itself
      --
      -- This is Nothing if we connected to an existing RPC server
      -- ('connectToRpcServer') rather than started a new server
      -- ('forkRpcServer')
    , rpcProc :: Maybe ProcessHandle
      -- | IORef containing the server response stream
    , rpcResponseR :: Stream Response
      -- | Server state
    , rpcState :: MVar RpcClientSideState
      -- | Identity of this server (for debugging purposes)
    , rpcIdentity :: String
    }
  
  -- | RPC server state
  data RpcClientSideState =
      -- | The server is running.
      RpcRunning
      -- | The server was stopped, either manually or because of an exception
    | RpcStopped Ex.SomeException
  
  -- | Fork an RPC server as a separate process
  --
  -- @forkRpcServer exec args@ starts executable @exec@ with arguments
  -- @args ++ args'@ where @args'@ are internal arguments generated by
  -- 'forkRpcServer'. These internal arguments should be passed as arguments
  -- to 'rpcServer'.
  --
  -- As a typical example, you might pass @["--server"]@ as @args@, and the
  -- 'main' function of @exec@ might look like
  --
  -- > main = do
  -- >   args <- getArgs
  -- >   case args of
  -- >     "--server" : args' ->
  -- >       rpcServer args' <<your request handler>>
  -- >     _ ->
  -- >       <<deal with other cases>>
  forkRpcServer :: FilePath        -- ^ Filename of the executable
                -> [String]        -- ^ Arguments
                -> Maybe FilePath  -- ^ Working directory
                -> Maybe [(String, String)] -- ^ Environment
                -> IO RpcServer
  forkRpcServer path args workingDir menv = do
    request <- makeSocket
    response <- makeSocket
  
    tmpDir <- Dir.getTemporaryDirectory
    (errorLogPath, errorLogHandle) <- openTempFile tmpDir "rpc-server-.log"
    hClose errorLogHandle
  
    ports <- mapM socketPort [request, response]
    let args' = args
             ++ [errorLogPath]
             ++ map portToString ports
  
    fullPath <- pathToExecutable path
    (Nothing, Nothing, Nothing, ph) <- createProcess (proc fullPath args') {
                                           cwd = workingDir,
                                           env = menv
                                         }
  
    request'  <- acceptHandle request
    response' <- acceptHandle response
  
    st    <- newMVar RpcRunning
    input <- newStream response'
    return RpcServer {
        rpcRequestW  = request'
      , rpcErrorLog  = errorLogPath
      , rpcProc      = Just ph
      , rpcState     = st
      , rpcResponseR = input
      , rpcIdentity  = path
      }
    where
      pathToExecutable :: FilePath -> IO FilePath
      pathToExecutable relPath = do
        fullPath    <- Dir.canonicalizePath relPath
        permissions <- Dir.getPermissions fullPath
        if Dir.executable permissions
          then return fullPath
          else Ex.throwIO . userError $ relPath ++ " not executable"
  
  -- | Connect to an existing RPC server
  --
  -- It is the responsibility of the caller to make sure that each triplet
  -- of named pipes is only used for RPC connection.
  connectToRpcServer :: WriteChannel   -- ^ stdin
                     -> ReadChannel   -- ^ stdout
                     -> FilePath   -- ^ logfile for storing exceptions
                     -> (RpcServer -> IO a)
                     -> IO a
  connectToRpcServer (WriteChannel request) (ReadChannel response) errorLog act =
    Ex.bracket (connectToPort request) hClose $ \request'  ->
    Ex.bracket (connectToPort response) hClose $ \response' -> do
      st    <- newMVar RpcRunning
      input <- newStream response'
      act $ RpcServer {
          rpcRequestW  = request'
        , rpcErrorLog  = errorLog
        , rpcProc      = Nothing
        , rpcState     = st
        , rpcResponseR = input
        , rpcIdentity  = (show request) -- TODO is it possible to make a better identifier?
        }
  
  -- | Specialized form of 'rpcConversation' to do single request and wait for
  -- a single response.
  rpc :: (Typeable req, Typeable resp, Binary req, Binary resp) => RpcServer -> req -> IO resp
  rpc server req = rpcConversation server $ \RpcConversation{..} -> put req >> get
  
  -- | Run an RPC conversation. If the handler throws an exception during
  -- the conversation the server is terminated.
  rpcConversation :: RpcServer
                  -> (RpcConversation -> IO a)
                  -> IO a
  rpcConversation server handler = withRpcServer server $ \st ->
    case st of
      RpcRunning -> do
        -- We want to be able to detect when a conversation is used out of scope
        inScope <- newIORef True
  
        -- Call the handler, update the state, and return the result
        a <- handler . conversation $ do isInScope <- readIORef inScope
                                         unless isInScope $
                                           Ex.throwIO illscopedConversationException
  
        -- Record that the conversation is no longer in scope and return
        writeIORef inScope False
        return (RpcRunning, a)
      RpcStopped ex ->
        Ex.throwIO ex
    where
      conversation :: IO () -> RpcConversation
      conversation verifyScope = RpcConversation {
          put = \req -> do
                   verifyScope
                   mapIOToExternal server $ do
                     let msg = encode $ Request (IncBS $ encode req)
                     hPutFlush (rpcRequestW server) msg
        , get = do verifyScope
                   mapIOToExternal server $ do
                     Response resp <- nextInStream (rpcResponseR server)
                     Ex.evaluate $ decode (unIncBS resp)
        }
  
  illscopedConversationException :: Ex.IOException
  illscopedConversationException =
    userError "Attempt to use RPC conversation outside its scope"
  
  -- | Shut down the RPC server
  --
  -- This simply kills the remote process. If you want to shut down the remote
  -- process cleanly you must implement your own termination protocol before
  -- calling 'shutdown'.
  shutdown :: RpcServer -> IO ()
  shutdown server = withRpcServer server $ \_ -> do
    terminate server
    ignoreIOExceptions $ Dir.removeFile (rpcErrorLog server)
    let ex = Ex.toException (userError "Manual shutdown")
    return (RpcStopped ex, ())
  
  -- | Force shutdown.
  --
  -- In order to faciliate a force shutdown while another thread may be
  -- communicating with the RPC server, we _try_ to update the MVar underlying
  -- the RPC server, but if we fail, we terminate the server anyway. This means
  -- that this may leave the 'RpcServer' in an invalid state -- so you shouldn't
  -- be using it anymore after calling forceShutdown!
  forceShutdown :: RpcServer -> IO ()
  forceShutdown server = Ex.mask_ $ do
    mst <- tryTakeMVar (rpcState server)
  
    ignoreIOExceptions $ forceTerminate server
    ignoreIOExceptions $ Dir.removeFile (rpcErrorLog server)
  
    case mst of
      Nothing -> -- We failed to take the MVar. Shrug.
        return ()
      Just _ -> do
        let ex = Ex.toException (userError "Forced manual shutdown")
        $putMVar (rpcState server) (RpcStopped ex)
  
  -- | Terminate the RPC connection
  --
  -- If we connected using 'forkRpcServer' (rather than 'connectToRpcServer')
  -- we wait for the remote process to terminate.
  terminate :: RpcServer -> IO ()
  terminate server = do
      ignoreIOExceptions $ hPutFlush (rpcRequestW server) (encode RequestShutdown)
      case rpcProc server of
        Just ph -> void $ waitForProcess ph
        Nothing -> return ()
  
  -- | Force-terminate the external process
  --
  -- Throws an exception when we are connected to an existing RPC server
  forceTerminate :: RpcServer -> IO ()
  forceTerminate server =
      case rpcProc server of
        Just ph -> killProcessHandle ph
        Nothing ->
          Ex.throwIO $ userError "forceTerminate: parallel connection"
  
  -- | Like modifyMVar, but terminate the server on exceptions
  withRpcServer :: RpcServer
                -> (RpcClientSideState -> IO (RpcClientSideState, a))
                -> IO a
  withRpcServer server io =
    Ex.mask $ \restore -> do
      st <- $takeMVar (rpcState server)
  
      mResult <- Ex.try $ restore (io st)
  
      case mResult of
        Right (st', a) -> do
          $putMVar (rpcState server) st'
          return a
        Left ex -> do
     --     terminate server
          $putMVar (rpcState server) (RpcStopped (Ex.toException (userError (rpcIdentity server ++ ": " ++ show (ex :: Ex.SomeException)))))
          Ex.throwIO ex
  
  -- | Get the exit code of the RPC server, unless still running.
  --
  -- Thross an exception for connections to existing RPC servers.
  getRpcExitCode :: RpcServer -> IO (Maybe ExitCode)
  getRpcExitCode RpcServer{rpcProc} =
    case rpcProc of
      Just ph -> getProcessExitCode ph
      Nothing -> Ex.throwIO $ userError "getRpcExitCode: parallel connection"
  
  {------------------------------------------------------------------------------
    Aux
  ------------------------------------------------------------------------------}
  
  -- | Map IO exceptions to external exceptions, using the error written
  -- by the server (if any)
  mapIOToExternal :: RpcServer -> IO a -> IO a
  mapIOToExternal server p = Ex.catch p $ \ex -> do
    let _ = ex :: Ex.IOException
    merr <- readFile (rpcErrorLog server)
    if null merr
      then Ex.throwIO (serverKilledException (Just ex))
      else Ex.throwIO (ExternalException merr (Just ex))
  
  findProgram :: LogFunc -> ProgramSearchPath -> FilePath -> IO (Maybe FilePath)
  findProgram logFunc searchPath prog = do
      shownPath <- renderPath searchPath
      $logInfo $ "Searching for " <> Text.pack prog <> " on this path: " <> Text.pack shownPath
      mres <- findProgramOnSearchPath normal searchPath prog
      $logInfo $ case mres of
        Nothing -> "Failed to find " <> Text.pack prog
        Just res -> "Found " <> Text.pack prog <> " - using this one: " <> Text.pack res
      return mres
    where
      renderPath = fmap (intercalate ":" . catMaybes) . mapM pathEntryString
      pathEntryString (ProgramSearchPathDir fp) = return (Just fp)
      pathEntryString ProgramSearchPathDefault = lookupEnv "PATH"
  