path: root/src/main.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ViewPatterns #-}
module Main (main) where

import Blessings.Text (Blessings(Plain,SGR),pp)
import Control.Applicative ((<|>))
import Control.Concurrent
import Control.Monad (forM)
import Control.Monad (forever)
import Data.Bits ((.|.),testBit)
import Data.Default (def)
import Data.Function ((&))
import Data.List.Extra ((!!?))
import Data.Maybe (catMaybes,fromMaybe)
import Data.Monoid.Extra (mintercalate)
import Data.Set (Set)
import Data.Text (Text)
import Foreign.C.Types (CLong)
import Much.Screen (Screen(Screen), withScreen)
import Pager.Types
import Scanner
import State (State(..))
import System.Environment (getArgs)
import System.IO
import System.Posix.Signals (Handler(Catch), Signal, installHandler, sigINT)
import qualified Blessings.Internal as Blessings
import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Text.Extra as Text
import qualified Data.Text.IO as Text
import qualified Data.Text.Read as Text
import qualified Graphics.X11 as X11
import qualified Graphics.X11.Extra as X11
import qualified Graphics.X11.Xlib.Extras as X11
import qualified Graphics.X11.Xlib.Extras.Extra as X11
import qualified Hack.Buffer as Buffer
import qualified Hack.Buffer.Extra as Buffer
import qualified Pager.Sixelerator as Pager
import qualified System.Console.Terminal.Size as Term


atom_NET_ACTIVE_WINDOW :: X11.Atom
atom_NET_ACTIVE_WINDOW = X11.unsafeInternAtom "_NET_ACTIVE_WINDOW" True

atom_NET_CLIENT_LIST :: X11.Atom
atom_NET_CLIENT_LIST = X11.unsafeInternAtom "_NET_CLIENT_LIST" True

atom_NET_CURRENT_DESKTOP :: X11.Atom
atom_NET_CURRENT_DESKTOP = X11.unsafeInternAtom "_NET_CURRENT_DESKTOP" True

atom_NET_DESKTOP_NAMES :: X11.Atom
atom_NET_DESKTOP_NAMES = X11.unsafeInternAtom "_NET_DESKTOP_NAMES" True

atom_NET_WM_DESKTOP :: X11.Atom
atom_NET_WM_DESKTOP = X11.unsafeInternAtom "_NET_WM_DESKTOP" True

atom_NET_WM_NAME :: X11.Atom
atom_NET_WM_NAME = X11.unsafeInternAtom "_NET_WM_NAME" True

atom_WIN_WORKSPACE :: X11.Atom
atom_WIN_WORKSPACE = X11.unsafeInternAtom "_WIN_WORKSPACE" True

atom_WIN_WORKSPACE_NAMES :: X11.Atom
atom_WIN_WORKSPACE_NAMES = X11.unsafeInternAtom "_WIN_WORKSPACE_NAMES" True

atom_WM_NAME :: X11.Atom
atom_WM_NAME = X11.unsafeInternAtom "WM_NAME" True


getActiveWindow :: X11.Display -> IO (Maybe X11.Window)
getActiveWindow d =
  (fmap (fromIntegral . head) <$>) $
      X11.getWindowProperty32 d atom_NET_ACTIVE_WINDOW w
  where w = X11.defaultRootWindow d

getCurrentDesktop :: X11.Display -> IO (Maybe CLong)
getCurrentDesktop d =
  (fmap head <$>) $
      X11.getWindowProperty32 d atom_NET_CURRENT_DESKTOP w <|>
      X11.getWindowProperty32 d atom_WIN_WORKSPACE w
  where w = X11.defaultRootWindow d

getDesktopNames :: X11.Display -> IO (Maybe [Text])
getDesktopNames d = do
  (fmap (init . Text.split (=='\NUL')) <$>) $
      X11.getWindowPropertyText d atom_NET_DESKTOP_NAMES w <|>
      X11.getWindowPropertyText d atom_WIN_WORKSPACE_NAMES w
  where w = X11.defaultRootWindow d

