    {-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances, TupleSections #-}
    {-# LANGUAGE CPP #-}
    {-# OPTIONS_GHC -O2 #-}
    {- Building this module with -O0 causes streams not to fuse and too much
     - memory to be used. -}

    -- |
    -- Copyright: 2015 Joey Hess <id@joeyh.name>
    -- License: BSD-2-clause
    --
    -- Concurrent output handling, internals.
    --
    -- May change at any time.

    module System.Console.Concurrent.Internal where

    import System.IO
    #ifndef mingw32_HOST_OS
    import System.Posix.IO
    #endif
    import System.Directory
    import System.Exit
    import Control.Monad
    import Control.Monad.IO.Class (liftIO, MonadIO)
    import System.IO.Unsafe (unsafePerformIO)
    import Control.Concurrent
    import Control.Concurrent.STM
    import Control.Concurrent.Async
    import Data.Maybe
    import Data.List
    import Data.Monoid
    import qualified System.Process as P
    import qualified Data.Text as T
    import qualified Data.Text.IO as T
    import Control.Applicative
    import Prelude

    import Utility.Monad
    import Utility.Exception

    data OutputHandle = OutputHandle
    	{ outputLock :: TMVar Lock
    	, outputBuffer :: TMVar OutputBuffer
    	, errorBuffer :: TMVar OutputBuffer
    	, outputThreads :: TMVar Integer
    	, processWaiters :: TMVar [Async ()]
    	, waitForProcessLock :: TMVar ()
    	}

    data Lock = Locked

    -- | A shared global variable for the OutputHandle.
    {-# NOINLINE globalOutputHandle #-}
    globalOutputHandle :: OutputHandle
    globalOutputHandle = unsafePerformIO $ OutputHandle
    	<$> newEmptyTMVarIO
    	<*> newTMVarIO (OutputBuffer [])
    	<*> newTMVarIO (OutputBuffer [])
    	<*> newTMVarIO 0
    	<*> newTMVarIO []
    	<*> newEmptyTMVarIO

    -- | Holds a lock while performing an action. This allows the action to
    -- perform its own output to the console, without using functions from this
    -- module.
    --
    -- While this is running, other threads that try to lockOutput will block.
    -- Any calls to `outputConcurrent` and `createProcessConcurrent` will not
    -- block, but the output will be buffered and displayed only once the
    -- action is done.
    lockOutput :: (MonadIO m, MonadMask m) => m a -> m a
    lockOutput = bracket_ (liftIO takeOutputLock) (liftIO dropOutputLock)

    -- | Blocks until we have the output lock.
    takeOutputLock :: IO ()
    takeOutputLock = void $ takeOutputLock' True

    -- | Tries to take the output lock, without blocking.
    tryTakeOutputLock :: IO Bool
    tryTakeOutputLock = takeOutputLock' False

    withLock :: (TMVar Lock -> STM a) -> IO a
    withLock a = atomically $ a (outputLock globalOutputHandle)

    takeOutputLock' :: Bool -> IO Bool
    takeOutputLock' block = do
    	locked <- withLock $ \l -> do
    		v <- tryTakeTMVar l
    		case v of
    			Just Locked
    				| block -> retry
    				| otherwise -> do
    					-- Restore value we took.
    					putTMVar l Locked
    					return False
    			Nothing -> do
    				putTMVar l Locked
    				return True
    	when locked $ do
    		(outbuf, errbuf) <- atomically $ (,)
    			<$> swapTMVar (outputBuffer globalOutputHandle) (OutputBuffer [])
    			<*> swapTMVar (errorBuffer globalOutputHandle) (OutputBuffer [])
    		emitOutputBuffer StdOut outbuf
    		emitOutputBuffer StdErr errbuf
    	return locked

    -- | Only safe to call after taking the output lock.
    dropOutputLock :: IO ()
    dropOutputLock = withLock $ void . takeTMVar

    -- | Use this around any actions that use `outputConcurrent`
    -- or `createProcessConcurrent`
    --
    -- This is necessary to ensure that buffered concurrent output actually
    -- gets displayed before the program exits.
    withConcurrentOutput :: (MonadIO m, MonadMask m) => m a -> m a
    withConcurrentOutput a = a `finally` liftIO flushConcurrentOutput

    -- | Blocks until any processes started by `createProcessConcurrent` have
    -- finished, and any buffered output is displayed. Also blocks while
    -- `lockOutput` is is use.
    --
    -- `withConcurrentOutput` calls this at the end, so you do not normally
    -- need to use this.
    flushConcurrentOutput :: IO ()
    flushConcurrentOutput = do
    	atomically $ do
    		r <- takeTMVar (outputThreads globalOutputHandle)
    		if r <= 0
    			then putTMVar (outputThreads globalOutputHandle) r
    			else retry
    	-- Take output lock to wait for anything else that might be
    	-- currently generating output.
    	lockOutput $ return ()

    -- | Values that can be output.
    class Outputable v where
    	toOutput :: v -> T.Text

    instance Outputable T.Text where
    	toOutput = id

    instance Outputable String where
    	toOutput = toOutput . T.pack

    -- | Displays a value to stdout.
    --
    -- No newline is appended to the value, so if you want a newline, be sure
    -- to include it yourself.
    --
    -- Uses locking to ensure that the whole output occurs atomically
    -- even when other threads are concurrently generating output.
    --
    -- When something else is writing to the console at the same time, this does
    -- not block. It buffers the value, so it will be displayed once the other
    -- writer is done.
    outputConcurrent :: Outputable v => v -> IO ()
    outputConcurrent = outputConcurrent' StdOut

    -- | Like `outputConcurrent`, but displays to stderr.
    --
    -- (Does not throw an exception.)
    errorConcurrent :: Outputable v => v -> IO ()
    errorConcurrent = outputConcurrent' StdErr

    outputConcurrent' :: Outputable v => StdHandle -> v -> IO ()
    outputConcurrent' stdh v = bracket setup cleanup go
      where
    	setup = tryTakeOutputLock
    	cleanup False = return ()
    	cleanup True = dropOutputLock
    	go True = do
    		T.hPutStr h (toOutput v)
    		hFlush h
    	go False = do
    		oldbuf <- atomically $ takeTMVar bv
    		newbuf <- addOutputBuffer (Output (toOutput v)) oldbuf
    		atomically $ putTMVar bv newbuf
    	h = toHandle stdh
    	bv = bufferFor stdh

    newtype ConcurrentProcessHandle = ConcurrentProcessHandle P.ProcessHandle

    toConcurrentProcessHandle :: (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle) -> (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
    toConcurrentProcessHandle (i, o, e, h) = (i, o, e, ConcurrentProcessHandle h)

    -- | Use this to wait for processes started with
    -- `createProcessConcurrent` and `createProcessForeground`, and get their
    -- exit status.
    --
    -- Note that such processes are actually automatically waited for
    -- internally, so not calling this explicitly will not result
    -- in zombie processes. This behavior differs from `P.waitForProcess`
    waitForProcessConcurrent :: ConcurrentProcessHandle -> IO ExitCode
    waitForProcessConcurrent (ConcurrentProcessHandle h) =
    	bracket lock unlock checkexit
      where
    	lck = waitForProcessLock globalOutputHandle
    	lock = atomically $ tryPutTMVar lck ()
    	unlock True = atomically $ takeTMVar lck
    	unlock False = return ()
    	checkexit locked = maybe (waitsome locked) return
    		=<< P.getProcessExitCode h
    	waitsome True = do
    		let v = processWaiters globalOutputHandle
    		l <- atomically $ readTMVar v
    		if null l
    			-- Avoid waitAny [] which blocks forever
    			then P.waitForProcess h
    			else do
    				-- Wait for any of the running
    				-- processes to exit. It may or may not
    				-- be the one corresponding to the
    				-- ProcessHandle. If it is,
    				-- getProcessExitCode will succeed.
    				void $ tryIO $ waitAny l
    				checkexit True
    	waitsome False = do
    		-- Another thread took the lck first. Wait for that thread to
    		-- wait for one of the running processes to exit.
    		atomically $ do
    			putTMVar lck ()
    			takeTMVar lck
    		checkexit False

    -- Registers an action that waits for a process to exit,
    -- adding it to the processWaiters list, and removing it once the action
    -- completes.
    asyncProcessWaiter :: IO () -> IO ()
    asyncProcessWaiter waitaction = do
    	regdone <- newEmptyTMVarIO
    	waiter <- async $ do
    		self <- atomically (takeTMVar regdone)
    		waitaction `finally` unregister self
    	register waiter regdone
      where
    	v = processWaiters globalOutputHandle
      	register waiter regdone = atomically $ do
    		l <- takeTMVar v
    		putTMVar v (waiter:l)
    		putTMVar regdone waiter
    	unregister waiter = atomically $ do
    		l <- takeTMVar v
    		putTMVar v (filter (/= waiter) l)

    -- | Wrapper around `System.Process.createProcess` that prevents
    -- multiple processes that are running concurrently from writing
    -- to stdout/stderr at the same time.
    --
    -- If the process does not output to stdout or stderr, it's run
    -- by createProcess entirely as usual. Only processes that can generate
    -- output are handled specially:
    --
    -- A process is allowed to write to stdout and stderr in the usual
    -- way, assuming it can successfully take the output lock.
    --
    -- When the output lock is held (ie, by another concurrent process,
    -- or because `outputConcurrent` is being called at the same time),
    -- the process is instead run with its stdout and stderr
    -- redirected to a buffer. The buffered output will be displayed as soon
    -- as the output lock becomes free.
    --
    -- Currently only available on Unix systems, not Windows.
    #ifndef mingw32_HOST_OS
    createProcessConcurrent :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
    createProcessConcurrent p
    	| willOutput (P.std_out p) || willOutput (P.std_err p) =
    		ifM tryTakeOutputLock
    			( fgProcess p
    			, bgProcess p
    			)
    	| otherwise = do
    		r@(_, _, _, h) <- P.createProcess p
    		asyncProcessWaiter $
    			void $ tryIO $ P.waitForProcess h
    		return (toConcurrentProcessHandle r)
    #endif

    -- | Wrapper around `System.Process.createProcess` that makes sure a process
    -- is run in the foreground, with direct access to stdout and stderr.
    -- Useful when eg, running an interactive process.
    createProcessForeground :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
    createProcessForeground p = do
    	takeOutputLock
    	fgProcess p

    fgProcess :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
    fgProcess p = do
    	r@(_, _, _, h) <- P.createProcess p
    		`onException` dropOutputLock
    	registerOutputThread
    	-- Wait for the process to exit and drop the lock.
    	asyncProcessWaiter $ do
    		void $ tryIO $ P.waitForProcess h
    		unregisterOutputThread
    		dropOutputLock
    	return (toConcurrentProcessHandle r)

    #ifndef mingw32_HOST_OS
    bgProcess :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
    bgProcess p = do
    	(toouth, fromouth) <- pipe
    	(toerrh, fromerrh) <- pipe
    	let p' = p
    		{ P.std_out = rediroutput (P.std_out p) toouth
    		, P.std_err = rediroutput (P.std_err p) toerrh
    		}
    	registerOutputThread
    	r@(_, _, _, h) <- P.createProcess p'
    		`onException` unregisterOutputThread
    	asyncProcessWaiter $ void $ tryIO $ P.waitForProcess h
    	outbuf <- setupOutputBuffer StdOut toouth (P.std_out p) fromouth
    	errbuf <- setupOutputBuffer StdErr toerrh (P.std_err p) fromerrh
    	void $ async $ bufferWriter [outbuf, errbuf]
    	return (toConcurrentProcessHandle r)
      where
    	pipe = do
    		(from, to) <- createPipe
    		(,) <$> fdToHandle to <*> fdToHandle from
    	rediroutput ss h
    		| willOutput ss = P.UseHandle h
    		| otherwise = ss
    #endif

    willOutput :: P.StdStream -> Bool
    willOutput P.Inherit = True
    willOutput _ = False

    -- | Buffered output.
    data OutputBuffer = OutputBuffer [OutputBufferedActivity]
    	deriving (Eq)

    data StdHandle = StdOut | StdErr

    toHandle :: StdHandle -> Handle
    toHandle StdOut = stdout
    toHandle StdErr = stderr

    bufferFor :: StdHandle -> TMVar OutputBuffer
    bufferFor StdOut = outputBuffer globalOutputHandle
    bufferFor StdErr = errorBuffer globalOutputHandle

    data OutputBufferedActivity
    	= Output T.Text
    	| InTempFile
    		{ tempFile :: FilePath
    		, endsInNewLine :: Bool
    		}
    	deriving (Eq)

    data AtEnd = AtEnd
    	deriving Eq

    data BufSig = BufSig

    setupOutputBuffer :: StdHandle -> Handle -> P.StdStream -> Handle -> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
    setupOutputBuffer h toh ss fromh = do
    	hClose toh
    	buf <- newMVar (OutputBuffer [])
    	bufsig <- atomically newEmptyTMVar
    	bufend <- atomically newEmptyTMVar
    	void $ async $ outputDrainer ss fromh buf bufsig bufend
    	return (h, buf, bufsig, bufend)

    -- Drain output from the handle, and buffer it.
    outputDrainer :: P.StdStream -> Handle -> MVar OutputBuffer -> TMVar BufSig -> TMVar AtEnd -> IO ()
    outputDrainer ss fromh buf bufsig bufend
    	| willOutput ss = go
    	| otherwise = atend
      where
    	go = do
    		t <- T.hGetChunk fromh
    		if T.null t
    			then atend
    			else do
    				modifyMVar_ buf $ addOutputBuffer (Output t)
    				changed
    				go
    	atend = do
    		atomically $ putTMVar bufend AtEnd
    		hClose fromh
    	changed = atomically $ do
    		void $ tryTakeTMVar bufsig
    		putTMVar bufsig BufSig

    registerOutputThread :: IO ()
    registerOutputThread = do
    	let v = outputThreads globalOutputHandle
    	atomically $ putTMVar v . succ =<< takeTMVar v

    unregisterOutputThread :: IO ()
    unregisterOutputThread = do
    	let v = outputThreads globalOutputHandle
    	atomically $ putTMVar v . pred =<< takeTMVar v

    -- Wait to lock output, and once we can, display everything
    -- that's put into the buffers, until the end.
    --
    -- If end is reached before lock is taken, instead add the command's
    -- buffers to the global outputBuffer and errorBuffer.
    bufferWriter :: [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)] -> IO ()
    bufferWriter ts = do
    	activitysig <- atomically newEmptyTMVar
    	worker1 <- async $ lockOutput $
    		ifM (atomically $ tryPutTMVar activitysig ())
    			( void $ mapConcurrently displaybuf ts
    			, noop -- buffers already moved to global
    			)
    	worker2 <- async $ void $ globalbuf activitysig worker1
    	void $ async $ do
    		void $ waitCatch worker1
    		void $ waitCatch worker2
    		unregisterOutputThread
      where
    	displaybuf v@(outh, buf, bufsig, bufend) = do
    		change <- atomically $
    			(Right <$> takeTMVar bufsig)
    				`orElse`
    			(Left <$> takeTMVar bufend)
    		l <- takeMVar buf
    		putMVar buf (OutputBuffer [])
    		emitOutputBuffer outh l
    		case change of
    			Right BufSig -> displaybuf v
    			Left AtEnd -> return ()
    	globalbuf activitysig worker1 = do
    		ok <- atomically $ do
    			-- signal we're going to handle it
    			-- (returns false if the displaybuf already did)
    			ok <- tryPutTMVar activitysig ()
    			-- wait for end of all buffers
    			when ok $
    				mapM_ (\(_outh, _buf, _bufsig, bufend) -> takeTMVar bufend) ts
    			return ok
    		when ok $ do
    			-- add all of the command's buffered output to the
    			-- global output buffer, atomically
    			bs <- forM ts $ \(outh, buf, _bufsig, _bufend) ->
    				(outh,) <$> takeMVar buf
    			atomically $
    				forM_ bs $ \(outh, b) ->
    					bufferOutputSTM' outh b
    			-- worker1 might be blocked waiting for the output
    			-- lock, and we've already done its job, so cancel it
    			cancel worker1

    -- Adds a value to the OutputBuffer. When adding Output to a Handle,
    -- it's cheaper to combine it with any already buffered Output to that
    -- same Handle.
    --
    -- When the total buffered Output exceeds 1 mb in size, it's moved out of
    -- memory, to a temp file. This should only happen rarely, but is done to
    -- avoid some verbose process unexpectedly causing excessive memory use.
    addOutputBuffer :: OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
    addOutputBuffer (Output t) (OutputBuffer buf)
    	| T.length t' <= 1048576 = return $ OutputBuffer (Output t' : other)
    	| otherwise = do
    		tmpdir <- getTemporaryDirectory
    		(tmp, h) <- openTempFile tmpdir "output.tmp"
    		let !endnl = endsNewLine t'
    		let i = InTempFile
    			{ tempFile = tmp
    			, endsInNewLine = endnl
    			}
    		T.hPutStr h t'
    		hClose h
    		return $ OutputBuffer (i : other)
      where
    	!t' = T.concat (mapMaybe getOutput this) <> t
    	!(this, other) = partition isOutput buf
    	isOutput v = case v of
    		Output _ -> True
    		_ -> False
    	getOutput v = case v of
    		Output t'' -> Just t''
    		_ -> Nothing
    addOutputBuffer v (OutputBuffer buf) = return $ OutputBuffer (v:buf)

    -- | Adds a value to the output buffer for later display.
    --
    -- Note that buffering large quantities of data this way will keep it
    -- resident in memory until it can be displayed. While `outputConcurrent`
    -- uses temp files if the buffer gets too big, this STM function cannot do
    -- so.
    bufferOutputSTM :: Outputable v => StdHandle -> v -> STM ()
    bufferOutputSTM h v = bufferOutputSTM' h (OutputBuffer [Output (toOutput v)])

    bufferOutputSTM' :: StdHandle -> OutputBuffer -> STM ()
    bufferOutputSTM' h (OutputBuffer newbuf) = do
    	(OutputBuffer buf) <- takeTMVar bv
    	putTMVar bv (OutputBuffer (newbuf ++ buf))
      where
    	bv = bufferFor h

    -- | A STM action that waits for some buffered output to become
    -- available, and returns it.
    --
    -- The function can select a subset of output when only some is desired;
    -- the fst part is returned and the snd is left in the buffer.
    --
    -- This will prevent it from being displayed in the usual way, so you'll
    -- need to use `emitOutputBuffer` to display it yourself.
    outputBufferWaiterSTM :: (OutputBuffer -> (OutputBuffer, OutputBuffer)) -> STM (StdHandle, OutputBuffer)
    outputBufferWaiterSTM selector = waitgetbuf StdOut `orElse` waitgetbuf StdErr
      where
    	waitgetbuf h = do
    		let bv = bufferFor h
    		(selected, rest) <- selector <$> takeTMVar bv
    		when (selected == OutputBuffer [])
    			retry
    		putTMVar bv rest
    		return (h, selected)

    waitAnyBuffer :: OutputBuffer -> (OutputBuffer, OutputBuffer)
    waitAnyBuffer b = (b, OutputBuffer [])

    -- | Use with `outputBufferWaiterSTM` to make it only return buffered
    -- output that ends with a newline. Anything buffered without a newline
    -- is left in the buffer.
    waitCompleteLines :: OutputBuffer -> (OutputBuffer, OutputBuffer)
    waitCompleteLines (OutputBuffer l) =
    	let (selected, rest) = span completeline l
    	in (OutputBuffer selected, OutputBuffer rest)
      where
    	completeline (v@(InTempFile {})) = endsInNewLine v
    	completeline (Output b) = endsNewLine b

    endsNewLine :: T.Text -> Bool
    endsNewLine t = not (T.null t) && T.last t == '\n'

    -- | Emits the content of the OutputBuffer to the Handle
    --
    -- If you use this, you should use `lockOutput` to ensure you're the only
    -- thread writing to the console.
    emitOutputBuffer :: StdHandle -> OutputBuffer -> IO ()
    emitOutputBuffer stdh (OutputBuffer l) =
    	forM_ (reverse l) $ \ba -> case ba of
    		Output t -> emit t
    		InTempFile tmp _ -> do
    			emit =<< T.readFile tmp
    			void $ tryWhenExists $ removeFile tmp
      where
    	outh = toHandle stdh
    	emit t = void $ tryIO $ do
    		T.hPutStr outh t
    		hFlush outh