{-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances #-}
{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -O2 #-}
{- This module does a lot of calculation that can be expensive, so optimise
 - it well -}

-- | 
-- Copyright: 2015 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
-- > import System.Process
-- > 
-- > main = displayConsoleRegions $ do
-- > 	mapConcurrently download [1..5]
-- >		`concurrently` mapM_ message [1..10]
-- >		`concurrently` createProcessConcurrent (proc "echo" ["hello world"])
-- > 
-- > 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
-- > hello world
-- > 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
-- > hello world
-- > Message 2
-- > Download 1 done!
-- > Message 3
-- > Download 2 ... ...
-- > Download 3 ... ...

module System.Console.Regions (
	-- * Types
	ConsoleRegion,
	RegionLayout(..),
	ToRegionContent(..),
	RegionContent(..),
	LiftRegion(..),
	-- * Initialization
	displayConsoleRegions,
	withConsoleRegion,
	openConsoleRegion,
	newConsoleRegion,
	closeConsoleRegion,
	-- * Region content and display
	setConsoleRegion,
	appendConsoleRegion,
	finishConsoleRegion,
	getConsoleRegion,
	tuneDisplay,
	-- * STM region contents
	--
	-- | The `ToRegionContent` instance for `STM` `Text` can be used to
	-- make regions that automatically update whenever there's
	-- a change to any of the STM values that they use.
	--
	-- For example, a region that displays the screen size,
	-- and automatically refreshes it:
	--
	-- > import qualified Data.Text as T
	--
	-- > r <- openConsoleRegion Linear s
	-- > setConsoleRegion r $ do
	-- > 	w <- readTVar consoleWidth
	-- > 	h <- readTVar consoleHeight
	-- > 	return $ T.pack $ unwords
	-- > 		[ "size:"
	-- >		, show w
	-- > 		, "x"
	-- >		, show h
	-- > 		]
	-- >
	consoleWidth,
	consoleHeight,
	regionList,
	waitDisplayChange,
) where

import Data.Monoid
import Data.String
import Data.Char
import qualified Data.Text as T
import qualified Data.Text.Lazy as L
import qualified Data.Text.IO as T
import Data.Text (Text)
import Control.Monad
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 Text.Read
import Data.List
#ifndef mingw32_HOST_OS
import System.Posix.Signals
import System.Posix.Signals.Exts
#endif
import Control.Applicative
import Prelude

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
-- > ddddeeeefffffff -- [InLine]
-- > fffffggggg.....       (expanded to multiple lines)
-- > 
data RegionLayout = Linear | InLine ConsoleRegion
	deriving (Eq)

-- | A handle allowing access to a region of the console.
newtype ConsoleRegion = ConsoleRegion (TVar R)
	deriving (Eq)

data R = R
	{ regionContent :: RegionContent
	, regionRender :: (Text -> STM Text)
	, regionLayout :: RegionLayout
	, regionChildren :: TVar [ConsoleRegion]
	}

newtype RegionContent = RegionContent (STM Text)

-- | All the regions that are currently displayed on the screen.
--
-- 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.
{-# NOINLINE regionList #-}
regionList :: TMVar [ConsoleRegion]
regionList = unsafePerformIO newEmptyTMVarIO

-- | On Unix systems, this TVar is automatically updated when the
-- terminal is resized. On Windows, it is only initialized on program start
-- with the current terminal size.
{-# NOINLINE consoleSize #-}
consoleSize :: TVar (Console.Window Int)
consoleSize = unsafePerformIO $ newTVarIO $ 
	Console.Window { Console.width = 80, Console.height = 25}

type Width = Int

-- | Gets the width of the console.
--
-- On Unix, this is automatically updated when the terminal is resized.
-- On Windows, it is only initialized on program start.
consoleWidth :: STM Int
consoleWidth = munge . Console.width <$> readTVar consoleSize
  where
#ifndef mingw32_HOST_OS
	munge = id
#else
	-- On Windows, writing to the right-most column caused some
	-- problimatic wrap, so avoid it.
	munge = pred
#endif

-- | Get the height of the console.
consoleHeight :: STM Int
consoleHeight = Console.height <$> readTVar consoleSize

-- | Check if `displayConsoleRegions` is running.
regionDisplayEnabled :: IO Bool
regionDisplayEnabled = atomically $ not <$> isEmptyTMVar regionList

-- | Many actions in this module can be run in either the IO monad
-- or the STM monad. Using STM allows making several changes to the
-- displayed regions atomically, with the display updated a single time.
class LiftRegion m where
	liftRegion :: STM a -> m a

instance LiftRegion STM where
	liftRegion = id

instance LiftRegion IO where
	liftRegion = atomically

-- | Values that can be displayed in a region.
class ToRegionContent v where
	toRegionContent :: v -> RegionContent

instance ToRegionContent String where
	toRegionContent = fromOutput

instance ToRegionContent Text where
	toRegionContent = fromOutput

-- | Note that using a lazy Text in a region will buffer it all in memory.
instance ToRegionContent L.Text where
	toRegionContent = fromOutput

fromOutput :: Outputable v => v -> RegionContent
fromOutput = RegionContent . pure . toOutput

-- | Makes a STM action be run to get the content of a region.
--
-- Any change to the values that action reads will result in an immediate
-- refresh of the display.
instance ToRegionContent (STM Text) where
	toRegionContent = RegionContent

-- | Sets the value of a console region. This will cause the
-- console to be updated to display the new value.
--
-- It's fine for the value to be longer than the terminal is wide,
-- or to include newlines ('\n'). Regions expand to multiple lines as
-- necessary.
--
-- The value can include ANSI SGR escape sequences for changing
-- the colors etc of all or part of a region.
-- 
-- Other ANSI escape sequences, especially those doing cursor
-- movement, will mess up the layouts of regions. Caveat emptor.
setConsoleRegion :: (ToRegionContent v, LiftRegion m) => ConsoleRegion -> v -> m ()
setConsoleRegion r v = liftRegion $
	modifyRegion r $ const $ pure $ toRegionContent v

-- | Appends a value to the current value of a console region.
--
-- > appendConsoleRegion progress "." -- add another dot to progress display
appendConsoleRegion :: (Outputable v, LiftRegion m) => ConsoleRegion -> v -> m ()
appendConsoleRegion r v = liftRegion $
	modifyRegion r $ \(RegionContent a) ->
		return $ RegionContent $ do
			t <- a
			return (t <> toOutput v)

modifyRegion :: ConsoleRegion -> (RegionContent -> STM RegionContent) -> STM ()
modifyRegion (ConsoleRegion tv) f = do
	r <- readTVar tv
	rc <- f (regionContent r)
	let r' = r { regionContent = rc }
	writeTVar tv r'

readRegionContent :: RegionContent -> STM Text
readRegionContent (RegionContent a) = a

resizeRegion :: Width -> ConsoleRegion -> STM [Text]
resizeRegion width (ConsoleRegion tv) = do
	r <- readTVar tv
	ls <- calcRegionLines r width
	return ls

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

-- | Opens a new console region.
openConsoleRegion :: LiftRegion m => RegionLayout -> m ConsoleRegion
openConsoleRegion ly = liftRegion $ do
	h <- newConsoleRegion ly T.empty
	case ly of
		Linear -> do
			ml <- tryTakeTMVar regionList
			case ml of
				Just l -> 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

-- | Makes a new region, but does not add it to the display.
newConsoleRegion :: (LiftRegion m) => ToRegionContent v => RegionLayout -> v -> m ConsoleRegion
newConsoleRegion ly v = liftRegion $ do
	cs <- newTVar mempty
	let r = R
		{ regionContent = RegionContent $ return mempty
		, regionRender = pure
		, regionLayout = ly
		, regionChildren = cs
		}
	h <- ConsoleRegion <$> newTVar r
	displayChildren h
	setConsoleRegion h v
	return h

displayChildren :: ConsoleRegion -> STM ()
displayChildren p@(ConsoleRegion tv) = tuneDisplay p $ \t -> do
	children <- readTVar . regionChildren =<< readTVar tv
	ct <- T.concat <$> mapM getc children
	return $ t <> ct
  where
	getc (ConsoleRegion cv) = do
		c <- readTVar cv
		regionRender c =<< readRegionContent (regionContent c)

-- | Closes a console region. Once closed, the region is removed from the
-- display.
closeConsoleRegion :: LiftRegion m => ConsoleRegion -> m ()
closeConsoleRegion h@(ConsoleRegion tv) = liftRegion $ 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. When Nothing is passed,
-- displays the current value of the console region.
finishConsoleRegion :: (Outputable v, LiftRegion m) => ConsoleRegion -> v -> m ()
finishConsoleRegion h v = liftRegion $ do
	closeConsoleRegion h
	bufferOutputSTM StdOut (toOutput v <> fromString "\n")

-- | Gets the current content of a console region.
getConsoleRegion :: LiftRegion m => ConsoleRegion -> m Text
getConsoleRegion (ConsoleRegion tv) = liftRegion $
	readRegionContent . regionContent =<< readTVar tv

-- | Changes how a console region displays.
--
-- Each time the region's value changes, the STM action is provided
-- with the current value of the region, and returns the value to display.
--
-- For example, this will prevent a region from ever displaying more
-- than 10 characters wide, and will make it display text reversed:
-- 
-- > tuneDisplay myregion $ pure . T.take 10
-- > tuneDisplay myregion $ pure . T.reverse
--
-- Note that repeated calls to tuneDisplay are cumulative.
--
-- Normally, the STM action should avoid retrying, as that would
-- block all display updates.
tuneDisplay :: LiftRegion m => ConsoleRegion -> (Text -> STM Text) -> m ()
tuneDisplay (ConsoleRegion tv) renderer = liftRegion $ do
	r <- readTVar tv
	let rr = \t -> renderer =<< regionRender r t
	let r' = r { regionRender = rr }
	writeTVar tv r'

addChild :: ConsoleRegion -> ConsoleRegion -> STM ()
addChild child _parent@(ConsoleRegion pv) = do
	cv <- regionChildren <$> readTVar pv
	children <- readTVar cv
	let !children' = filter (/= child) children ++ [child]
	writeTVar cv children'

removeChild :: ConsoleRegion -> ConsoleRegion -> STM ()
removeChild child _parent@(ConsoleRegion pv) = do
	cv <- regionChildren <$> readTVar pv
	modifyTVar' cv (filter (/= child))

-- | 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
		when isterm trackConsoleWidth
		da <- async $ displayThread isterm 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 :: IO ()
trackConsoleWidth = do
	let getsz = maybe noop (atomically . writeTVar consoleSize)
		=<< Console.size
	getsz
	installResizeHandler (Just getsz)

data DisplayChange
	= BufferChange BufferSnapshot 
	| RegionChange RegionSnapshot
	| RegionListChange RegionSnapshot
	| TerminalResize Width
	| Shutdown
	| DisplayChangeBarrier Barrier

type BufferSnapshot = (StdHandle, OutputBuffer)
type RegionSnapshot = ([ConsoleRegion], [R], [[Text]])
type Barrier = Integer

-- | This is a broadcast TChan, which gets a DisplayChange written to it
-- after the display has been updated. It can be used to wait for something
-- to be displayed.
{-# NOINLINE displayUpdateNotifier #-}
displayUpdateNotifier :: TChan DisplayChange
displayUpdateNotifier = unsafePerformIO $ newBroadcastTChanIO

{-# NOINLINE displayChangeBarrier #-}
displayChangeBarrier :: TVar Barrier
displayChangeBarrier = unsafePerformIO $ newTVarIO 0

-- | Runs a STM action, and waits for the display to be fully updated
-- before returning.
waitDisplayChange :: STM a -> IO a
waitDisplayChange a = do
	r <- atomically a
	c <- atomically $ dupTChan displayUpdateNotifier
	b <- atomically $ do
		!b <- succ <$> readTVar displayChangeBarrier
		writeTVar displayChangeBarrier b
		return b
	atomically $ waitchange c b
	return r
  where
	waitchange c b = do
		change <- readTChan c
		case change of
			DisplayChangeBarrier b' | b' == b -> return ()
			_ -> waitchange c b

displayThread :: Bool -> TSem -> IO ()
displayThread isterm endsignal = do
	origwidth <- atomically consoleWidth
	origbarrier <- atomically (readTVar displayChangeBarrier)
	go ([], [], []) origwidth origbarrier
  where
	go origsnapshot@(orighandles, origregions, origlines) origwidth origbarrier = do
		let waitwidthchange = do
			w <- consoleWidth
			if w == origwidth then retry else return w
		let waitbarrierchange = do
			b <- readTVar displayChangeBarrier
			if b /= origbarrier
				then return b
				else retry
		let waitanychange =
			(RegionChange <$> regionWaiter origsnapshot origwidth)
				`orElse`
			(RegionListChange <$> regionListWaiter origsnapshot)
				`orElse`
			(BufferChange <$> outputBufferWaiterSTM waitCompleteLines)
				`orElse`
			(TerminalResize <$> waitwidthchange)
				`orElse`
			(waitTSem endsignal >> pure Shutdown)
				`orElse`
			-- Must come last, so the changes above are
			-- processed before barriers.
			(DisplayChangeBarrier <$> waitbarrierchange)
		(change, height) <- atomically $ (,)
			<$> waitanychange
			<*> consoleHeight
		let onscreen = take (height - 1) . concat
		let update snapshot@(_, _, newlines) = do
			when isterm $
				changedLines (onscreen origlines) (onscreen newlines)
			return $ go snapshot origwidth origbarrier
		next <- case change of
			RegionChange snapshot -> update snapshot
			RegionListChange snapshot -> update snapshot
			BufferChange (h, buf) -> do
				-- Note that even when every available line
				-- is dedicated to visible regions, the
				-- buffer is still displayed. It would be
				-- more efficient to not display it, but
				-- this makes it available in scroll back.
				let origlines' = onscreen origlines
				inAreaAbove isterm (length origlines') origlines' $
					emitOutputBuffer h buf
				return $ go origsnapshot origwidth origbarrier
			TerminalResize newwidth -> do
				newlines <- atomically (mapM (resizeRegion newwidth) orighandles)
				when isterm $ do
					resizeRecovery (onscreen newlines)
				return $ go (orighandles, origregions, newlines) newwidth origbarrier
			Shutdown ->
				return $ return ()
			DisplayChangeBarrier b ->
				return $ go origsnapshot origwidth b
		hFlush stdout
		atomically $ writeTChan displayUpdateNotifier change
		next

readRegions :: [ConsoleRegion] -> STM [R]
readRegions = mapM (\(ConsoleRegion h) -> readTVar h)

-- | Wait for any changes to the region list, eg adding or removing a region.
regionListWaiter :: RegionSnapshot -> STM RegionSnapshot
regionListWaiter (orighandles, _origregions, origlines) = do
	handles <- readTMVar regionList
	if handles == orighandles
		then retry
		else do
			rs <- readRegions handles
			return (handles, rs, origlines)

-- Wait for any changes to any of the contents of regions currently in the
-- region list.
regionWaiter :: RegionSnapshot -> Width -> STM RegionSnapshot
regionWaiter (orighandles, _origregions, origlines) width = do
	rs <- readRegions orighandles
	newlines <- mapM getr rs
	unless (newlines /= origlines)
		retry
	return (orighandles, rs, newlines)
  where
	getr r = calcRegionLines r width

-- This is not an optimal screen update like curses can do, but it's
-- pretty efficient, most of the time!
changedLines :: [Text] -> [Text] -> IO ()
changedLines origlines newlines
	| delta == 0 = do
		-- The total number of lines is unchanged, so update
		-- whichever ones have changed, and leave the rest as-is.
		diffUpdate origlines newlines
	| delta > 0 = do
		-- Added more lines, so output each, with a
		-- newline, thus scrolling the old lines up
		-- the screen. (We can do this, because the cursor
		-- is left below the first line.)
		let addedlines = reverse (take delta newlines)
		displayLines addedlines
		-- Some existing lines may have also changed..
		let scrolledlines = addedlines ++ origlines
		diffUpdate scrolledlines newlines
	| otherwise = do
		-- Some lines were removed. Move up that many lines,
		-- clearing each line, and update any changed lines.
		replicateM_ (abs delta) $ do
			setCursorColumn 0
			cursorUp 1
			clearLine
		diffUpdate (drop (abs delta) origlines) newlines
  where
	delta = length newlines - length origlines

diffUpdate :: [Text] -> [Text] -> IO ()
diffUpdate old new = updateLines (zip (zip new changed) old)
  where
	changed = map (uncurry (/=)) (zip new old) ++ repeat True

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

-- Displays lines that are paired with True, and skips over the rest.
-- Cursor is assumed to be just below the first line at the
-- beginning, and is put back there at the end.
updateLines :: [((Text, Bool), Text)] -> IO ()
updateLines l
	| null l' = noop
	| otherwise = do
		forM_ l' $ \((newt, offset), oldt) -> do
			setCursorColumn 0
			cursorUp offset
#ifndef mingw32_HOST_OS
			T.hPutStr stdout $
				genLineUpdate $ calcLineUpdate oldt newt
#else
			-- Windows does not support ansi characters
			-- emitted in a string, so do a full line
			-- redraw.
			T.hPutStr stdout newt
			clearFromCursorToLineEnd
#endif
		cursorDown (sum (map (snd . fst) l'))
		setCursorColumn 0
  where
	l' = changeOffsets l 1 []

-- Recover from a resize by redrawing all region lines.
--
-- The resize can change the position of the cursor, which would garble
-- the display going forward. To fix, the cursor is moved to the top of
-- the screen, which is cleared, and all regions are redrawn from there.
resizeRecovery :: [Text] -> IO ()
resizeRecovery newlines = do
	setCursorPosition 0 0
	inAreaAbove True 0 newlines $
		return ()

-- Move cursor up before the lines, performs some output there,
-- which will scroll down and overwrite the lines, so 
-- redraws all the lines below.
inAreaAbove :: Bool -> Int -> [Text] -> IO () -> IO ()
inAreaAbove isterm numlines ls outputter = do
	when isterm $ do
		unless (numlines < 1) $ do
			setCursorColumn 0
			cursorUp $ numlines
		clearFromCursorToScreenEnd
	-- Flush stdout now, because the outputter may write to stderr, so
	-- the cursor needs to be moved first.
	hFlush stdout
	outputter
	when isterm $ do
		setCursorColumn 0 -- just in case the output lacked a newline
		displayLines (reverse ls)

displayLines :: [Text] -> IO ()
displayLines = mapM_ $ \l -> do
	T.hPutStr stdout l
	putChar '\n'

installResizeHandler :: Maybe (IO ()) -> IO ()
#ifndef mingw32_HOST_OS
installResizeHandler h = void $
	installHandler windowChange (maybe Default Catch h) Nothing
#else
installResizeHandler _ = return ()
#endif

calcRegionLines :: R -> Width -> STM [Text]
calcRegionLines r width = do
	t <- regionRender r =<< readRegionContent (regionContent r)
	return $ reverse $ calcLines t width

-- | Splits a Text into the lines it would display using when output onto
-- a console with a given width, starting from the first column.
--
-- ANSI SGR sequences are handled specially, so that color, etc settings
-- work despite the lines being split up, and the lines can be output
-- indepedently. For example, "foooREDbar bazRESET" when split into lines
-- becomes ["fooREDbarRESET", "RED bazRESET"]
calcLines :: Text -> Width -> [Text]
calcLines t width
	| width < 1 || T.null t = [t] -- even an empty text is 1 line high
	| otherwise = calcLines' width [] [] 0 1 (T.length t) t

calcLines' :: Int -> [Text] -> [Text] -> Int -> Int -> Int -> Text -> [Text]
calcLines' width collectedlines collectedSGR i displaysize len t
	| i >= len = if i > 0
		then reverse (finishline t)
		else reverse collectedlines
	| t1 == '\n' = calcLines' width (finishline $ T.init currline)
		[] 0 1 (T.length rest) (contSGR rest)
	-- ANSI escape sequences do not take up space on screen.
	| t1 == '\ESC' && i+1 < len = case T.index t (i+1) of
		'[' -> skipansi endCSI True
		']' -> skipansi endOSC False
		_ -> calcLines' width collectedlines collectedSGR (i+1) displaysize len t
	-- Control characters do not take up space on screen.
	| isControl t1 = calcLines' width collectedlines collectedSGR (i+1) displaysize len t
	| displaysize >= width = calcLines' width (finishline currline)
		[] 0 1 (T.length rest) (contSGR rest)
	| otherwise = calcLines' width collectedlines collectedSGR (i+1) (displaysize+1) len t
  where
	t1 = T.index t i
	(currline, rest) = T.splitAt (i+1) t

	skipansi toend isCSI = case T.findIndex toend (T.drop (i+2) t) of
		Just csiend -> calcLines' width collectedlines 
			(addSGR (csiend+2)) (i+2+csiend) (displaysize-1) len t
		Nothing -> reverse (finishline t)
	  where
		addSGR csiend
			| not isCSI = collectedSGR
			| ansicode == resetSGR = []
			| not (T.null ansicode) && T.last ansicode == endSGR =
				ansicode : collectedSGR
			| otherwise = collectedSGR
		  where
			ansicode = T.take (csiend + 1) (T.drop i t)
	finishline l = closeSGR l : collectedlines
	-- Close any open SGR codes at end of line
	closeSGR l
		| null collectedSGR = l
		| otherwise = l <> resetSGR
	-- Continue any open SGR codes from previous line
	contSGR l = mconcat (reverse collectedSGR) <> l

resetSGR :: Text
resetSGR = T.pack (setSGRCode [Reset])

endCSI :: Char -> Bool
endCSI c = let o = ord c in o >= 64 && o < 127

endOSC :: Char -> Bool
endOSC c = c == '\BEL'

endSGR :: Char
endSGR = 'm'

-- | Finds the least expensive output to make a console that was displaying
-- the old line display the new line. Cursor starts at far left.
--
-- Basically, loop through and find spans where the old and new line are
-- the same. Generate cursorForwardCode ANSI sequences to skip over those
-- spans, unless such a sequence would be longer than the span it's skipping.
--
-- Since ANSI sequences can be present in the line, need to take them
-- into account. Generally, each of the sequences in new has to be included,
-- even if old contained the same sequence:
--
-- > old: GREENfoofoofooREDbarbarbarRESETbaz
-- > new: GREENfoofoofooREDxarbarbaxRESETbaz
-- > ret: GREEN-------->REDx------>yRESET
--
-- (The first GREEN does not effect any output text, so it can be elided.)
-- 
-- Also, despite old having the same second span as new, in the same
-- location, that span has to be re-emitted because its color changed:
-- 
-- > old: GREENfoofooREDbarbarbarbarbar
-- > new: GREENfoofoofooTANbarbarbar
-- > ret: GREEN----->fooTANbarbarbarCLEARREST
--
-- Also note above that the sequence has to clear the rest of the line,
-- since the new line is shorter than the old.
calcLineUpdate :: Text -> Text -> [LineUpdate]
calcLineUpdate old new = 
	reverse $ go
		(advanceLine old [] [])
		(advanceLine new [] [])
  where
	go (Just _, _, _, _) (Nothing, _, past, _) = ClearToEnd : past
	go (Nothing, _, _, _) (Nothing, _, past, _) = past
	go (Nothing, _, _, _) (Just n, ns, past, _) =
		Display ns : Display (T.singleton n) : past
	go (Just o, os, _, oinvis) (Just n, ns, past, ninvis)
		| o == n && oinvis == ninvis = go
			(advanceLine os [] oinvis)
			(advanceLine ns (Skip [o] : past) ninvis)
		| otherwise = go
			(advanceLine os [] oinvis)
			(advanceLine ns (Display (T.singleton n) : past) ninvis)

type Past = [LineUpdate]
type Invis = [LineUpdate]

-- Find next character of t that is not a ANSI escape sequence
-- or control char. Any such passed on the way to the character
-- are prepended to past, and added to invis.
--
-- resetSGR is handled specially; it causes all SGRs to be removed from
-- invis, It's still prepended to past.
advanceLine :: Text -> Past -> Invis -> (Maybe Char, Text, Past, Invis)
advanceLine t past invis
	| T.null t = (Nothing, T.empty, past, invis)
	| otherwise = case T.head t of
		'\ESC' -> case T.drop 1 t of
			t' | T.null t' -> advanceLine (T.drop 1 t)
				(Skip "\ESC":past) (Skip "\ESC":invis)
			   | otherwise -> case T.head t' of
			   	'[' -> skipansi endCSI
				']' -> skipansi endOSC
				c -> (Just c, T.drop 2 t, Skip "\ESC":past, Skip "\ESC":invis)
		c | isControl c -> advanceLine (T.drop 1 t) (Skip [c]:past) (Skip [c]:invis)
		  | otherwise -> (Just c, T.drop 1 t, past, invis)
  where
	skipansi toend = case T.findIndex toend (T.drop 2 t) of
		Just csiend -> 
			let sgr = SGR (T.take (csiend+3) t)
			in advanceLine (T.drop (csiend+3) t)
				(sgr:past) (addsgr sgr invis)
		Nothing -> (Nothing, T.empty, past, invis)
	addsgr (SGR sgrt) l
		| sgrt == resetSGR = filter (not . isSGR) l
	addsgr s l = s:l

data LineUpdate = Display Text | Skip [Char] | SGR Text | ClearToEnd
	deriving (Eq, Show)

isSGR :: LineUpdate -> Bool
isSGR (SGR _) = True
isSGR _ = False

genLineUpdate :: [LineUpdate] -> Text
genLineUpdate l = T.concat $ map tot (optimiseLineUpdate l)
  where
	tot (Display t) = t
	tot (Skip s)
		-- length (cursorForwardCode 1) == 4 so there's no point
		-- generating that for a skip of less than 5.
		| len < 5 = T.pack s
		| otherwise = T.pack (cursorForwardCode len)
	  where
		len = length s
	tot (SGR t) = t
	tot ClearToEnd = T.pack clearFromCursorToLineEndCode

optimiseLineUpdate :: [LineUpdate] -> [LineUpdate]
optimiseLineUpdate = go []
  where
	-- elide trailing Skips
	go (Skip _:rest) [] = go rest []
	-- elide SGRs at the end of the line, except for the reset SGR
	go (SGR t:rest) [] | t /= resetSGR = go rest []
	go c [] = reverse c
	-- combine adjacent SGRs and Skips
	go c (SGR t1:Skip s:SGR t2:rest) = tryharder c (SGR (combineSGR t1 t2):Skip s:rest)
	go c (Skip s:Skip s':rest) = tryharder c (Skip (s++s'):rest)
	go c (SGR t1:SGR t2:rest) = tryharder c (SGR (combineSGR t1 t2):rest)
	go c (v:rest) = go (v:c) rest
	tryharder c l = go [] (reverse c ++ l)

-- Parse and combine 2 ANSI SGR sequences into one.
combineSGR :: Text -> Text -> Text
combineSGR a b = case combineSGRCodes (codes a) (codes b) of
	Nothing -> a <> b
	Just cs -> T.pack $ "\ESC[" ++ intercalate ";" (map show cs) ++ "m"
  where
	codes = map (readMaybe . T.unpack) .
		T.split (== ';') . T.drop 2 . T.init

-- Prefers values from the second sequence when there's a conflict with
-- values from the first sequence.
combineSGRCodes :: [Maybe Int] -> [Maybe Int] -> Maybe [Int]
combineSGRCodes as bs =
	map snd . nubBy (\a b -> fst a == fst b) <$> mapM range (reverse bs ++ reverse as)
  where
	range Nothing = Nothing
	range (Just x)
		| x >= 30 && x <= 37 = Just (Foreground, x)
		| x >= 40 && x <= 47 = Just (Background, x)
		| x >= 90 && x <= 97 = Just (Foreground, x)
		| x >= 100 && x <= 107 = Just (Background, x)
		| otherwise = Nothing