getGeometry :: X11.Display -> X11.Window -> IO Geometry
getGeometry d w = do
  (_, x, y, width, height, _, _) <- X11.getGeometry d w
  return Geometry
    { geometry_x = fromIntegral x
    , geometry_y = fromIntegral y
    , geometry_width = fromIntegral width
    , geometry_height = fromIntegral height
    }

getWindowDesktop :: X11.Display -> X11.Window -> IO (Maybe CLong)
getWindowDesktop d w =
  (fmap head <$>) $
    X11.getWindowProperty32 d atom_NET_WM_DESKTOP w <|>
    X11.getWindowProperty32 d atom_WIN_WORKSPACE w

getWindowTitle :: X11.Display -> X11.Window -> IO (Maybe Text)
getWindowTitle d w =
    X11.getWindowPropertyText d atom_NET_WM_NAME w <|>
    X11.getWindowPropertyText d atom_WM_NAME w

getWorkspaces :: X11.Display -> Geometry -> Set X11.Window -> IO [Workspace]
getWorkspaces display screenGeometry focusWindows = do
    let rootWindow = X11.defaultRootWindow display

    currentDesktop <- fromMaybe 0 <$> getCurrentDesktop display

    workspaces <- do
      names <- zip [0..] . fromMaybe [] <$> getDesktopNames display
      ws <-
        forM names $ \(index, name) -> do
          return Workspace
              { workspace_geometry = screenGeometry
              , workspace_focused = currentDesktop == index
              , workspace_name = name
              , workspace_windows = []
              }
      return $ Map.fromList $ zip [0..] ws

    clientList <-
      maybe [] (map fromIntegral) <$>
        X11.getWindowProperty32 display atom_NET_CLIENT_LIST rootWindow

    let
        f w = do
            title <- getWindowTitle display w
            desktop <- fromMaybe 0 <$> getWindowDesktop display w
            geometry <- getGeometry display w

            wm_hints <- X11.getWMHints display w
            let urgent = testBit (X11.wmh_flags wm_hints) X11.urgencyHintBit

            let
              window =
                Window
                  { window_id = fromIntegral w
                  , window_title = fromMaybe "" title
                  , window_geometry = geometry
                  , window_focused = Set.member w focusWindows
                  , window_urgent = urgent
                  }

            return ( window, desktop )

    clientList' <- mapM f clientList

    return
      $ map (\ws -> ws { workspace_windows =
                            uncurry (<>) $
                            List.partition window_focused (workspace_windows ws)
                      })
      $ Map.elems
      $ foldr
        (\(w, i) ->
          Map.adjust (\ws -> ws { workspace_windows = w : workspace_windows ws })
            i
        )
        workspaces
        clientList'


