Common.hs

module UI.Widgets.Common
  ( module UI.Widgets.Common
  , module UI.Terminal.IO
  , module Control.Monad
  , module Data.Text
  , module Control.Monad.IO.Class
  , module Control.Monad.State.Strict
  , module Data.Constraint
  , module GHC.Stack
  ) where

import Common
import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TChan
import Control.Exception
import Control.Monad.IO.Class
import Control.Monad.Loops (iterateWhile)
import Data.Bits
import Data.Constraint
import Data.Kind (Type)
import Data.Map.Strict as M hiding (keys)
import Data.Maybe
import Data.Text as T
import Data.Text hiding (lines)
import qualified Data.Text as C
import Data.Text.IO as T
import Data.Typeable (Proxy(..), Typeable, cast, typeRep)
import qualified Data.Vector.Mutable as MV
import Data.Word
import GHC.Stack
import qualified System.IO as SIO
import System.Random
import qualified System.Terminal as TERM
import UI.Chars
import UI.Terminal.IO

import Control.Monad
import Control.Monad.State.Strict
import DiffRender.DiffRender
import Highlighter.Highlighter
import qualified System.Console.ANSI as A
import qualified System.IO as S

data WidgetState = WidgetState
  { wsWidgets       :: Map Int SomeWidget
  , wsCursorWidget  :: Maybe SomeKeyInputWidget -- Widget that authoritativly decide the status/location of cursor. Does not decide what widgets receive keyboard input
  , wsCursorVisible :: Bool
  }

data UIState = UIState
  { usWidgetState :: WidgetState
  , usDiffRender  :: DiffRender
  }

liftWSMod :: (WidgetState -> WidgetState) -> (UIState -> UIState)
liftWSMod fn = (\us -> us { usWidgetState = fn $ usWidgetState us })

class HasCharScreen m where
  csInitialize :: Dimensions -> m ()
  csClear :: m ()
  csDraw :: m ()
  csPutText :: StyledText -> m ()
  csSetCursorPosition :: Int -> Int -> m ()

setCursorVisibility :: WidgetC m => Bool ->  m ()
setCursorVisibility b = modify $ liftWSMod (\ws -> ws { wsCursorVisible = b })

emptyWidgetState :: WidgetState
emptyWidgetState = do
  WidgetState mempty Nothing True

emptyUIState :: Dimensions -> IO UIState
emptyUIState dim = do
  dfr <- emptyDiffRender dim
  pure $ UIState emptyWidgetState dfr

type WidgetM m a = MonadIO m => StateT UIState m a

runWidgetM' :: MonadIO m => WidgetM m a -> m (a, UIState)
runWidgetM' act = do
  ws <- liftIO $ emptyUIState $ Dimensions 0 0
  runWidgetM'' ws act

runWidgetM'' :: MonadIO m => UIState -> WidgetM m a -> m (a, UIState)
runWidgetM'' ws act = flip runStateT ws act

runWidgetM :: MonadIO m => WidgetM m a -> m a
runWidgetM act =  fst <$> runWidgetM' act

type WidgetC m =
  ( HasCallStack
  , HasCharScreen m
  , HasRandom m
  , HasLog m
  , HasTerminal m
  , MonadState UIState m
  , MonadIO m
  )

getScreenBounds :: WidgetC m => m Dimensions
getScreenBounds = do
  screenState <- (dfScreenState . usDiffRender) <$> get
  let
    screenLines = ssLines screenState
    screenColumns = ssColumns screenState
  pure $ Dimensions screenColumns (MV.length screenLines)

instance MonadIO m => HasRandom (StateT UIState m) where
  getRandom = liftIO randomIO

