{-# LANGUAGE DeriveDataTypeable #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  XMonad.Actions.DynamicWorkspaceGroups
-- Copyright   :  (c) Brent Yorgey 2009
-- License     :  BSD-style (see LICENSE)
--
-- Maintainer  :  <byorgey@gmail.com>
-- Stability   :  experimental
-- Portability :  unportable
--
-- Dynamically manage \"workspace groups\", sets of workspaces being
-- used together for some common task or purpose, to allow switching
-- between workspace groups in a single action.  Note that this only
-- makes sense for multi-head setups.
--
-----------------------------------------------------------------------------

module XMonad.Actions.DynamicWorkspaceGroups
    ( -- * Usage
      -- $usage

      WSGroupId

    , addWSGroup
    , addCurrentWSGroup
    , forgetWSGroup
    , viewWSGroup

    , promptWSGroupView
    , promptWSGroupAdd
    , promptWSGroupForget

    , WSGPrompt
    ) where

import Data.List (find)
import Control.Arrow ((&&&))
import qualified Data.Map as M

import XMonad
import qualified XMonad.StackSet as W

import XMonad.Prompt
import qualified XMonad.Util.ExtensibleState as XS

-- $usage
-- You can use this module by importing it into your ~\/.xmonad\/xmonad.hs file:
--
-- > import XMonad.Actions.DynamicWorkspaceGroups
--
-- Then add keybindings like the following (this example uses
-- "XMonad.Util.EZConfig"-style keybindings, but this is not necessary):
--
-- >    , ("M-y n", promptWSGroupAdd myXPConfig "Name this group: ")
-- >    , ("M-y g", promptWSGroupView myXPConfig "Go to group: ")
-- >    , ("M-y d", promptWSGroupForget myXPConfig "Forget group: ")
--

type WSGroup = [(ScreenId,WorkspaceId)]

type WSGroupId = String

data WSGroupStorage = WSG { unWSG :: M.Map WSGroupId WSGroup }
  deriving (Typeable, Read, Show)

withWSG :: (M.Map WSGroupId WSGroup -> M.Map WSGroupId WSGroup) -> WSGroupStorage -> WSGroupStorage
withWSG f = WSG . f . unWSG

instance ExtensionClass WSGroupStorage where
  initialValue = WSG $ M.empty
  extensionType = PersistentExtension

-- | Add a new workspace group with the given name.
addWSGroup :: WSGroupId -> [WorkspaceId] -> X ()
addWSGroup name wids = withWindowSet $ \w -> do
  let wss  = map ((W.tag . W.workspace) &&& W.screen) $ W.screens w
      wmap = mapM (strength . (flip lookup wss &&& id)) wids
  case wmap of
    Just ps -> XS.modify . withWSG . M.insert name $ ps
    Nothing -> return ()
 where strength (ma, b) = ma >>= \a -> return (a,b)

-- | Give a name to the current workspace group.
addCurrentWSGroup :: WSGroupId -> X ()
addCurrentWSGroup name = withWindowSet $ \w ->
  addWSGroup name $ map (W.tag . W.workspace) (W.current w : W.visible w)

-- | Delete the named workspace group from the list of workspace
--   groups.  Note that this has no effect on the workspaces involved;
--   it simply forgets the given name.
forgetWSGroup :: WSGroupId -> X ()
forgetWSGroup = XS.modify . withWSG . M.delete

-- | View the workspace group with the given name.
viewWSGroup :: WSGroupId -> X ()
viewWSGroup name = do
  WSG m <- XS.get
  case M.lookup name m of
    Just grp -> mapM_ (uncurry viewWS) grp
    Nothing -> return ()

-- | View the given workspace on the given screen.
viewWS :: ScreenId -> WorkspaceId -> X ()
viewWS sid wid = do
  mw <- findScreenWS sid
  case mw of
    Just w -> do
      windows $ W.view w
      windows $ W.greedyView wid
    Nothing -> return ()

-- | Find the workspace which is currently on the given screen.
findScreenWS :: ScreenId -> X (Maybe WorkspaceId)
findScreenWS sid = withWindowSet $
  return . fmap (W.tag . W.workspace) . find ((==sid) . W.screen) . W.screens

data WSGPrompt = WSGPrompt String

instance XPrompt WSGPrompt where
  showXPrompt (WSGPrompt s) = s

-- | Prompt for a workspace group to view.
promptWSGroupView :: XPConfig -> String -> X ()
promptWSGroupView xp s = do
  gs <- fmap (M.keys . unWSG) XS.get
  mkXPrompt (WSGPrompt s) xp (mkComplFunFromList' gs) viewWSGroup

-- | Prompt for a name for the current workspace group.
promptWSGroupAdd :: XPConfig -> String -> X ()
promptWSGroupAdd xp s =
  mkXPrompt (WSGPrompt s) xp (const $ return []) addCurrentWSGroup

-- | Prompt for a workspace group to forget.
promptWSGroupForget :: XPConfig -> String -> X ()
promptWSGroupForget xp s = do
  gs <- fmap (M.keys . unWSG) XS.get
  mkXPrompt (WSGPrompt s) xp (mkComplFunFromList' gs) forgetWSGroup