main :: IO ()
main = do
    args <- getArgs
    let
      ( commandFromArgs, focusWindows ) =
        let readInt s =
              case Text.decimal (Text.pack s) of
                Right (i, "") ->
                  i
                _ ->
                  (-1)
        in
        case args of
          "shift" : focusWindows_ ->
            ( ShiftWindowToWorkspace undefined, map readInt focusWindows_ )

          "shiftview" : focusWindows_ ->
            ( ShiftWindowToAndViewWorkspace undefined, map readInt focusWindows_ )

          "view" : focusWindows_ ->
            ( ViewWorkspace, map readInt focusWindows_ )

          _ ->
            error $ "bad arguments: " <> show args


    Just activeWindow <- X11.withDefaultDisplay getActiveWindow

    screenGeometry <-
      X11.withDefaultDisplay $ \display -> do
        let rootWindow = X11.defaultRootWindow display
        getGeometry display rootWindow

    workspaces <-
      X11.withDefaultDisplay $ \display -> do
        getWorkspaces display screenGeometry (Set.fromList focusWindows)

    let screen0 = Screen False NoBuffering (BlockBuffering $ Just 4096)
                [ 1000 -- X & Y on button press and release
                , 1005 -- UTF-8 mouse mode
                , 1047 -- use alternate screen buffer
                ,   80 -- enable sixel scrolling 
                ]
                [   25 -- hide cursor
                ]
    result <- do
      withFile "/dev/tty" ReadWriteMode $ \i ->
        withFile "/dev/tty" WriteMode $ \o ->
          withScreen i o screen0 $ \_ -> do
            (putEvent, getEvent) <- do
                v <- newEmptyMVar
                return (putMVar v, takeMVar v)

            let q1 =
                  updateFoundWorkspaces $ def
                    { command = commandFromArgs
                    , screenHeight = geometry_height screenGeometry
                    , screenWidth = geometry_width screenGeometry
                    , termBorder = 2
                    , workspaces =
                        let
                          f workspace@Workspace{workspace_name} = ( workspace_name, workspace )
                        in
                        Map.fromList (map f workspaces)
                    }
                signalHandlers =
                    [ (sigINT, putEvent EShutdown)
                    , (28, winchHandler i putEvent)
                    ]

            installHandlers signalHandlers

            threadIds <- mapM forkIO
                [
                  forever $
                    scan i >>= putEvent . EScan
                ]

            winchHandler i putEvent

            result <- run o getEvent q1

            mapM_ killThread threadIds

            return result


    case snd result of
      FocusWorkspace name -> do
          case command (fst result) of
            ViewWorkspace -> do
              X11.withDefaultDisplay $ \d -> do
                let Just s = name `List.elemIndex` map workspace_name workspaces
                switchDesktop d (fromIntegral s)

            ShiftWindowToWorkspace _ -> do
              X11.withDefaultDisplay $ \d ->
                let
                  Just s = name `List.elemIndex` map workspace_name workspaces
                in
                windowToDesktop d activeWindow (fromIntegral s)

            ShiftWindowToAndViewWorkspace _ -> do
              X11.withDefaultDisplay $ \d -> do
                let Just s = name `List.elemIndex` map workspace_name workspaces
                windowToDesktop d activeWindow (fromIntegral s)
                switchDesktop d (fromIntegral s)

      _ ->
        return ()


switchDesktop :: X11.Display -> CLong -> IO ()
switchDesktop d s =
  X11.allocaXEvent $ \e -> do
    X11.setEventType e X11.clientMessage
    X11.setClientMessageEvent' e w atom_NET_CURRENT_DESKTOP 32 [fromIntegral s,0,0,0,0]
    X11.sendEvent d w False mask e
  where
    w = X11.defaultRootWindow d
    mask = X11.structureNotifyMask


windowToDesktop :: X11.Display -> X11.Window -> CLong -> IO ()
windowToDesktop d w s =
  X11.allocaXEvent $ \e -> do
    X11.setEventType e X11.clientMessage
    X11.setClientMessageEvent' e (fromIntegral w) atom_NET_WM_DESKTOP 32 [fromIntegral s,0,0,0,0]
    X11.sendEvent d (fromIntegral w) True mask e
  where
    mask = X11.substructureRedirectMask .|. X11.substructureNotifyMask


run :: Handle -> IO Event -> State -> IO (State, Action)
run o getEvent = rec . Right where
    rec = \case
        Right q ->
            redraw o q >> getEvent >>= processEvent q >>= rec
        Left q ->
            return q


installHandlers :: [(Signal, IO ())] -> IO ()
installHandlers =
    mapM_ (\(s, h) -> installHandler s (Catch h) Nothing)


