DiffRender.hs

module DiffRender.DiffRender where

import Common
import Control.Monad.IO.Class
import Data.List as DL
import Data.Text as T
import Data.Text.IO as T
import Data.Vector.Mutable (IOVector)
import qualified Data.Vector.Mutable as MV
import System.Console.ANSI (Color(..), ColorIntensity(..), ConsoleLayer(..),
                            SGR(..), Underlining(..), setSGRCode)
import qualified System.Console.ANSI as A
import qualified System.IO as S
import Test.Common

data DiffRender = DiffRender
  { dfScreenStateBack :: ScreenState
  , dfScreenState     :: ScreenState
  }

emptyDiffRender :: Dimensions -> IO DiffRender
emptyDiffRender Dimensions { diH = rows, diW = cols } = do
  ss <- emptyScreenState rows cols
  ssBack <- emptyScreenState rows cols
  pure $ DiffRender { dfScreenStateBack = ssBack, dfScreenState = ss }

emptyScreenState :: Int -> Int -> IO ScreenState
emptyScreenState rows cols = do
  stLines <- MV.generate rows (\_ -> [Plain (T.replicate cols " ")])
  pure (ScreenState stLines (ScreenPos 0 0) cols True)

dfSetCursorPosition :: Int -> Int -> DiffRender -> DiffRender
dfSetCursorPosition x y dfr  =
  let
    screenState = dfScreenStateBack dfr
    screenLines = ssLines screenState
    screenColumns = ssColumns screenState
  in if (x >= 0 && x < screenColumns) && (y >= 0 && y < (MV.length screenLines))
    then dfr { dfScreenStateBack = screenState { ssCursorOverflow = False, ssCursorPos = ScreenPos x y }}
    else dfr { dfScreenStateBack = screenState { ssCursorOverflow = True }}

dfPutText :: MonadIO m => StyledText -> DiffRender -> m ()
dfPutText t (DiffRender { dfScreenStateBack = ScreenState {ssLines = ssLns, ssCursorOverflow = cursorOverflow, ssCursorPos = ScreenPos cx cy} }) =  do
    -- Write stuff to the backbuffer. If the cursor is in an overflow position, then do nothing.
    if cursorOverflow
      then pure ()
      else liftIO $ flip (MV.modify ssLns) cy $ \l -> stInsert l cx t

dfDraw :: MonadIO m => Maybe (ScreenPos, Dimensions) -> DiffRender -> m DiffRender
dfDraw mScreenParams dfr@(DiffRender { dfScreenStateBack = (ssLines -> ssb), dfScreenState = (ssLines -> ss) }) =  do
  let
    (screenOffsetX, screenOffsetY) = case mScreenParams of
      Just (sp, _) -> (sX sp, sY sp)
      Nothing -> (0, 0)

  liftIO $ MV.imapM_ (\idx neLine -> do
    oldLine <- MV.read ss idx
    if (oldLine /= neLine)
      then do
        A.setCursorPosition (idx + screenOffsetY) screenOffsetX
        -- mapM_ (\x -> do T.putStr x; S.hFlush stdout; threadDelay 10000;) (stRender <$> neLine)
        mapM_ T.putStr (stRender <$> neLine)
        S.hFlush S.stdout
      else pure ()
    ) ssb
  pure $ dfr { dfScreenState = dfScreenStateBack dfr, dfScreenStateBack = dfScreenState dfr }

dfClear :: MonadIO m => DiffRender -> m ()
dfClear dfIn = do
  let bb = dfScreenStateBack dfIn
  liftIO $ MV.set (ssLines bb) [Plain (T.replicate (ssColumns bb) " ")]

dfInitialize :: MonadIO m => Dimensions -> DiffRender -> m DiffRender
dfInitialize (Dimensions cols rows) dfIn = do
    -- Initialize the screen memory for the dimensions
    -- and initialize to whitespaces.
    (ss, ssBack) <- liftIO $ do
      ss <- emptyScreenState rows cols
      ssBack <- emptyScreenState rows cols
      pure (ss, ssBack)
    pure $ dfIn { dfScreenStateBack = ssBack, dfScreenState = ss}

data ScreenState = ScreenState
  { ssLines          :: IOVector [StyledText]
  , ssCursorPos      :: ScreenPos
  , ssColumns        :: Int
  , ssCursorOverflow :: Bool
  }

screenStateDimension :: ScreenState -> Dimensions
screenStateDimension ScreenState {..} = Dimensions ssColumns (MV.length ssLines)

diffRenderToDimension :: DiffRender -> Dimensions
diffRenderToDimension DiffRender {..} = screenStateDimension dfScreenState

