{-# LANGUAGE BangPatterns, TupleSections #-}

-- | 
-- Copyright: 2013 Joey Hess <id@joeyh.name>
-- License: BSD-2-clause
-- 
-- Console regions are displayed near the bottom of the console, and can be
-- updated concurrently by threads. Any other output displayed using
-- `outputConcurrent` and `createProcessConcurrent`
-- will scroll up above the open console regions.
--
-- For example, this program:
--
-- > import Control.Concurrent.Async
-- > import Control.Concurrent
-- > import System.Console.Concurrent
-- > import System.Console.Regions
-- > 
-- > main = displayConsoleRegions $ do
-- > 	mapConcurrently download [1..5] `concurrently` mapM_ message [1..10]
-- > 
-- > message :: Int -> IO ()
-- > message n = do
-- > 	threadDelay 500000
-- > 	outputConcurrent ("Message " ++ show n ++ "\n")
-- > 
-- > download :: Int -> IO ()
-- > download n = withConsoleRegion Linear $ \r -> do
-- > 	setConsoleRegion r basemsg
-- > 	go n r
-- >   where
-- > 	basemsg = "Download " ++ show n
-- >	go c r
-- >		| c < 1 = finishConsoleRegion r (basemsg ++ " done!")
-- > 		| otherwise = do
-- > 			threadDelay 1000000
-- > 			appendConsoleRegion r " ... "
-- > 			go (c-1) r
--
-- Will display like this:
--
-- > Message 1
-- > Message 2
-- > Download 1 ...
-- > Download 2 ...
-- > Download 3 ...
--
-- Once the 1st download has finished, and another message has displayed,
-- the console will update like this:
--
-- > Message 1
-- > Message 2
-- > Download 1 done!
-- > Message 3
-- > Download 2 ... ...
-- > Download 3 ... ...

module System.Console.Regions (
	-- * Initialization
	displayConsoleRegions,
	ConsoleRegionHandle,
	RegionLayout(..),
	withConsoleRegion,
	openConsoleRegion,
	closeConsoleRegion,
	-- * Output
	setConsoleRegion,
	appendConsoleRegion,
	finishConsoleRegion,
	-- * STM interface
	--
	-- | These actions can be composed into a STM transaction; 
	-- once the transaction completes the console will be updated
	-- a single time to reflect all the changes made.
	openConsoleRegionSTM,
	closeConsoleRegionSTM,
	setConsoleRegionSTM,
	appendConsoleRegionSTM,
	updateRegionListSTM,
) where

import Data.Monoid
import Data.Maybe
import Data.String
import Data.Char
import qualified Data.ByteString as B
import Control.Monad
import Control.Applicative
import Control.Monad.IO.Class (liftIO, MonadIO)
import Control.Concurrent.STM
import Control.Concurrent.STM.TSem
import Control.Concurrent.Async
import System.Console.ANSI
import qualified System.Console.Terminal.Size as Console
import System.IO
import System.IO.Unsafe (unsafePerformIO)
import System.Posix.Signals
import System.Posix.Signals.Exts

import System.Console.Concurrent
import Utility.Monad
import Utility.Exception

-- | Controls how a region is laid out in the console.
--
-- Here's an annotated example of how the console layout works.
--
-- > scrolling......
-- > scrolling......
-- > scrolling......
-- > aaaaaa......... -- Linear
-- > bbbbbbbbbbbbbbb -- Linear
-- > bbb............       (expanded to multiple lines)
-- > ccccccccc...... -- Linear
-- > ddd eee fffffff -- [InLine]
-- > ffff ggggg.....       (expanded to multiple lines)
-- > 
data RegionLayout = Linear | InLine ConsoleRegionHandle
	deriving (Eq)

newtype ConsoleRegionHandle = ConsoleRegionHandle (TVar Region)
	deriving (Eq)

type Width = Int
type Height = Int

data Region = Region
	{ regionContent :: B.ByteString
	, regionHeight :: (Width -> Height)
	-- ^ Function from console width to the height of this region
	, regionLayout :: RegionLayout
	, regionChildren :: Maybe [ConsoleRegionHandle]
	}

instance Eq Region where
	a == b = regionContent a == regionContent b
		&& regionLayout a == regionLayout b

type RegionList = TMVar [ConsoleRegionHandle]

-- | A shared global list of regions.
{-# NOINLINE regionList #-}
regionList :: RegionList
regionList = unsafePerformIO newEmptyTMVarIO

-- | Updates the list of regions. The list is ordered from the bottom of
-- the screen up. Reordering it will change the order in which regions are
-- displayed. It's also fine to remove, duplicate, or add new regions to the
-- list.
updateRegionListSTM :: ([ConsoleRegionHandle] -> [ConsoleRegionHandle]) -> STM ()
updateRegionListSTM f = 
	maybe noop (putTMVar regionList . f) =<< tryTakeTMVar regionList

-- | The RegionList TMVar is left empty when `displayConsoleRegions`
-- is not running.
regionDisplayEnabled :: IO Bool
regionDisplayEnabled = atomically $ not <$> isEmptyTMVar regionList

-- | Sets the value to display within a console region.
{- TODO
-- It's fine for the value to be longer than the terminal is wide.
-- Regions are laid out according to the size of their contents,
-- and expand to multiple lines if necessary.
--
-- The value of a `Linear` region can contain newlines ('\n').
-- And it's ok to include ANSI escape sequences for changing colors,
-- or setting the terminal title. However, ANSI cursor movement sequences
-- will mess up the  layouts of regions, '\r' is not handled, and
-- other control characters or ANSI codes may confuse the region
-- display. Caveat emptor.
-}
setConsoleRegion :: Outputable v => ConsoleRegionHandle -> v -> IO ()
setConsoleRegion h = atomically . setConsoleRegionSTM h

setConsoleRegionSTM :: Outputable v => ConsoleRegionHandle -> v -> STM ()
setConsoleRegionSTM (ConsoleRegionHandle tv) v = do
	r <- readTVar tv
	writeTVar tv (modifyRegion r (const (toOutput v)))
	case regionLayout r of
		Linear -> return ()
		InLine p -> refreshParent p

-- | Appends the value to whatever was already on display within a console
-- region.
appendConsoleRegion :: Outputable v => ConsoleRegionHandle -> v -> IO ()
appendConsoleRegion h = atomically . appendConsoleRegionSTM h

appendConsoleRegionSTM :: Outputable v => ConsoleRegionHandle -> v -> STM ()
appendConsoleRegionSTM (ConsoleRegionHandle tv) v = do
	r <- readTVar tv
	writeTVar tv (modifyRegion r (<> toOutput v))
	case regionLayout r of
		Linear -> return ()
		InLine p -> refreshParent p

modifyRegion :: Region -> (B.ByteString -> B.ByteString) -> Region
modifyRegion r f = r { regionContent = c, regionHeight = calcRegionHeight c }
  where
	!c = f (regionContent r)

-- | Runs the action with a new console region, closing the region when
-- the action finishes or on exception.
withConsoleRegion :: (MonadIO m, MonadMask m) => RegionLayout -> (ConsoleRegionHandle -> m a) -> m a
withConsoleRegion l = bracketIO (openConsoleRegion l) closeConsoleRegion

-- | Opens a new console region for output.
openConsoleRegion :: RegionLayout -> IO ConsoleRegionHandle
openConsoleRegion = atomically . openConsoleRegionSTM

-- | STM version of `openConsoleRegion`. Allows atomically opening multiple
-- regions at the same time, which guarantees they are on adjacent lines.
--
-- > [r1, r2, r3] <- atomically $
-- >	replicateM 3 (openConsoleRegionSTM Linear)
openConsoleRegionSTM :: RegionLayout -> STM ConsoleRegionHandle
openConsoleRegionSTM ly = do
	let r = Region
		{ regionContent = mempty
		, regionHeight = calcRegionHeight mempty
		, regionLayout = ly
		, regionChildren = Nothing
		}
	h <- ConsoleRegionHandle <$> newTVar r
	case ly of
		Linear -> do
			v <- tryTakeTMVar regionList
			case v of
				Just l -> do putTMVar regionList (h:l)
				-- displayConsoleRegions is not active, so
				-- it's not put on any list, and won't display
				Nothing -> return ()
		InLine parent -> addChild h parent

	return h

-- | Closes a console region. Once closed, the region is removed from the
-- display.
closeConsoleRegion :: ConsoleRegionHandle -> IO ()
closeConsoleRegion = atomically . closeConsoleRegionSTM

closeConsoleRegionSTM :: ConsoleRegionHandle -> STM ()
closeConsoleRegionSTM h@(ConsoleRegionHandle tv) = do
	v <- tryTakeTMVar regionList
	case v of
		Just l ->
			let !l' = filter (/= h) l
			in putTMVar regionList l'
		_ -> return ()
	ly <- regionLayout <$> readTVar tv
	case ly of
		Linear -> return ()
		InLine parent -> removeChild h parent

-- | Closes the console region and displays the passed value in the
-- scrolling area above the active console regions.
finishConsoleRegion :: Outputable v => ConsoleRegionHandle -> v -> IO ()
finishConsoleRegion h = atomically . finishConsoleRegionSTM h

finishConsoleRegionSTM :: Outputable v => ConsoleRegionHandle -> v -> STM ()
finishConsoleRegionSTM h v = do
	closeConsoleRegionSTM h
	bufferOutputSTM StdOut (toOutput v <> fromString "\n")

removeChild :: ConsoleRegionHandle -> ConsoleRegionHandle -> STM ()
removeChild child parent@(ConsoleRegionHandle pv) = do
	modifyTVar' pv $ \p -> case regionChildren p of
		Nothing -> p
		Just l -> p { regionChildren = Just $ filter (/= child) l }
	refreshParent parent

addChild :: ConsoleRegionHandle -> ConsoleRegionHandle -> STM ()
addChild child parent@(ConsoleRegionHandle pv) = do
	modifyTVar' pv $ \p -> p
		{ regionChildren = Just $ child : filter (/= child) (fromMaybe [] (regionChildren p)) }
	refreshParent parent

refreshParent :: ConsoleRegionHandle -> STM ()
refreshParent parent@(ConsoleRegionHandle pv) = do
	p <- readTVar pv
	case regionChildren p of
		Nothing -> return ()
		Just l -> do
			cs <- forM l $ \child@(ConsoleRegionHandle cv) -> do
				refreshParent child
				regionContent <$> readTVar cv
			let p' = p { regionContent = mconcat cs }
			writeTVar pv p'

-- | Handles all display for the other functions in this module.
--
-- Note that this uses `lockOutput`, so it takes over all output to the
-- console while the passed IO action is running. As well as displaying
-- the console regions, this handles display of anything buffered by
-- `outputConcurrent` and `createProcessConcurrent`.
--
-- When standard output is not an ANSI capable terminal,
-- console regions are not displayed.
displayConsoleRegions :: (MonadIO m, MonadMask m) => m a -> m a
displayConsoleRegions a = ifM (liftIO regionDisplayEnabled)
	( a -- displayConsoleRegions is already running
	, lockOutput $ bracket setup cleanup (const a)
	)
  where
	setup = liftIO $ do
		atomically $ putTMVar regionList []
		endsignal <- atomically $ do
			s <- newTSem 1
			waitTSem s
			return s
		isterm <- liftIO $ hSupportsANSI stdout
		cwidth <- atomically newEmptyTMVar
		when isterm $
			trackConsoleWidth cwidth
		da <- async $ displayThread isterm cwidth endsignal
		return (isterm, da, endsignal)
	cleanup (isterm, da, endsignal) = liftIO $ do
		atomically $ signalTSem endsignal
		void $ wait da
		void $ atomically $ takeTMVar regionList
		when isterm $
			installResizeHandler Nothing

trackConsoleWidth :: TMVar Width -> IO ()
trackConsoleWidth cwidth = do
	let getwidth = do
		v <- Console.size
		case v of
			Nothing -> return ()
			Just (Console.Window _height width) ->
				atomically $ void $ do
					void $ tryTakeTMVar cwidth
					putTMVar cwidth width
	getwidth
	installResizeHandler $ Just getwidth

data DisplayChange
	= BufferChange [(StdHandle, OutputBuffer)]
	| RegionChange RegionSnapshot
	| TerminalResize (Maybe Width)
	| EndSignal ()

type RegionSnapshot = ([ConsoleRegionHandle], [Region], Maybe Width)

displayThread :: Bool -> TMVar Width -> TSem -> IO ()
displayThread isterm cwidth endsignal = do
	origwidth <- atomically $ tryReadTMVar cwidth
	go ([], [], origwidth)
  where
	go origsnapshot@(orighandles, origregions, origwidth) = do
		let waitwidthchange = do
			v <- tryReadTMVar cwidth
			if v == origwidth then retry else return v
		change <- atomically $
			(RegionChange <$> regionWaiter origsnapshot)
				`orElse`
			(RegionChange <$> regionListWaiter origsnapshot)
				`orElse`
			(BufferChange <$> outputBufferWaiterSTM waitCompleteLines)
				`orElse`
			(TerminalResize <$> waitwidthchange)
				`orElse`
			(EndSignal <$> waitTSem endsignal)
		case change of
			RegionChange snapshot@(_, regions, _) -> do
				when isterm $
					changedRegions origregions regions
				go snapshot
			BufferChange buffers -> do
				inAreaAbove isterm origregions $
					mapM_ (uncurry emitOutputBuffer) buffers
				go origsnapshot
			TerminalResize width -> do
				when isterm $
					-- force redraw of all regions
					inAreaAbove isterm origregions $
						return ()
				go (orighandles, origregions, width)
			EndSignal () -> return ()

readRegions :: [ConsoleRegionHandle] -> STM [Region]
readRegions = mapM (\(ConsoleRegionHandle h) -> readTVar h)

-- Wait for any changes to the region list, eg adding or removing a handle.
regionListWaiter :: RegionSnapshot -> STM RegionSnapshot
regionListWaiter (orighandles, _origregions, origwidth) = do
	handles <- readTMVar regionList
	if handles == orighandles
		then retry
		else (handles,,origwidth) <$> readRegions handles
		
-- Wait for any changes to any of the regions currently in the region list.
regionWaiter :: RegionSnapshot -> STM RegionSnapshot
regionWaiter (orighandles, origregions, origwidth) = do
	rs <- readRegions orighandles
	if rs == origregions
		then retry
		else return (orighandles, rs, origwidth)

-- This is not an optimal screen update like curses can do, but it's
-- pretty efficient, most of the time! The only particularly
-- expensive part is removing a region, which typically reorders
-- the regions and so requires redrawing them all.
changedRegions :: [Region] -> [Region] -> IO ()
changedRegions origregions regions
	| delta == 0 = do
		-- The total number of regions is unchanged, so update
		-- whichever ones have changed, and leave the rest as-is.
		diffUpdate origregions regions
	| delta > 0 = do
		-- Added more regions, so output each, with a
		-- newline, thus scrolling the old regions up
		-- the screen. (We can do this, because the cursor
		-- is left below the first region.)
		let newregions = reverse (take delta regions)
		displayRegions newregions
		hFlush stdout
		-- Some existing regions may have also changed..
		let scrolledregions = newregions ++ origregions
		diffUpdate scrolledregions regions
	| otherwise = do
		-- Some regions were removed. Move up that many lines,
		-- clearing each line, and update any changed regions.
		replicateM_ (abs delta) $ do
			cursorUpLine 1
			clearLine
		diffUpdate (drop (abs delta) origregions) regions
  where
	delta = length regions - length origregions

-- TODO Rather than writing the whole text of a region, find
-- a more efficient update, reusing parts of the old content of the region.
diffUpdate :: [Region] -> [Region] -> IO ()
diffUpdate origregions regions = updateRegions (zip regions changed)
  where
	changed = map (uncurry (/=)) (zip regions origregions) ++ repeat True

-- Displays regions that are paired with True, and skips over the rest.
-- Cursor is assumed to be just below the line of the first region at the
-- beginning, and is put back there at the end.
updateRegions :: [(Region, Bool)] -> IO ()
updateRegions l
	| null l' = noop
	| otherwise = do
		forM_ l' $ \(r, offset) -> do
			cursorUpLine offset
			clearLine
			B.hPut stdout (regionContent r)
			hFlush stdout
		cursorDownLine (sum (map snd l'))
		setCursorColumn 0
		hFlush stdout
  where
	l' = changeOffsets l 1 []

changeOffsets :: [(r, Bool)] -> Int -> [(r, Int)] -> [(r, Int)]
changeOffsets [] _ c = reverse c
changeOffsets ((r, changed):rs) n c
	| changed = changeOffsets rs 1 ((r, n):c)
	| otherwise = changeOffsets rs (succ n) c

-- Move cursor up before the regions, performs some output there,
-- which will scroll down and overwrite the regions, so 
-- redraws all the regions below.
inAreaAbove :: Bool -> [Region] -> IO () -> IO ()
inAreaAbove isterm regions outputter = do
	when isterm $ do
		unless (null regions) $
			cursorUpLine (length regions)
		clearFromCursorToScreenEnd
	outputter
	when isterm $ do
		setCursorColumn 0 -- just in case the output lacked a newline
		displayRegions (reverse regions)
	hFlush stdout

displayRegions :: [Region] -> IO ()
displayRegions = mapM_ $ \r -> do
	B.hPut stdout (regionContent r)
	putChar '\n'

installResizeHandler :: Maybe (IO ()) -> IO ()
installResizeHandler h = void $
	installHandler windowChange (maybe Default Catch h) Nothing

calcRegionHeight :: B.ByteString -> Width -> Height
calcRegionHeight b width = sum $ map (calcLineHeight width) (B.split wnl b)
  where
	wnl = fromIntegral (ord '\n')

-- Calculate the height a line would occupy if output onto a console
-- with a given width, starting from the first column.
--
-- Note that this counts each byte, is not aware of multibyte characters,
-- so will over-estimate a bit for those.
--
-- Also an issue are ANSI escape sequences for eg, setting colors or the
-- title. This is handled by stripping out common ANSI escape sequences
-- and control characters.
calcLineHeight :: Width -> B.ByteString -> Height
calcLineHeight width b
	| width < 1 || B.null b = 1 -- even an empty line is 1 line high
	| otherwise =
		let (q,r) = (B.length b - countInvisibleBytes b) `quotRem` width
		in q + if r > 0 then 1 else 0

-- ANSI sequences and control characters.
countInvisibleBytes :: B.ByteString -> Int
countInvisibleBytes = go 0 . breakesc
  where
	go c (beforeesc, b)
		| B.length b <= 1 = B.length b + c'
		| headis csi (B.drop 1 b) = countseq breakcsi
		| headis osc (B.drop 1 b) = countseq breakosc
		-- found an ESC but apparently not in an ANSI sequence;
		-- count it as 1 invisible control character
		| otherwise = go (c'+1) (breakesc (B.drop 1 b))
	  where
		c' = c + countControlChars beforeesc
		countseq breaker =
			let (inseq, b') = breaker (B.drop 2 b)
			-- add 1 for the ESC and one for the head char
			-- that introduced the sequence, plus the length
			-- of the rest of the sequence, plus 1 for 
			-- the end char of the sequence.
			in go (c'+1+1+B.length inseq+1) (breakesc (B.drop 1 b'))

	esc = fromIntegral (ord '\ESC')
	bel = fromIntegral (ord '\BEL')
	breakesc = B.break (== esc)
	headis c b = B.head b ==c
	csi = fromIntegral (ord '[') -- Control Sequence Introducer
	csiend c = c >= 64 && c < 127
	breakcsi = B.break csiend
	osc = fromIntegral (ord ']') -- Operating system command
	breakosc = B.break (== bel)

countControlChars :: B.ByteString -> Int
countControlChars = length . filter iscontrol8 . B.unpack
  where
	iscontrol8 c = c < 32 || c == 127 -- del