instance MonadIO m => HasCharScreen (StateT UIState m) where
  csInitialize dim = do
    emptyDfr <- liftIO $ emptyDiffRender dim
    modify (\ws -> ws { usDiffRender = emptyDfr })

  csClear = usDiffRender <$> get >>= dfClear

  csDraw = do
    nDfr <- usDiffRender <$> get >>= dfDraw Nothing
    modify (\ws -> ws { usDiffRender = nDfr })

    (wsCursorVisible . usWidgetState) <$> get >>= \case
      False -> pure ()
      True ->
        ((wsCursorWidget . usWidgetState) <$> get) >>= \case
          Just (SomeKeyInputWidget fref) -> getCursorInfo fref >>= \case
            Just (cl, csst) -> do
              liftIO $ A.setCursorPosition (sY cl) (sX cl)
              putTextFlush $ cursorStyleCode csst
            Nothing -> pure ()
          Nothing -> pure ()

  csPutText t = usDiffRender <$> get >>= dfPutText t

  csSetCursorPosition x y = do
    dfr <- usDiffRender <$> get
    modify (\ws -> ws { usDiffRender = dfSetCursorPosition x y dfr })

getTerminalSizeIO :: IO (Maybe (Int, Int))
getTerminalSizeIO = do
  A.getTerminalSize >>= \case
    Just (y, x) -> pure $ Just (x, y)
    Nothing     -> pure Nothing

instance MonadIO m => HasTerminal (StateT UIState m) where
  setCursorPosition x y = do
    liftIO $ A.setCursorPosition y x
    hFlush
  hideCursor = liftIO A.hideCursor
  showCursor = do
    liftIO A.showCursor
    hFlush
  putText t = liftIO $ do
    T.putStr t
  putTextFlush t = do
    putText t
    hFlush
  hFlush = liftIO $ S.hFlush S.stdout
  hSetEcho h b = liftIO $ S.hSetEcho h b
  hGetChar = liftIO $ S.hGetChar S.stdin
  hSetBuffering h b = liftIO $ S.hSetBuffering h b
  hWaitForInput = liftIO $ S.hWaitForInput stdin 0
  clearscreen = do
    liftIO A.clearScreen
    hFlush
  clearline = liftIO $ A.hClearFromCursorToLineEnd stdout

instance MonadIO m => HasLog (StateT UIState m) where
  appendLog a = liftIO (appendLog a)

-- Below, the type parameter `a` is left in case we need to use tagged
-- references, like an IORef.
newtype WRef (a :: Type) = WRef Int
  deriving newtype (Eq, Ord)