data Style = FgBg Color Color| Fg Color | Bg Color | TextUnderline | NoStyle
  deriving (Eq, Show)

data StyledText = StyledText Style [StyledText] | Plain Text
  deriving (Show, Eq)

renderLines :: [[StyledText]] -> Text
renderLines st = T.intercalate "\n" $ (\i -> (T.concat $ stRender <$> i)) <$> st

instance HasGen StyledText where
  getGen = recursive choice
    [Plain . T.pack <$> txGen]
    [StyledText NoStyle <$> list (linear 1 100) getGen]
    where
      txGen = list (linear 1 100) (choice [lower, upper])

-- Punch a hole starting from offset of size length. Columns start
-- from 0. Returns chunks left and right of the hole.
punchHole :: (Int, Int) -> [StyledText] -> ([StyledText], [StyledText])
punchHole (tk, ln) sts = (stTake tk sts, stDrop (tk + ln) sts)

stInsert :: [StyledText] -> Int -> StyledText -> [StyledText]
stInsert sts cx t =
  let (lr, rg) = punchHole (cx, stLength t) sts
  in (lr <> [t] <> rg)

stDrop :: Int -> [StyledText] -> [StyledText]
stDrop l stsIn = snd $ DL.foldl' stDrop' (l, []) stsIn
  where
    stDrop' :: (Int, [StyledText]) -> StyledText -> (Int, [StyledText])
    stDrop' (0, sts) st = (0, sts <> [st])
    stDrop' (s, sts) (Plain t) = let
      r = T.drop s t
      rlen = T.length r
      in if rlen > 0
        then (s - (T.length t - rlen), sts <> [Plain r])
        else (s - T.length t, sts)
    stDrop' (s, sts) (StyledText st sts') =
     case DL.foldl' stDrop' (s, []) sts' of
      (s', sts''@(_:_)) -> (s', sts <> [StyledText st sts''])
      (s', _)           -> (s', sts)

stTake :: Int -> [StyledText] -> [StyledText]
stTake l stsIn = snd $ DL.foldl' stTake' (l, []) stsIn
  where
    stTake' :: (Int, [StyledText]) -> StyledText -> (Int, [StyledText])
    stTake' (0, sts) _ = (0, sts)
    stTake' (s, sts) (Plain t) = let
      r = T.take s t
      in (s - T.length r, sts <> [Plain r])
    stTake' (s, sts) (StyledText st sts') = let
      (s', sts'') = DL.foldl' stTake' (s, []) sts'
      in (s', sts <> [StyledText st sts''])

applyStyleToRange :: (StyledText -> StyledText) -> Int -> (Int, Int) -> [StyledText] -> [StyledText]
applyStyleToRange fn sStart (rStart, rEnd) segments =
  let
    sEnd = sStart + stTotalLength segments - 1
  in if rEnd >= sStart && rStart <= sEnd
    then
      let
        r1 = max 0 (rStart - sStart)
        r2 = rEnd - (max sStart rStart) + 1
      in (stTake r1 segments) <> (fn <$> (stTake r2 (stDrop r1 segments))) <> (stDrop (r1 + r2) segments)
    else segments

stLength :: StyledText -> Int
stLength (Plain t)          = T.length t
stLength (StyledText _ sts) = sum (stLength <$> sts)

stTotalLength :: [StyledText] -> Int
stTotalLength sts = sum $ stLength <$> sts

-- This implimentation is not optimal and does not correctly
-- render nested styles. But this appear to be good enough for now.
styleToPrefix :: Style -> Text
styleToPrefix st = T.pack $ case st of
  FgBg fg bg -> setSGRCode [SetColor Foreground Vivid fg] <> setSGRCode [SetColor Background Vivid bg]
  Bg bg -> setSGRCode [SetColor Background Vivid bg]
  Fg fg -> setSGRCode [SetColor Foreground Vivid fg]
  TextUnderline -> setSGRCode [SetUnderlining SingleUnderline]
  NoStyle -> mempty

stRender :: StyledText -> Text
stRender st = stRender' NoStyle st

mergeStyles :: Style -> Style -> Style
mergeStyles NoStyle a     = a
mergeStyles a NoStyle     = a
mergeStyles (Fg a) (Bg b) = (FgBg a b)
mergeStyles (Bg a) (Fg b) = (FgBg b a)
mergeStyles _ a           = a

stRender' :: Style -> StyledText -> Text
stRender' _ (Plain t) = T.replace "\n" " " t
stRender' pst (StyledText st sts) = let
  suffix = T.pack $ setSGRCode []
  prefix = styleToPrefix (mergeStyles pst st)
  in T.concat (((\x -> prefix <> x <> suffix) . stRender' st) <$> sts)