path: root/src/Rahm/Desktop/WorkspaceWheel.hs

{-# LANGUAGE MultiWayIf #-}

module Rahm.Desktop.WorkspaceWheel
  ( displayWorkspaceWheel,
    displayWorkspaceWheelAtPoint,
    displayWorkspaceWheelInCenter,
  )
where

import Control.Arrow (Arrow (first, second))
import Control.Concurrent (threadDelay)
import Control.Monad (forM_)
import Control.Monad.Trans.Maybe (MaybeT (runMaybeT))
import Data.Char (toLower)
import Data.Function (fix)
import Data.List (find, sortBy)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, isNothing)
import Data.Monoid (Endo (Endo, appEndo))
import Data.Ord (comparing)
import Data.Word
import qualified Graphics.X11 as X
import Rahm.Desktop.Common (click, gotoWorkspace, pointerLocation, pointerScreen)
import Rahm.Desktop.Keys.Wml (getAndResetWindowSelection, getAndResetWindowSelectionOrCurrent)
import Rahm.Desktop.Layout.PinWindow (withWindowsUnpinned)
import Rahm.Desktop.Marking (setAlternateWindows, setAlternateWorkspace)
import Rahm.Desktop.StackSet as W
import Rahm.Desktop.Submap (ButtonOrKeyEvent (ButtonPress, KeyPress), nextButton, nextButtonOrKeyEvent)
import System.IO
import Text.Printf (printf)
import XMonad
import qualified XMonad.Operations as X
import XMonad.Util.Run (spawnPipe)

dzenSize :: (Integral a) => a
dzenSize = 100

data ScreenType = Current | Visible | Hidden deriving (Eq)

data WorkspaceWheelType = Existing | Alphabet Bool Bool | Digits

displayWorkspaceWheel :: X ()
displayWorkspaceWheel =
  uncurry (displayWorkspaceWheelAtPoint workspaceWheelTypeRing)
    =<< pointerLocation

data Ring a = Ring Int [a] (Maybe a) [a]

workspaceWheelTypeRing :: Ring WorkspaceWheelType
workspaceWheelTypeRing =
  ringAdv $
    mkRing
      [ Existing,
        Alphabet False False,
        Alphabet False True,
        Alphabet True False,
        Alphabet True True,
        Digits
      ]

mkRing :: [a] -> Ring a
mkRing lst =
  let as = cycle lst
      rev = cycle (reverse lst)
   in Ring (length lst) rev Nothing as

ringCur :: Ring a -> Maybe a
ringCur (Ring _ _ a _) = a

ringAdv :: Ring a -> Ring a
ringAdv (Ring s p Nothing (n : ns)) = Ring s p (Just n) ns
ringAdv (Ring s p (Just a) (n : ns)) = Ring s (a : p) (Just n) ns

ringPrev :: Ring a -> Ring a
ringPrev (Ring s (p : ps) Nothing ns) = Ring s ps (Just p) ns
ringPrev (Ring s (p : ps) (Just a) ns) = Ring s ps (Just p) (a : ns)

ringFastForwardTo :: (a -> Bool) -> Ring a -> Ring a
ringFastForwardTo fn ring@(Ring s _ a _) =
  case ringFastForwardTo' fn s ring of
    ring'@(Ring s' p' (Just a') n') | isNothing a -> Ring s' (a' : p') Nothing n'
    ring' -> ring'
  where
    ringFastForwardTo' _ 0 ring = ring
    ringFastForwardTo' fn _ ring@(Ring _ _ (Just a) _) | fn a = ring
    ringFastForwardTo' fn n ring = ringFastForwardTo' fn (n - 1) (ringAdv ring)

displayWorkspaceWheelInCenter :: X ()
displayWorkspaceWheelInCenter =
  mapM_
    ( \(W.screenDetail -> (SD (Rectangle x y w h))) ->
        displayWorkspaceWheelAtPoint
          workspaceWheelTypeRing
          (fromIntegral $ x + (fromIntegral w `div` 2))
          (fromIntegral $ y + (fromIntegral h `div` 2))
    )
    =<< pointerScreen

