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

import Prelude hiding ((!!))

import Control.Monad.Trans.Maybe
import Control.Arrow (second, (&&&))
import qualified XMonad.StackSet as W
import XMonad

import Data.List.Safe ((!!))

import Rahm.Desktop.Common
import Rahm.Desktop.History
import XMonad.Actions.DynamicWorkspaces
import Data.List (sortOn, sort, sortBy, find)
import Data.Maybe (mapMaybe, fromMaybe)
import Data.Char (isUpper, toUpper, toLower, isAlphaNum)

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)

shiftToWorkspace :: WorkspaceId -> X ()
shiftToWorkspace t = do
  addHiddenWorkspace t
  windows . W.shift $ t


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

swapWorkspace :: WorkspaceId -> X ()
swapWorkspace toWorkspace = do
  addHiddenWorkspace toWorkspace
  windows $ \ss -> do
    let fromWorkspace = W.tag $ W.workspace $ W.current ss in
      W.StackSet (swapSc fromWorkspace toWorkspace $ W.current ss)
                (map (swapSc fromWorkspace toWorkspace) $ W.visible ss)
                (map (swapWs fromWorkspace toWorkspace) $ W.hidden ss)
                (W.floating ss)
  where
    swapSc fromWorkspace toWorkspace (W.Screen ws a b) =
        W.Screen (swapWs fromWorkspace toWorkspace ws) a b

    swapWs fromWorkspace toWorkspace ws@(W.Workspace t' l s)
        | t' == fromWorkspace = W.Workspace toWorkspace l s
        | t' == toWorkspace = W.Workspace fromWorkspace l s
        | otherwise = ws

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)