processEvent :: State -> Event -> IO (Either (State, Action) State)
processEvent q = \case
    EScan (ScanKey s) -> do
        let
            key = Text.pack s
            (q', action) = keymap key q

            realizeAction = \case
              None ->
                return $ Right q'

              FocusWorkspace name -> do
                return $ Left (q', FocusWorkspace name)

        realizeAction action

    EScan mouseInfo@ScanMouse{} ->
        Right <$> mousemap mouseInfo q
    EShutdown ->
        return $ Left (q,None)
    EResize w h ->
        return $ Right q
            { termWidth = w, termHeight = h
            , flashMessage = Plain $ "resize " <> Text.show (w,h)

            , workspaceViewportHeight = newWorkspaceViewportHeight
            , workspaceViewportOffset = newWorkspaceViewportOffset
            }
      where
        newWorkspaceViewportHeight = h - 2 {- input line + status line -}

        newWorkspaceViewportOffset =
          if newWorkspaceViewportHeight > workspaceViewportHeight q then
            max 0 $ workspaceViewportOffset q + (workspaceViewportHeight q - newWorkspaceViewportHeight)

          else if newWorkspaceViewportHeight <= workspaceCursor q - workspaceViewportOffset q then
            workspaceViewportOffset q + (workspaceViewportHeight q - newWorkspaceViewportHeight)

          else
            workspaceViewportOffset q


moveWorkspaceCursor :: Int -> State -> State
moveWorkspaceCursor i q@State{..} =
    q
      { workspaceCursor = newWorkspaceCursor
      , workspaceViewportOffset = newWorkspaceViewportOffset
      }
  where
    newWorkspaceCursor = max 0 $ min (workspaceCursor + i) $ length foundWorkspaces - 1

    newWorkspaceViewportOffset =
      if newWorkspaceCursor < workspaceViewportOffset then
        newWorkspaceCursor

      else if newWorkspaceCursor >= workspaceViewportOffset + workspaceViewportHeight then
        newWorkspaceCursor - workspaceViewportHeight + 1

      else
        workspaceViewportOffset


setCount :: Int -> State -> State
setCount i q = q { count = i }

keymap :: Text -> State -> ( State, Action )

keymap s
      | [ "\ESC[4"
        , Text.decimal -> Right (termHeightPixels, "")
        , Text.unsnoc -> Just (Text.decimal -> Right (termWidthPixels, "") , 't')
        ] <- Text.split (==';') s
      = \q ->
    ( q { termHeightPixels, termWidthPixels }
    , None
    )

keymap s
      | [ "\ESC[6"
        , Text.decimal -> Right (charHeight, "")
        , Text.unsnoc -> Just (Text.decimal -> Right (charWidth, "") , 't')
        ] <- Text.split (==';') s
      = \q ->
    ( q { charHeight, charWidth }
    , None
    )

-- Up
keymap "\ESC[A" = \q@State{..} ->
    ( moveWorkspaceCursor count q & setCount 1
    , None
    )

-- Down
keymap "\ESC[B" = \q@State{..} ->
    ( moveWorkspaceCursor (-count) q & setCount 1
    , None
    )

-- PgUp
keymap "\ESC[5~" = \q@State{..} ->
    ( moveWorkspaceCursor (count * max 1 (workspaceViewportHeight - 1)) q & setCount 1
    , None
    )

-- PgDn
keymap "\ESC[6~" = \q@State{..} ->
    ( moveWorkspaceCursor (-count * max 1 (workspaceViewportHeight - 1)) q & setCount 1
    , None
    )

-- Right
keymap "\ESC[C" = \q@State{..} ->
    ( q { buffer = Buffer.move Buffer.CharsForward 1 buffer }
    , None
    )

-- Left
keymap "\ESC[D" = \q@State{..} ->
    ( q { buffer = Buffer.move Buffer.CharsBackward 1 buffer }
    , None
    )

-- Home
keymap "\ESC[H" = \q@State{..} ->
    ( q { ex_offsetY = ex_offsetY - 1 }
    , None
    )
keymap "\ESC[7~" = \q@State{..} ->
    ( q { ex_offsetY = ex_offsetY - 1 }
    , None
    )

-- End
keymap "\ESC[F" = \q@State{..} ->
    ( q { ex_offsetY = ex_offsetY + 1 }
    , None
    )
keymap "\ESC[8~" = \q@State{..} ->
    ( q { ex_offsetY = ex_offsetY + 1 }
    , None
    )

-- Backspace
keymap "\b" = \q@State{..} ->
    ( updateFoundWorkspaces q { buffer = Buffer.delete Buffer.CharsBackward 1 buffer }
    , None
    )
keymap "\DEL" = \q@State{..} ->
    ( updateFoundWorkspaces q { buffer = Buffer.delete Buffer.CharsBackward 1 buffer }
    , None
    )

keymap "\n" = \q ->
    ( q
    , maybe None (FocusWorkspace . workspace_name) (lookupCursoredWorkspace q)
    )

keymap s
      | Text.length s == 1
      , [c] <- Text.unpack s
      , Char.isPrint c
      = \q0 ->
    let
      q = updateFoundWorkspaces q0 { buffer = Buffer.insertChar c (buffer q0) }
    in
    ( q
    , case lookupCursoredWorkspace q of
        Just Workspace{workspace_focused,workspace_name} ->
          if length (foundWorkspaces q) == 1 then
            if not workspace_focused then
              FocusWorkspace workspace_name

            else
              None

          else
            None

        Nothing ->
          None
    )

keymap s = \q ->
    ( displayKey s q
    , None
    )

updateFoundWorkspaces :: State -> State
updateFoundWorkspaces q@State{..} =
    q { foundWorkspaces = newFoundWorkspaces
      , workspaceCursor = max 0 $ min workspaceCursor $ length newFoundWorkspaces - 1
      }
  where
    f = Text.isPrefixOf (Text.pack (Buffer.showBuffer buffer))

    newFoundWorkspaces =
      map fst . List.sortOn (workspaceSortKey . snd) . filter (f . fst) . Map.toList $ workspaces

    -- smaller is "more important"
    workspaceSortKey ws@Workspace{..} =
      ( not $ isWorkspaceUrgent ws    -- urgent workspaces are most importsnt
      , workspace_focused             -- focused workspace is least important
      , List.null workspace_windows   -- non-empty workspaces are more important
      , workspace_name                -- sort by name
      )

    isWorkspaceUrgent = List.any window_urgent . workspace_windows


mousemap :: Scan -> State -> IO State
mousemap info = displayMouse info

displayKey :: Text -> State -> State
displayKey s q = q { flashMessage = Plain $ Text.show s }

displayMouse :: Scan -> State -> IO State
displayMouse info q =
    return q { flashMessage = SGR [38,5,202] $ Plain $ Text.show info }


winchHandler :: Handle -> (Event -> IO ()) -> IO ()
winchHandler h putEvent = do
    hPutStr h "\ESC[14t" -- query terminal width / height (in pixels)
    hPutStr h "\ESC[16t" -- query character width / height
    Term.hSize h >>= \case
        Just Term.Window {Term.width = width, Term.height = height} ->
            putEvent $ EResize width height
        Nothing ->
            return ()


redraw :: Handle -> State -> IO ()
redraw o q@State{..} = do
    Text.hPutStr o . pp $
        "\ESC[H"
        <> (mintercalate "\n" $ map eraseRight2 $ render0 q)
        <> workspacePreview
    hFlush o
  where

    workspacePreview :: Blessings Text
    workspacePreview =
      let
        previewWidth = paddingRight * charWidth :: Int
        previewHeight = round $ (fromIntegral previewWidth :: Double) / fromIntegral screenWidth * fromIntegral screenHeight :: Int
        previewGeometry = Geometry
            { geometry_width = fromIntegral previewWidth
            , geometry_height = fromIntegral previewHeight
            , geometry_x = fromIntegral $ termWidthPixels - previewWidth
            , geometry_y = fromIntegral $ ex_offsetY
            }
      in
      fromMaybe "" (renderWorkspacePreview previewGeometry q <$> lookupCursoredWorkspace q)


    renderWorkspacePreview :: Geometry -> State -> Workspace -> Blessings Text
    renderWorkspacePreview geometry qq =
        Plain . Text.decodeUtf8With (\_ _ -> Nothing) . Pager.renderWorkspacePreview geometry qq


    maxWidth = termWidth - paddingRight

    paddingRight = 10

    eraseRight2 s =
        if Blessings.length s < maxWidth then
          s <> SGR [38,5,234] (Plain $ Text.pack $ replicate (maxWidth - Blessings.length s) '@')
        else
          s


render0 :: State -> [Blessings Text]
render0 q@State{..} =
    map (Blessings.take maxWidth) ((shownWorkspacesPadding <> shownWorkspaces) `join` (shownWindowsPadding <> fromMaybe mempty shownWindows)) <>
    [prompt <> inputLine] <>
    [statusLine]
  where

    prompt = SGR [38,5,147] "> "

    inputLine = Blessings.take (maxWidth - Blessings.length prompt) (renderBuffer q)

    statusLine =
        ls <> sp <> rs
      where
        ln = Blessings.length ls
        ls =
          (Blessings.take termWidth
            (SGR [38,5,242] flashMessage <> Plain (Text.show count)))

        sn = termWidth - ln - rn
        sp =
          (SGR [38,5,234] $ Plain $ Text.pack $ replicate sn '#')

        rn = Blessings.length rs
        rs =
          case command of
            ViewWorkspace -> SGR [38,5,236] "view"
            ShiftWindowToWorkspace _ -> SGR [38,5,236] "shift"
            ShiftWindowToAndViewWorkspace _ -> SGR [38,5,236] "sh+vi"


    maxWidth = termWidth - paddingRight

    paddingRight = 10

    n = termHeight - 2

    shownWorkspaces :: [Blessings Text]
    shownWorkspaces =
        reverse $ map showWorkspace $ zip [workspaceViewportOffset..] (take n (drop workspaceViewportOffset $ foundWorkspaces'))
      where
        foundWorkspaces' = catMaybes $ map (flip Map.lookup workspaces) foundWorkspaces

    shownWorkspacesPadding = replicate (n - length shownWorkspaces) (SGR [38,5,234] "~")

    showWorkspace :: (Int, Workspace) -> Blessings Text
    showWorkspace (index, Workspace{..}) =
      let
        isUrgent = any window_urgent workspace_windows

        (ls,rs) = Text.splitAt (Buffer.length buffer) workspace_name
        marker =
          if index == workspaceCursor then
            SGR [38,5,177] "> "
          else
            Plain "  "

        fgColor =
          if isUrgent then
            SGR [38,5,196]
          else if workspace_focused then
            SGR [38,5,238]
          else if length workspace_windows == 0 then
            SGR [38,5,246]
          else
            id
      in
      marker <> SGR [48,5,023] (fgColor (Plain ls)) <> fgColor (Plain rs)

    shownWindows :: Maybe [Blessings Text]

    shownWindows =
      take (length shownWorkspaces) . map (SGR [38,5,236] . Plain . window_title) . workspace_windows <$> lookupCursoredWorkspace q

    shownWindowsPadding = replicate (n - maybe 0 length shownWindows) (SGR [38,5,234] "~")


renderBuffer :: State -> Blessings Text
renderBuffer State{buffer=(ls0,rs0)} =
    let
      ls = Text.pack ls0
      rs = Text.pack rs0
    in
    case Text.uncons rs of
      Just (c, rs') ->
        Plain ls <> SGR [48,5,200] (Plain $ Text.singleton c) <> Plain rs'

      Nothing ->
        let
          c = ' '
        in
        Plain ls <> SGR [48,5,200] (Plain $ Text.singleton c) <> Plain rs


join :: [Blessings Text] -> [Blessings Text] -> [Blessings Text]
join ls rs =
    zipWith (<>) lsFilled rs
  where
    lsWidth = maximum $ map Blessings.length ls
    lsFilled = map (lsFill lsWidth) ls
    lsFill n s =
        if Blessings.length s < n then
          s <> SGR [38,5,234] (Plain (Text.pack (replicate (n - Blessings.length s) '#') <> Text.singleton '|'))
        else
          s <> SGR [38,5,234] (Plain (Text.singleton '|'))


lookupCursoredWorkspace :: State -> Maybe Workspace
lookupCursoredWorkspace State{..} =
  flip Map.lookup workspaces =<< foundWorkspaces !!? workspaceCursor