displayWorkspaceWheelAtPoint :: Ring WorkspaceWheelType -> Int -> Int -> X ()
displayWorkspaceWheelAtPoint wheelType px py = do
  click

  (W.StackSet cur vis other _floating) <- gets windowset
  screen <- fromMaybe cur <$> pointerScreen

  allWorkspaces <- getWorkspaces (fromMaybe Existing $ ringCur wheelType)

  let (SD rect) = W.screenDetail screen
      rectsToWorkspaces = calculatePositions px py rect allWorkspaces
  -- rectsToWorkspaces = map (second (first W.tag)) dzenRects

  dzenPipes <- mapM (uncurry launchDzen) rectsToWorkspaces

  after <-
    ( fix $ \recur ring -> do
        evt <- runMaybeT nextButtonOrKeyEvent
        let stop = return ()
        let retry = recur ring

        case evt of
          Nothing -> return stop -- timeout
          Just (ButtonPress _ b)
            | b == button4 ->
                return (displayWorkspaceWheelAtPoint (ringAdv wheelType) px py)
            | b == button5 ->
                return (displayWorkspaceWheelAtPoint (ringPrev wheelType) px py)
          Just (ButtonPress _ b) ->
            withDisplay $ \d -> do
              (_, _, _, fi -> cx, fi -> cy, _, _, _) <-
                io $ X.queryPointer d =<< rootWindow d 0
              case find (\(r, _) -> pointWithin cx cy r) rectsToWorkspaces of
                Just (_, (ws, _)) ->
                  if
                    | b == button1 -> do
                        gotoWorkspace ws
                        return (displayWorkspaceWheelAtPoint wheelType px py)
                    | b == button3 -> do
                        selection <- getAndResetWindowSelectionOrCurrent
                        withWindowsUnpinned selection $
                          windows $
                            appEndo $
                              mconcat (map (Endo . W.shiftWin ws) selection)

                        return (displayWorkspaceWheelAtPoint wheelType px py)
                    | b == 13 -> do
                        winSet <- gets windowset
                        let mWs = find ((== ws) . W.tag) (W.workspaces winSet)
                        forM_ mWs $ \ws -> do
                          let allWindows = W.integrate' (W.stack ws)
                          setAlternateWindows allWindows
                          mapM_ (`setAlternateWorkspace` W.tag ws) allWindows
                          mapM_ killWindow allWindows

                          -- Preemptively delete the windows so we show the updates
                          -- immediately.
                          windows $
                            appEndo $
                              mconcat $
                                map (Endo . W.delete) allWindows

                        X.refresh
                        return (displayWorkspaceWheelAtPoint wheelType px py)
                    | otherwise -> return stop
                Nothing -> return stop
          Just (KeyPress _ sym _ _) -> do
            if
              | sym == xK_l -> do
                  mapM_ deselect (ringCur ring)
                  let ring' = ringAdv ring
                  mapM_ select (ringCur ring')
                  recur ring'
              | sym == xK_h -> do
                  mapM_ deselect (ringCur ring)
                  let ring' = ringPrev ring
                  mapM_ select (ringCur ring')
                  recur ring'
              | sym == xK_space || sym == xK_Return -> do
                  mapM_
                    (\(_, (_, (ws, _))) -> gotoWorkspace ws)
                    (ringCur ring)
                  return stop
              | otherwise -> return stop
      )
      ( ringFastForwardTo (\(_, (_, (_, n))) -> n == Current) $
          mkRing $
            zip dzenPipes rectsToWorkspaces
      )

  io $ do
    mapM_ hClose dzenPipes

  after
  where
    fi :: (Integral a, Num b) => a -> b
    fi = fromIntegral

    toInt :: (Integral a) => a -> Int
    toInt = fromIntegral

    deselect :: (Handle, (Rectangle, (String, ScreenType))) -> X ()
    deselect (pipe, (rect, (wsName, screenType))) = io $ hPutStrLn pipe wsName

    select :: (Handle, (Rectangle, (String, ScreenType))) -> X ()
    select (pipe, (rect, (wsName, screenType))) =
      io $ hPutStrLn pipe $ printf "^fg(#ff0)%s" wsName

    getWorkspaces :: WorkspaceWheelType -> X [(WorkspaceId, ScreenType)]
    getWorkspaces Existing = do
      (W.StackSet cur vis other _floating) <- gets windowset
      return $
        map (first W.tag) $
          sortBy (comparing (W.tag . fst)) $
            ((W.workspace cur, Current) : map ((,Visible) . W.workspace) vis)
              ++ filter
                (\(w, _) -> (not . null . W.integrate' . W.stack) w && W.tag w /= "*")
                (map (,Hidden) other)
    getWorkspaces (Alphabet isGreek b) =
      getWorkspacesForAlphabet
        ( map (if b then id else toLower) $
            if isGreek then "ΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ" else "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        )

    getWorkspaces Digits =
      getWorkspacesForAlphabet "0123456789"

    getWorkspacesForAlphabet a = do
      (W.StackSet cur vis other _floating) <- gets windowset
      let screenTypeById =
            Map.fromList $
              map (first (W.tag . W.workspace)) $
                (cur, Current) : map (,Visible) vis
      return $
        map
          (\c -> ([c], fromMaybe Hidden $ Map.lookup [c] screenTypeById))
          a

    -- Returns, starting x, dzen width
    calculatePositions :: Int -> Int -> Rectangle -> [a] -> [(Rectangle, a)]
    calculatePositions
      pointerX'
      pointerY'
      (Rectangle (toInt -> x) (toInt -> y) (toInt -> w) (toInt -> h))
      as =
        let radius :: (Num a) => a
            radius = 300

            nas = max 2 (length as)
            dPhi = 2 * pi / fromIntegral nas

            circDiff = sqrt $ (radius * cos dPhi - radius) ** 2 + (radius * sin dPhi) ** 2

            size :: (Integral a) => a
            size = round $ min 100 (circDiff * 0.5)

            sizeD2 :: (Integral a) => a
            sizeD2 = fi size `div` 2

            pointerX =
              max (radius + sizeD2 + x) $
                min pointerX' (x + w - radius - sizeD2)

            pointerY =
              max (radius + sizeD2 + y) $
                min pointerY' (y + h - radius - sizeD2)
         in zipWith
              ( \phi a ->
                  let xpos = radius * cos phi + fi pointerX
                      ypos = radius * sin phi + fi pointerY
                      r =
                        Rectangle
                          (round xpos - sizeD2)
                          (round ypos - sizeD2)
                          (fi size)
                          (fi size)
                   in (r, a)
              )
              [3 * pi / 2, 3 * pi / 2 + dPhi ..]
              as

    embedRectangle :: Rectangle -> Int -> Int -> Rectangle
    embedRectangle r@(Rectangle x y w h) w' h' | w' > fi w || h' > fi h = r
    embedRectangle (Rectangle x y w h) w' h' =
      Rectangle
        (x + fi (w `div` 2) - fi (w' `div` 2))
        (y + fi (h `div` 2) - fi (h' `div` 2))
        (fi w')
        (fi h')

    launchDzen :: Rectangle -> (String, ScreenType) -> X Handle
    launchDzen (Rectangle x y w h) (wsTag, screenType) = do
      pipe <-
        spawnPipe
          ( printf
              "dzen2 -fn \"Monofur Nerd Font:size=15\" -x %d -y %d -w %d -h %d -bg '#222' -fg '%s'"
              x
              y
              w
              h
              ( case screenType of
                  Current -> "#ff8888"
                  Visible -> "#8888ff"
                  Hidden -> "#aaa"
              )
          )

      io $ do
        hPutStrLn pipe wsTag
        hFlush pipe

      return pipe