strToKeyEvent :: String -> [KeyEvent]
-- gnome-terminal
strToKeyEvent ('\DEL': rst)  =  (KeyCtrl False False False Backspace):strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'A' : rst)  =  (KeyCtrl False False False ArrowUp) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'B' : rst)  =  (KeyCtrl False False False ArrowDown) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'H' : rst)  =  (KeyCtrl False False False Home) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'F' : rst)  =  (KeyCtrl False False False End) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'C' : rst)  =  (KeyCtrl False False False ArrowRight) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'D' : rst)  =  (KeyCtrl False False False ArrowLeft) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '2' : '~': rst)  =  (KeyCtrl False False False Insert) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '3' : '~': rst)  =  (KeyCtrl False False False Del) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1' : '5': '~': rst)  =  (KeyCtrl False False False (Fun 5)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1' : '7': '~': rst)  =  (KeyCtrl False False False (Fun 6)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '2' : '1': '~': rst)  =  (KeyCtrl False False False (Fun 10)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'P': rst)  =  (KeyCtrl False False False (Fun 1)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'Q': rst)  =  (KeyCtrl False False False (Fun 2)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'R': rst)  =  (KeyCtrl False False False (Fun 3)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': '9': ';' : '5' : '~' : rst)  = (KeyCtrl True False False (Fun 8)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'C' : rst)  = (KeyCtrl True False False (ArrowRight)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'D' : rst)  = (KeyCtrl True False False (ArrowLeft)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'B' : rst)  = (KeyCtrl True False False (ArrowDown)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'A' : rst)  = (KeyCtrl True False False (ArrowUp)) : strToKeyEvent rst

-- Linux term
strToKeyEvent ('\ESC': '[' : '[' : 'E': rst)  =  (KeyCtrl False False False (Fun 5)) : strToKeyEvent rst
--
strToKeyEvent ('\ESC': c : rst )  =  (KeyChar False False True c) : strToKeyEvent rst
strToKeyEvent ('\ESC' : rst)  =  (KeyCtrl False False False Esc) : strToKeyEvent rst
strToKeyEvent ('\n': rst)  =  (KeyCtrl False False False Return) : strToKeyEvent rst
strToKeyEvent str       =  KeyChar False False False <$> str

-- Some terminals sets the eighth bit, instead of sending
-- Escape code to indicate alt key sequences. So this function
-- checks for the 8th bit, and if it is set, sets the alt modifier
-- and changes the original char to one with 8th bit unset.
convertFrom8Bit :: KeyEvent -> KeyEvent
convertFrom8Bit kc@(KeyChar c s _ char) = let
  cw = (toEnum $ fromEnum char) :: Word8
  mask' = complement (bit 7)
  in if ((bit 7 .&. cw) > 0) then (KeyChar c s True (toEnum $ fromEnum (cw .&. mask'))) else kc
convertFrom8Bit c = c

readTerminalEvent :: (TERM.MonadScreen m, TERM.MonadInput m, MonadIO m) => m [TerminalEvent]
readTerminalEvent = TERM.awaitEvent >>= \case
    Left _ -> do
      pure [TerminalInterrupt]
    Right x ->
      case x of
        TERM.WindowEvent TERM.WindowSizeChanged -> do
          TERM.Size h w <- TERM.getWindowSize
          pure [TerminalResize w h]
        TERM.KeyEvent k mods -> pure $ TerminalKey <$> case k of
          TERM.CharKey c -> [convertFrom8Bit $ setModifiers mods $ KeyChar False False False c]
          TERM.ArrowKey TERM.Upwards -> [setModifiers mods $ KeyCtrl False False False ArrowUp]
          TERM.ArrowKey TERM.Downwards -> [setModifiers mods $ KeyCtrl False False False ArrowDown]
          TERM.ArrowKey TERM.Rightwards -> [setModifiers mods $ KeyCtrl False False False ArrowRight]
          TERM.ArrowKey TERM.Leftwards -> [setModifiers mods $ KeyCtrl False False False ArrowLeft]
          TERM.EscapeKey -> [setModifiers mods $ KeyCtrl False False False Esc]
          TERM.EnterKey -> [setModifiers mods $ KeyCtrl False False False Return]
          TERM.FunctionKey n -> [setModifiers mods $ KeyCtrl False False False (Fun n)]
          TERM.BackspaceKey -> [setModifiers mods $ KeyCtrl False False False Backspace]
          TERM.HomeKey -> [setModifiers mods $ KeyCtrl False False False Home]
          TERM.EndKey -> [setModifiers mods $ KeyCtrl False False False End]
          TERM.DeleteKey -> [setModifiers mods $ KeyCtrl False False False Del]
          _ -> []
        _ -> pure []
  where
    setModifiers :: TERM.Modifiers -> KeyEvent -> KeyEvent
    setModifiers modifiers (KeyCtrl _ _ _ c) = KeyCtrl ((modifiers .&. TERM.ctrlKey) /= mempty) ((modifiers .&. TERM.shiftKey) /= mempty) ((modifiers .&. TERM.altKey) /= mempty) c
    setModifiers modifiers (KeyChar _ _ _ c) = KeyChar ((modifiers .&. TERM.ctrlKey) /= mempty) ((modifiers .&. TERM.shiftKey) /= mempty) ((modifiers .&. TERM.altKey) /= mempty) c

readKey :: IO [KeyEvent]
readKey = do
  k <- readKey_
  pure $ strToKeyEvent k

readKey_ :: IO String
readKey_ = do
  char <- S.hGetChar S.stdin
  readRest [char]
  where
    readRest :: [Char] -> IO [Char]
    readRest t = S.hWaitForInput stdin 0 >>= \case
      True -> do
        c <- S.hGetChar S.stdin
        readRest (c:t)
      False -> pure $ Prelude.reverse t

uiLoop :: forall m event. (Show event, WidgetC m) => TChan event -> (event -> m Bool) -> m ()
uiLoop es cback = do
  void $ iterateWhile id $ do
    event <- liftIO $ atomically $ readTChan es
    r <- cback event
    pure r

cursorStyleCode :: CursorStyle -> Text
cursorStyleCode Bar       = T.pack $ "\ESC[5 q" <> A.showCursorCode
cursorStyleCode Underline = T.pack $ "\ESC[4 q" <> A.showCursorCode
cursorStyleCode Hidden    = T.pack A.hideCursorCode

readWRef :: forall a m. (WidgetC m, Widget a) => WRef a -> m a
readWRef (WRef ref) = do
  (fromMaybe (error "not found") . M.lookup ref . wsWidgets . usWidgetState) <$> get >>= \case
    SomeWidget w -> case cast w of
      Just a  -> pure a
      Nothing -> error "Unexpected type"

modifyWRef :: (WidgetC m, Widget a) => WRef a -> (a -> a) -> m ()
modifyWRef (WRef ref) fn =
  modify $ liftWSMod $ \s -> s { wsWidgets = M.update (Just . (modifySomeWidget fn)) ref $ wsWidgets s }

modifyWRefM :: (WidgetC m, Widget a) => WRef a -> (a -> m a) -> m ()
modifyWRefM (WRef ref) fn = do
  m <- (wsWidgets . usWidgetState) <$> get
  case M.lookup ref m of
    Just sw -> do
      nSw <- modifySomeWidgetM fn sw
      modify $ liftWSMod $ \s -> s { wsWidgets = M.update (\_ -> Just nSw) ref $ wsWidgets s }
    Nothing -> pure ()

modifySomeWidget :: Widget a => (a -> a) -> SomeWidget -> SomeWidget
modifySomeWidget fn (SomeWidget w) = case cast w of
  Just a  -> (SomeWidget (fn a))
  Nothing -> error "unexpected type"

modifySomeWidgetM :: (Monad m, Widget a) => (a -> m a) -> SomeWidget -> m SomeWidget
modifySomeWidgetM fn (SomeWidget w) = case cast w of
  Just a  -> do
    n <- fn a
    pure (SomeWidget n)
  Nothing -> error "unexpected type"

-- Insert the new widget at a random key in the Widget state map
-- and return the key.
newWRef :: (WidgetC m, Widget a) => a -> m (WRef a)
newWRef a = do
  ref <- getRandom
  modify $ liftWSMod $ \s -> s { wsWidgets = M.insert ref (SomeWidget a) $ wsWidgets s }
  pure (WRef ref)

data CtrlKey
  = Del
  | Esc
  | Insert
  | End
  | Home
  | ArrowLeft
  | ArrowRight
  | ArrowUp
  | ArrowDown
  | Backspace
  | Fun Int
  | Return
  deriving (Show, Ord, Eq)

data KeyEvent
  = KeyChar Bool Bool Bool Char -- Bool fields for modifiers for ctrl, shift, alt
  | KeyCtrl Bool Bool Bool CtrlKey
  deriving (Show, Eq, Ord)

data TerminalEvent
  = TerminalKey KeyEvent
  | TerminalResize Int Int
  | TerminalInterrupt
  deriving Show

pushTerminalEventToHandle :: SIO.Handle -> TerminalEvent -> IO ()
pushTerminalEventToHandle handle' kv = case kv of
  TerminalKey (KeyChar False False False c) -> do
    SIO.hPutChar handle' c
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Backspace) -> do
    SIO.hPutChar handle' '\BS'
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Esc) -> do
    SIO.hPutChar handle' '\ESC'
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Return) -> do
    SIO.hPutChar handle' '\n'
    SIO.hFlush handle'
  _ -> pass

data TerminalException
  = TerminalException Text
  deriving (Show)

instance Exception TerminalException

class HasRandom m where
  getRandom :: Random a => m a

class HasCursor m where
  getCursor :: m CursorInfo

class Layout a where
  addWidget :: (WidgetC m, Widget child) => WRef a -> Text -> WRef child -> m ()

class Drawable a where
  draw :: WidgetC m => WRef a -> m ()
  setVisibility :: WidgetC m => WRef a -> Bool -> m ()
  getVisibility :: WidgetC m => WRef a -> m Bool

class Moveable a where
  move ::  WidgetC m => WRef a -> ScreenPos -> m ()
  getPos :: WidgetC m => WRef a -> m ScreenPos
  getDim :: WidgetC m => WRef a -> m Dimensions
  resize :: WidgetC m => WRef a -> (Dimensions -> Dimensions) -> m ()

class Container a c | a -> c where
  setContent :: WidgetC m => WRef a -> c -> m ()
  getContent :: WidgetC m => WRef a -> m c

class Selectable a where
  getSelection :: WidgetC m => WRef a -> m Text

class (Typeable a, Drawable a) => KeyInput a where
  getCursorInfo :: WidgetC m => WRef a -> m (Maybe CursorInfo)
  handleInput :: WidgetC m => WRef a -> KeyEvent -> m ()

data WidgetCapability a (c :: Constraint) where
  KeyInputCap :: WRef a -> WidgetCapability a (KeyInput a)
  MoveableCap :: WRef a -> WidgetCapability a (Moveable a)
  SelectableCap :: WRef a -> WidgetCapability a (Selectable a)
  DrawableCap :: WRef a -> WidgetCapability a (Drawable a)
  ContainerCap :: Typeable cnt => WRef a -> Proxy cnt -> WidgetCapability a (Container a cnt)

class Typeable a => Widget a where
  hasCapability :: WidgetCapability a c -> Maybe (Dict c)

withCapability :: forall a c m b . (WidgetC m, Widget a, Typeable c) => WidgetCapability a c -> (c => m b)  -> m b
withCapability cap fn = case hasCapability cap of
  Just Dict -> fn
  Nothing   -> error ("No capability:" <> (show $ typeRep (Proxy @c)))

data SomeWidget where
  SomeWidget :: forall a. Widget a => a -> SomeWidget

data SomeWidgetRef where
  SomeWidgetRef :: forall a. (Typeable a, Widget a) => WRef a -> SomeWidgetRef

data SomeKeyInputWidget where
  SomeKeyInputWidget :: KeyInput a => WRef a -> SomeKeyInputWidget

wSetCursor :: (WidgetC m, HasTerminal m) => ScreenPos -> m ()
wSetCursor ScreenPos {..} = csSetCursorPosition sX sY

wSetCursorRel :: (HasTerminal m, WidgetC m) => ScreenPos -> ScreenPos -> m ScreenPos
wSetCursorRel o rel = do
  let n = moveRight (sX rel) $ moveDown (sY rel) o
  wSetCursor n
  pure n

drawTitleLine :: WidgetC m => ScreenPos -> Int -> Int -> Maybe Text -> m ()
drawTitleLine sp width titleOffset (fromMaybe "" -> title) = do
  wSetCursor sp
  let titleLength = T.length title
  csPutText $ StyledText NoStyle [Plain $ C.concat [C.replicate (titleOffset - 1) (C.singleton horizontalLine), C.singleton verticalRight], colorText A.White A.Blue title, Plain (C.singleton verticalLeft <> C.replicate (width - titleLength - titleOffset - 1) (C.singleton horizontalLine)) ]

drawTitleLineVertical :: WidgetC m => ScreenPos -> Int -> Int -> Maybe Text -> m ()
drawTitleLineVertical sp height titleOffset (fromMaybe "" -> title) = do
  wSetCursor sp
  let titleLength = T.length title

  forM_ [0..(titleOffset - 1)] (\r -> do
    wSetCursor $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    )

  forM_ (Prelude.zip [titleOffset ..] (T.unpack title)) (\(r, c) -> do
    wSetCursor $ moveDown r sp
    csPutText $ colorText A.White A.Blue $ C.singleton c
    )

  forM_ [(titleOffset + titleLength) .. (height - 1)] (\r -> do
    wSetCursor $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    )

drawBorderBox :: WidgetC m => ScreenPos -> Dimensions -> m ()
drawBorderBox sp Dimensions {..} = do
  wSetCursor sp
  csPutText $ Plain $ C.concat [C.singleton cornerLT, C.replicate (diW - 2) (C.singleton horizontalLine), C.singleton cornerRT]
  wSetCursor $ moveDown (diH - 1) sp
  csPutText $ Plain $ C.concat [C.singleton cornerLB, C.replicate (diW - 2) (C.singleton horizontalLine), C.singleton cornerRB]
  forM_ [1..(diH - 2)] (\r -> do
    wSetCursor $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    wSetCursor $ moveRight (diW - 1) $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    )