StackSet.hs

module YMonad.Core.StackSet (
  Stack (..),
  Workspace (..),
  Screen (..),
  WindowSet (..),
  emptyWorkspace,
  singleton,
  view,
  shift,
  focusWindow,
  insertUp,
  delete,
  focusUp,
  focusDown,
  swapUp,
  swapDown,
  peek,
  integrate,
  allWindows,
  modifyCurrentWorkspace,
) where

import Brolude (until)
import Data.Map.Strict qualified as M
import YMonad.Core.Common (Rect)

data Stack a = Stack
  { stackFocus :: a
  , stackUp :: [a]
  , stackDown :: [a]
  }
  deriving stock (Eq, Ord, Show)

data Workspace i l a = Workspace
  { workspaceId :: i
  , workspaceLayout :: l
  , workspaceStack :: Maybe (Stack a)
  , workspaceRestoreLayout :: Maybe l
  }
  deriving stock (Eq, Ord, Show)

data Screen sid i l a = Screen
  { screenWorkspace :: Workspace i l a
  , screenId :: sid
  , screenRect :: Rect
  }
  deriving stock (Eq, Ord, Show)

data WindowSet sid i l a = WindowSet
  { current :: Screen sid i l a
  , visible :: [Screen sid i l a]
  , hidden :: [Workspace i l a]
  , floating :: M.Map a Rect
  }
  deriving stock (Eq, Ord, Show)

emptyWorkspace :: i -> l -> Workspace i l a
emptyWorkspace wid lay =
  Workspace
    { workspaceId = wid
    , workspaceLayout = lay
    , workspaceStack = Nothing
    , workspaceRestoreLayout = Nothing
    }

singleton :: Screen sid i l a -> [Workspace i l a] -> WindowSet sid i l a
singleton scr hs =
  WindowSet
    { current = scr
    , visible = []
    , hidden = hs
    , floating = M.empty
    }

modifyCurrentWorkspace :: (Workspace i l a -> Workspace i l a) -> WindowSet sid i l a -> WindowSet sid i l a
modifyCurrentWorkspace f ws = ws {current = (current ws) {screenWorkspace = f (screenWorkspace (current ws))}}

withCurrentStack :: b -> (Stack a -> b) -> WindowSet sid i l a -> b
withCurrentStack dflt f = maybe dflt f . workspaceStack . screenWorkspace . current

modifyZ :: Maybe (Stack a) -> (Stack a -> Maybe (Stack a)) -> WindowSet sid i l a -> WindowSet sid i l a
modifyZ dflt f ws =
  ws
    { current =
        (current ws)
          { screenWorkspace =
              (screenWorkspace (current ws))
                { workspaceStack = withCurrentStack dflt f ws
                }
          }
    }

modifyZ' :: (Stack a -> Stack a) -> WindowSet sid i l a -> WindowSet sid i l a
modifyZ' f = modifyZ Nothing (Just . f)

peek :: WindowSet sid i l a -> Maybe a
peek = withCurrentStack Nothing (Just . stackFocus)

integrate :: Stack a -> [a]
integrate st = reverse (stackUp st) ++ stackFocus st : stackDown st

allWindows :: WindowSet sid i l a -> [a]
allWindows ws = concatMap workspaceWindows (workspaces ws)
  where
    workspaceWindows wk = maybe [] integrate (workspaceStack wk)

workspaces :: WindowSet sid i l a -> [Workspace i l a]
workspaces ws = screenWorkspace (current ws) : map screenWorkspace (visible ws) ++ hidden ws

member :: (Eq a) => a -> WindowSet sid i l a -> Bool
member win ws = win `elem` allWindows ws

workspaceMember :: (Eq i) => i -> WindowSet sid i l a -> Bool
workspaceMember wid ws = wid `elem` map workspaceId (workspaces ws)

findWorkspaceIdOfWindow :: (Eq a) => a -> WindowSet sid i l a -> Maybe i
findWorkspaceIdOfWindow win ws = workspaceId <$> find containsWindow (workspaces ws)
  where
    containsWindow wk = maybe False (elem win . integrate) (workspaceStack wk)

