-- Common ways to select workspaces
module Rahm.Desktop.Workspaces where

import Control.Arrow (second, (&&&))
import Control.Monad.Trans.Maybe
import Data.Char (isAlphaNum, isUpper, toLower, toUpper)
import Data.List (find, sort, sortBy, sortOn)
import Data.List.Safe ((!!))
import Data.Maybe (fromMaybe, mapMaybe)
import Rahm.Desktop.Common
import Rahm.Desktop.History
import qualified Rahm.Desktop.StackSet as W
import XMonad
import Prelude hiding ((!!))

newtype Selector = Selector (forall a. (a -> Bool) -> [a] -> Maybe a)

data WorkspaceState = Current | Hidden | Visible
  deriving (Ord, Eq, Enum)

-- Returns all the workspaces that are either visible, current or Hidden but
-- have windows and that workspace's state.
--
-- In other words, filters out workspaces that have no windows and are not
-- visible.
--
-- This function will sort the result by the workspace tag.
getPopulatedWorkspaces ::
  W.StackSet String l a sid sd -> [(WorkspaceState, W.Workspace String l a)]
getPopulatedWorkspaces (W.StackSet (W.Screen cur _ _) vis hi _) =
  filter ((/= "*") . W.tag . snd) $
    sortOn (W.tag . snd) $
      mapMaybe (\w@(W.Workspace _ _ s) -> fmap (const (Hidden, w)) s) hi
        ++ map (\(W.Screen w _ _) -> (Visible, w)) vis
        ++ [(Current, cur)]

next :: Selector
next = Selector $ \f l -> select f l l
  where
    select f (x : y : xs) _ | f x = Just y
    select f [x] (y : _) | f x = Just y
    select f (x : xs) orig = select f xs orig
    select f _ _ = Nothing

prev :: Selector
prev = Selector $ \f l ->
  let (Selector fn) = next in fn f (reverse l)

lastWorkspaceId :: X WorkspaceId
lastWorkspaceId =
  W.tag . snd . last <$> withWindowSet (return . getPopulatedWorkspaces)

firstWorkspaceId :: X WorkspaceId
firstWorkspaceId =
  W.tag . snd . head <$> withWindowSet (return . getPopulatedWorkspaces)

windowsInCurrentWorkspace :: X [Window]
windowsInCurrentWorkspace = withWindowSet $
  \(W.StackSet (W.Screen (W.Workspace _ _ s) _ _) _ _ _) -> do
    return $ W.integrate' s

getHorizontallyOrderedScreens ::
  W.StackSet wid l a ScreenId ScreenDetail ->
  [(Bool, W.Screen wid l a ScreenId ScreenDetail)]
-- ^ Returns a list of screens ordered from leftmost to rightmost.
getHorizontallyOrderedScreens windowSet =
  flip sortBy screens $ \sc1 sc2 ->
    let (SD (Rectangle x1 _ _ _)) = W.screenDetail (snd sc1)
        (SD (Rectangle x2 _ _ _)) = W.screenDetail (snd sc2)
     in x1 `compare` x2
  where
    screens = (True, W.current windowSet) : map (False,) (W.visible windowSet)

accompaningWorkspace :: WorkspaceId -> WorkspaceId
accompaningWorkspace [s] =
  return $
    if isUpper s
      then toLower s
      else toUpper s
accompaningWorkspace s = s

adjacentWorkspaceNotVisible :: Selector -> WorkspaceId -> X WorkspaceId
adjacentWorkspaceNotVisible (Selector selector) from =
  withWindowSet $ \ss ->
    let tags =
          sort $
            W.tag . snd
              <$> filter
                (\x -> fst x /= Visible)
                ( getPopulatedWorkspaces ss
                )
     in return $ fromMaybe from $ selector (== from) tags

adjacentWorkspace :: Selector -> WorkspaceId -> X WorkspaceId
adjacentWorkspace (Selector selector) from =
  withWindowSet $ \ss ->
    let tags = sort $ W.tag . snd <$> getPopulatedWorkspaces ss
     in return $ fromMaybe from $ selector (== from) tags

viewAdjacent :: Selector -> X ()
viewAdjacent sel =
  gotoWorkspace =<< (adjacentWorkspaceNotVisible sel =<< getCurrentWorkspace)

adjacentScreen :: Selector -> X WorkspaceId
adjacentScreen (Selector f) = do
  (screens, current) <-
    withWindowSet $ return . (getHorizontallyOrderedScreens &&& W.current)

  return $ W.tag $ W.workspace $ maybe current snd (f fst screens)

withScreen :: (WorkspaceId -> WindowSet -> WindowSet) -> Int -> X ()
withScreen fn n = do
  windows $ \windowSet ->
    case map snd (getHorizontallyOrderedScreens windowSet) !! n of
      Nothing -> windowSet
      Just screen -> fn (W.tag $ W.workspace screen) windowSet

workspaceWithWindow :: Window -> X (Maybe WorkspaceId)
workspaceWithWindow wid = withWindowSet $ \(W.StackSet c v h _) ->
  return $
    W.tag
      <$> find
        (\(W.Workspace _ _ stack) -> wid `elem` W.integrate' stack)
        (map W.workspace (c : v) ++ h)