{-# 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"