{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE Trustworthy #-}
-------------------------------------------------------------------------------
-- |
-- Module      :  System.Posix.Pty
-- Copyright   :  (C) 2013 Merijn Verstraaten
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Merijn Verstraaten <merijn@inconsistent.nl>
-- Stability   :  experimental
-- Portability :  haha
--
-- A module for interacting with subprocesses through a pseudo terminal (pty).
-- Provides functions for reading from, writing to and resizing pseudo
-- terminals. Re-exports most of "System.Posix.Terminal", providing wrappers
-- that work with the 'Pty' type where necessary.
-------------------------------------------------------------------------------
module System.Posix.Pty (
    -- * Subprocess Creation
      spawnWithPty
    -- * Data Structures
    , Pty
    , PtyControlCode (..)
    -- * Pty Interaction Functions
    , createPty
    , tryReadPty
    , readPty
    , writePty
    , resizePty
    , ptyDimensions
    -- * Re-exports of "System.Posix.Terminal"
    -- $posix-reexport
    , getTerminalAttributes
    , setTerminalAttributes
    , sendBreak
    , drainOutput
    , discardData
    , controlFlow
    , getTerminalProcessGroupID
    , getTerminalName
    , getSlaveTerminalName
    , module System.Posix.Terminal
    ) where

import Control.Applicative
import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BSI

import Unsafe.Coerce (unsafeCoerce)
import Foreign
import Foreign.C.String (CString, newCString, peekCString)
import Foreign.C.Types
import Foreign.C.Error (Errno(..), getErrno)

#if defined(linux_HOST_OS)
import Foreign.C.Error (eIO)
import System.IO.Error (catchIOError)
#endif

import System.IO (Handle)
import System.IO.Error (mkIOError, eofErrorType)
import System.Posix.IO.ByteString (fdToHandle, fdReadBuf)
import System.Posix.Types
import System.Process (ProcessHandle)
import System.Process.Internals (mkProcessHandle)

import qualified System.Posix.Terminal as T
import System.Posix.Terminal hiding
        ( getTerminalAttributes
        , setTerminalAttributes
        , sendBreak
        , drainOutput
        , discardData
        , controlFlow
        , getTerminalProcessGroupID
        , setTerminalProcessGroupID
        , queryTerminal
        , getTerminalName
        , openPseudoTerminal
        , getSlaveTerminalName)

-- | Abstract pseudo terminal type.
data Pty = Pty !Fd !Handle

-- | Pseudo terminal control information.
--
-- [Terminal read queue] The terminal read queue contains the data that was
-- written from the master terminal to the slave terminal, which was not read
-- from the slave yet.
--
-- [Terminal write queue] The terminal write queue contains the data that was
-- written from the slave terminal, which was not sent to the master yet.
data PtyControlCode = FlushRead     -- ^ Terminal read queue was flushed.
                    | FlushWrite    -- ^ Terminal write queue was flushed.
                    | OutputStopped -- ^ Terminal output was stopped.
                    | OutputStarted -- ^ Terminal output was restarted.
                    | DoStop        -- ^ Terminal stop and start characters are
                                    --   @^S@ and @^Q@ respectively.
                    | NoStop        -- ^ Terminal stop and start characters are
                                    --   NOT @^S@ and @^Q@.
                    deriving (Eq, Read, Show)

-- | Produces a 'Pty' if the file descriptor is associated with a terminal and
-- Nothing if not.
createPty :: Fd -> IO (Maybe Pty)
createPty fd = do
    isTerm <- T.queryTerminal fd
    if isTerm
       then Just . Pty fd <$> fdToHandle fd
       else return Nothing

-- | Read data from a fd using read() system call avoiding GHC.IO buffering.
--
-- Throws an 'IOError' of type 'eofErrorType'.
fdReadBS :: Fd -> ByteCount -> IO ByteString
fdReadBS fd n
  | n <= 0    = return BS.empty
  | otherwise = BSI.createAndTrim (fromIntegral n) fill
  where
    fill buf = do
      rc <- wrap (fdReadBuf fd buf n)
      case rc of
        _ | rc == 0 -> eof
          | otherwise -> return (fromIntegral rc)

    wrap :: IO a -> IO a
#if defined(linux_HOST_OS)
    -- Linux indicates slave pty EOF as EIO
    -- https://lkml.org/lkml/2009/4/8/578
    wrap action = catchIOError action $ \ioE -> do
      errno <- getErrno
      case errno of
          e | e == eIO -> eof
          _ -> ioError ioE
#else
    wrap = id
#endif

    eof = do
      hnd <- fdToHandle fd
      ioError $ mkIOError eofErrorType "eof" (Just hnd) Nothing

-- | Attempt to read data from a pseudo terminal. Produces either the data read
-- or a list of 'PtyControlCode'@s@ indicating which control status events that
-- have happened on the slave terminal.
--
-- Throws an 'IOError' of type 'eofErrorType' when the terminal has been
-- closed, for example when the subprocess has terminated.
tryReadPty :: Pty -> IO (Either [PtyControlCode] ByteString)
tryReadPty (Pty fd _) = do
    result <- fdReadBS fd 1024
    case BS.uncons result of
         Just (byte, rest)
            | byte == 0    -> return (Right rest)
            | BS.null rest -> return (Left $ byteToControlCode byte)
            | otherwise    -> ioError can'tHappen
         Nothing -> ioError can'tHappen
  where
    can'tHappen = userError "Uh-oh! Something different went horribly wrong!"

-- | The same as 'tryReadPty', but discards any control status events.
readPty :: Pty -> IO ByteString
readPty pty = tryReadPty pty >>= \case
                   Left _ -> readPty pty
                   Right bs -> return bs

-- | Write a 'ByteString' to the pseudo terminal, throws an 'IOError' when the
-- terminal has been closed, for example when the subprocess has terminated.
writePty :: Pty -> ByteString -> IO ()
writePty (Pty _ hnd) = BS.hPut hnd

-- | Set the pseudo terminal's dimensions to the specified width and height.
resizePty :: Pty -> (Int, Int) -> IO ()
resizePty (Pty fd _) (x, y) =
    set_pty_size fd x y >>= throwCErrorOnMinus1 "unable to set pty dimensions"

-- | Produces the pseudo terminal's current dimensions.
ptyDimensions :: Pty -> IO (Int, Int)
ptyDimensions (Pty fd _) = alloca $ \x -> alloca $ \y -> do
    get_pty_size fd x y >>= throwCErrorOnMinus1 "unable to get pty size"
    (,) <$> peek x <*> peek y

-- | Create a new process that is connected to the current process through a
-- pseudo terminal. If an environment is specified, then only the specified
-- environment variables will be set. If no environment is specified the
-- process will inherit its environment from the current process. Example:
--
-- > pty <- spawnWithPty (Just [("SHELL", "tcsh")]) True "ls" ["-l"] (20, 10)
--
-- This searches the user's PATH for a binary called @ls@, then runs this
-- binary with the commandline argument @-l@ in a terminal that is 20
-- characters wide and 10 characters high. The environment of @ls@ will
-- contains one variable, SHELL, which will be set to the value \"tcsh\".
spawnWithPty :: Maybe [(String, String)]    -- ^ Optional environment for the
                                            --   new process.
             -> Bool                        -- ^ Search for the executable in
                                            --   PATH?
             -> FilePath                    -- ^ Program's name.
             -> [String]                    -- ^ Command line arguments for the
                                            --   program.
             -> (Int, Int)                  -- ^ Initial dimensions for the
                                            --   pseudo terminal.
             -> IO (Pty, ProcessHandle)
spawnWithPty env' search path' argv' (x, y) = do
    path <- newCString path'
    argv <- mapM newCString argv'
    env <- maybe (return []) (mapM fuse) env'

    (ptyFd, cpid) <- forkExecWithPty x y path (fromBool search) argv env

    mapM_ free (env ++ argv)
    free path

    throwCErrorOnMinus1 "unable to fork or open new pty" ptyFd

    hnd <- fdToHandle ptyFd
    ph <- mkProcessHandle (unsafeCoerce cpid) False
    return (Pty ptyFd hnd, ph)
  where
    fuse (key, val) = newCString (key ++ "=" ++ val)

-- Module internal functions

getFd :: Pty -> Fd
getFd (Pty fd _) = fd

throwCErrorOnMinus1 :: (Eq a, Num a) => String -> a -> IO ()
throwCErrorOnMinus1 s i = when (i == -1) $ do
    errnoMsg <- getErrno >>= \(Errno code) -> (peekCString . strerror) code
    ioError . userError $ s ++ ": " ++ errnoMsg

forkExecWithPty :: Int
                -> Int
                -> CString
                -> CInt
                -> [CString]
                -> [CString]
                -> IO (Fd, CInt)
forkExecWithPty x y path search argv' env' = do
    argv <- newArray0 nullPtr (path:argv')
    env <- case env' of
                [] -> return nullPtr
                _  -> newArray0 nullPtr env'

    alloca $ \pid -> do
      result <- fork_exec_with_pty x y search path argv env pid
      free argv >> free env
      pid' <- peek pid
      return (result, pid')

byteToControlCode :: Word8 -> [PtyControlCode]
byteToControlCode i = map snd $ filter ((/=0) . (.&.i) . fst) codeMapping
    where codeMapping :: [(Word8, PtyControlCode)]
          codeMapping =
            [ (tiocPktFlushRead,  FlushRead)
            , (tiocPktFlushWrite, FlushWrite)
            , (tiocPktStop,       OutputStopped)
            , (tiocPktStart,      OutputStarted)
            , (tiocPktDoStop,     DoStop)
            , (tiocPktNoStop,     NoStop)
            ]

-- Foreign imports

tiocPktFlushRead :: Word8
tiocPktFlushRead = 1

tiocPktFlushWrite :: Word8
tiocPktFlushWrite = 2

tiocPktStop :: Word8
tiocPktStop = 4

tiocPktStart :: Word8
tiocPktStart = 8

tiocPktDoStop :: Word8
tiocPktDoStop = 32

tiocPktNoStop :: Word8
tiocPktNoStop = 16

foreign import ccall unsafe "string.h"
    strerror :: CInt -> CString

foreign import ccall "pty_size.h"
    set_pty_size :: Fd -> Int -> Int -> IO CInt

foreign import ccall "pty_size.h"
    get_pty_size :: Fd -> Ptr Int -> Ptr Int -> IO CInt

foreign import ccall "fork_exec_with_pty.h"
    fork_exec_with_pty :: Int
                       -> Int
                       -> CInt
                       -> CString
                       -> Ptr CString
                       -> Ptr CString
                       -> Ptr CInt
                       -> IO Fd

-- Pty specialised re-exports of System.Posix.Terminal

{- $posix-reexport
This module re-exports the entirety of "System.Posix.Terminal", with the
exception of the following functions:

[setTerminalProcessGroupID] This function can't be used after a process using
the slave terminal has been created, rendering it mostly useless for working
with 'Pty'@s@ created by this module.

[queryTerminal] Useless, 'Pty' is always a terminal.

[openPseudoTerminal] Only useful for the kind of tasks this module is supposed
abstract away.

In addition, some functions from "System.Posix.Terminal" work directly with
'Fd'@s@, these have been hidden and instead the following replacements working
on 'Pty'@s@ are exported.
-}

-- | See 'System.Posix.Terminal.getTerminalAttributes'.
getTerminalAttributes :: Pty -> IO TerminalAttributes
getTerminalAttributes = T.getTerminalAttributes . getFd

-- | See 'System.Posix.Terminal.setTerminalAttributes'.
setTerminalAttributes :: Pty -> TerminalAttributes -> TerminalState -> IO ()
setTerminalAttributes = T.setTerminalAttributes . getFd

-- | See 'System.Posix.Terminal.sendBreak'.
sendBreak :: Pty -> Int -> IO ()
sendBreak = T.sendBreak . getFd

-- | See 'System.Posix.Terminal.drainOutput'.
drainOutput :: Pty -> IO ()
drainOutput = T.drainOutput . getFd

-- | See 'System.Posix.Terminal.discardData'.
discardData :: Pty -> QueueSelector -> IO ()
discardData = T.discardData . getFd

-- | See 'System.Posix.Terminal.controlFlow'.
controlFlow :: Pty -> FlowAction -> IO ()
controlFlow = T.controlFlow . getFd

-- | See 'System.Posix.Terminal.getTerminalProcessGroupID'.
getTerminalProcessGroupID :: Pty -> IO ProcessGroupID
getTerminalProcessGroupID = T.getTerminalProcessGroupID . getFd

-- | See 'System.Posix.Terminal.getTerminalName'.
getTerminalName :: Pty -> IO FilePath
getTerminalName = T.getTerminalName . getFd

-- | See 'System.Posix.Terminal.getSlaveTerminalName'.
getSlaveTerminalName :: Pty -> IO FilePath
getSlaveTerminalName = T.getSlaveTerminalName . getFd