{- Copyright 2017 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

module Pty (
	Pty,
	runWithPty,
	readPty,
	writePty,
	inRawMode,
	withoutEcho,
	withTerminalAttributes,
) where

import System.Posix
import System.Posix.Pty
import qualified System.Console.Terminal.Size as Console
import System.Posix.Signals.Exts
import System.Process
import Control.Exception

-- | Run a program on a Pty.
-- 
-- While doing so, the outer pty has echo disabled (so the child can echo),
-- and has raw mode enabled (so the child pty can see all special characters).
--
-- A SIGWINCH handler is installed, to forward resizes to the Pty.
runWithPty :: String -> [String] -> ((Pty, ProcessHandle) -> IO a) -> IO a
runWithPty cmd params a = bracket setup cleanup go
  where
	setup = do
		as <- System.Posix.getTerminalAttributes stdInput
		sz <- Console.size
		(p, ph) <- spawnWithPty Nothing True cmd params
			(maybe 80 Console.width sz, maybe 25 Console.height sz)
		_ <- installHandler windowChange (Catch (forwardresize p)) Nothing
		-- Set the pty's terminal attributes to the same ones that
		-- the outer terminal had.
		System.Posix.Pty.setTerminalAttributes p as Immediately
		System.Posix.setTerminalAttributes stdInput (setRawMode as) Immediately
		return (p, ph, as)
	cleanup (p, ph, as) = do
		-- Needed in case the provided action throws an exception
		-- before it waits for the process.
		terminateProcess ph
		closePty p
		_ <- installHandler windowChange Default Nothing
		System.Posix.setTerminalAttributes stdInput as Immediately
	go (p, ph, _) = a (p, ph)
	forwardresize p = do
		msz <- Console.size
		case msz of
			Nothing -> return ()
			Just sz -> resizePty p (Console.width sz, Console.height sz)

withTerminalAttributes :: (TerminalAttributes -> TerminalAttributes) -> IO a -> IO a
withTerminalAttributes f a = bracket setup cleanup go
  where
	setup = do
		as <- System.Posix.getTerminalAttributes stdInput
		System.Posix.setTerminalAttributes stdInput (f as) Immediately
		return as
	cleanup as = System.Posix.setTerminalAttributes stdInput as Immediately
	go _ = a

-- | This is similar to cfmakeraw(3).
inRawMode :: IO a -> IO a
inRawMode = withTerminalAttributes setRawMode

setRawMode :: TerminalAttributes -> TerminalAttributes
setRawMode as = as
	`withoutMode` IgnoreBreak
	`withoutMode` InterruptOnBreak
	`withoutMode` CheckParity
	`withoutMode` StripHighBit
	`withoutMode` MapLFtoCR
	`withoutMode` IgnoreCR
	`withoutMode` MapCRtoLF
	`withoutMode` StartStopOutput
	`withoutMode` ProcessOutput
	`withoutMode` EnableEcho
	`withoutMode` EchoLF
	`withoutMode` ProcessInput
	`withoutMode` KeyboardInterrupts
	`withoutMode` ExtendedFunctions
	`withoutMode` EnableParity

-- | Disable terminal echo.
withoutEcho :: IO a -> IO a
withoutEcho = withTerminalAttributes $ \as -> as
	`withoutMode` EnableEcho