view :: (Eq i) => i -> WindowSet sid i l a -> WindowSet sid i l a
view wid ws
  | workspaceId (screenWorkspace (current ws)) == wid = ws
  | otherwise =
      case break ((== wid) . workspaceId . screenWorkspace) (visible ws) of
        (before, target : after) ->
          ws
            { current = target
            , visible = current ws : before ++ after
            }
        (_, []) ->
          case break ((== wid) . workspaceId) (hidden ws) of
            (before, target : after) ->
              let oldCurrent = screenWorkspace (current ws)
               in ws
                    { current = (current ws) {screenWorkspace = target}
                    , hidden = before ++ [oldCurrent] ++ after
                    }
            (_, []) -> ws

insertUp :: (Eq a) => a -> WindowSet sid i l a -> WindowSet sid i l a
insertUp win ws
  | member win ws = ws
  | otherwise = modifyZ (Just (Stack win [] [])) (\(Stack focus up down) -> Just (Stack win up (focus : down))) ws

focusWindow :: (Eq a, Eq i) => a -> WindowSet sid i l a -> WindowSet sid i l a
focusWindow win ws
  | Just win == peek ws = ws
  | otherwise =
      fromMaybe ws $ do
        wid <- findWorkspaceIdOfWindow win ws
        pure $ until ((Just win ==) . peek) focusUp (view wid ws)

focusUp :: WindowSet sid i l a -> WindowSet sid i l a
focusUp = modifyZ' focusUpStack

focusDown :: WindowSet sid i l a -> WindowSet sid i l a
focusDown = modifyZ' focusDownStack

swapUp :: WindowSet sid i l a -> WindowSet sid i l a
swapUp = modifyZ' swapUpStack

swapDown :: WindowSet sid i l a -> WindowSet sid i l a
swapDown = modifyZ' swapDownStack

delete :: (Ord a) => a -> WindowSet sid i l a -> WindowSet sid i l a
delete win = sink win . delete' win

delete' :: (Eq a) => a -> WindowSet sid i l a -> WindowSet sid i l a
delete' win ws =
  ws
    { current = removeFromScreen (current ws)
    , visible = map removeFromScreen (visible ws)
    , hidden = map removeFromWorkspace (hidden ws)
    }
  where
    removeFromWorkspace wk = wk {workspaceStack = workspaceStack wk >>= filterStack (/= win)}
    removeFromScreen scr = scr {screenWorkspace = removeFromWorkspace (screenWorkspace scr)}

sink :: (Ord a) => a -> WindowSet sid i l a -> WindowSet sid i l a
sink win ws = ws {floating = M.delete win (floating ws)}

filterStack :: (a -> Bool) -> Stack a -> Maybe (Stack a)
filterStack p (Stack focus up down) =
  case filter p (focus : down) of
    focus' : down' -> Just (Stack focus' (filter p up) down')
    [] ->
      case filter p up of
        focus' : up' -> Just (Stack focus' up' [])
        [] -> Nothing

focusUpStack :: Stack a -> Stack a
focusUpStack st = case stackUp st of
  up : ups -> Stack up ups (stackFocus st : stackDown st)
  [] ->
    case reverse (stackFocus st : stackDown st) of
      [] -> st
      focus : rest -> Stack focus rest []

focusDownStack :: Stack a -> Stack a
focusDownStack = reverseStack . focusUpStack . reverseStack

swapUpStack :: Stack a -> Stack a
swapUpStack st = case stackUp st of
  up : ups -> Stack (stackFocus st) ups (up : stackDown st)
  [] -> Stack (stackFocus st) (reverse (stackDown st)) []

swapDownStack :: Stack a -> Stack a
swapDownStack = reverseStack . swapUpStack . reverseStack

reverseStack :: Stack a -> Stack a
reverseStack st = Stack (stackFocus st) (stackDown st) (stackUp st)

shift :: (Eq a, Eq i) => i -> WindowSet sid i l a -> WindowSet sid i l a
shift wid ws = case peek ws of
  Nothing -> ws
  Just win -> shiftWin wid win ws

shiftWin :: (Eq a, Eq i) => i -> a -> WindowSet sid i l a -> WindowSet sid i l a
shiftWin wid win ws = case findWorkspaceIdOfWindow win ws of
  Just from
    | workspaceMember wid ws && wid /= from -> move from ws
  _ -> ws
  where
    move from = onWorkspace wid (insertUp win) . onWorkspace from (delete' win)

onWorkspace :: (Eq i) => i -> (WindowSet sid i l a -> WindowSet sid i l a) -> WindowSet sid i l a -> WindowSet sid i l a
onWorkspace wid f ws = view currentWid (f (view wid ws))
  where
    currentWid = workspaceId (screenWorkspace (current ws))