{-# LANGUAGE DeriveDataTypeable #-}

module DynamicPerScreenWorkspaces ( screenWorkspaceStorage
                                  , WorkspaceScreens
                                  , DynamicWorkspaceConfig(..)
                                  , dynamicRescreenHook
                                  , create
                                  , cleanup
                                  , focusWindow
                                  , view
                                  , viewOn
                                  , viewOnCurrent
                                  ) where

import XMonad hiding (hide)
import qualified XMonad.StackSet as W
import XMonad.Layout.LayoutModifier (LayoutModifier(..), ModifiedLayout(..))

import Control.Arrow ((&&&))
import Control.Monad
import qualified Data.Map as M
import Data.Maybe

import Storage


data DynamicWorkspaceStoreData = DWSD (M.Map Int ScreenId)
                                 deriving (Typeable, Show, Read)
instance StoreData DynamicWorkspaceStoreData

screenWorkspaceStorage :: (LayoutClass l a) => l a -> ModifiedLayout (Storage DynamicWorkspaceStoreData) l a
screenWorkspaceStorage = storage (DWSD M.empty)


type WorkspaceScreens = Int -> Int -> ScreenId

data DynamicWorkspaceConfig = DynamicWorkspaceConfig
                              { defaultWorkspaceScreen  :: !WorkspaceScreens
                              , workspaceTag            :: !(Int -> WorkspaceId)
                              }

defaultScreenWorkspaces :: DynamicWorkspaceConfig -> Int -> ScreenId -> [Int]
defaultScreenWorkspaces conf n s = filter ((== s) . defaultWorkspaceScreen conf n) [0..]


dynamicRescreenHook :: DynamicWorkspaceConfig -> [ScreenDetail] -> X ()
dynamicRescreenHook conf sds = do
  layout <- asks $ layoutHook . config
  windows $ cleanup' . \wset -> let (as,bs) = splitAt (length sds) $ map W.workspace (W.current wset:W.visible wset)
                         
                                    mkWS i = let tag = workspaceTag conf i
                                             in W.Workspace tag layout Nothing
                                    
                                    mkScreen sid (Just ws) sd = W.Screen ws sid sd
                                    mkScreen sid Nothing sd = W.Screen (mkWS $ nextEmpty' conf (length sds) sid wset) sid sd
                                    
                                    (s:ss) = zipWith3 mkScreen (map S [0..]) (map Just as ++ repeat Nothing) sds
                                in wset {W.current = s, W.visible = ss, W.hidden = bs ++ W.hidden wset}

workspaceScreen :: DynamicWorkspaceConfig -> Int -> X ScreenId
workspaceScreen conf i = do
  wset <- gets windowset
  d <- getStoreData
  
  let defaultScreen = defaultWorkspaceScreen conf (length $ W.screens wset) i
  return $ case d of
    Just (DWSD wsscreens) -> do
       case (M.lookup i wsscreens) of
         Just s  -> s
         Nothing -> defaultScreen
    _ ->
      defaultScreen

hide :: DynamicWorkspaceConfig -> Int -> WindowSet -> X WindowSet
hide conf i wset = gets (W.screens . windowset) >>= foldM (flip hideOn) wset
  where
    hideOn s | (workspaceTag conf i) == (W.tag $ W.workspace s) = viewNextEmpty conf (W.screen s)
    hideOn _ = return

cleanup' :: WindowSet -> WindowSet
cleanup' wset = wset { W.hidden = hidden' }
  where
    hidden' = filter (isJust . W.stack) $ W.hidden wset

viewOn' :: DynamicWorkspaceConfig -> ScreenId -> Int -> WindowSet -> X WindowSet
viewOn' conf sid i wset = do
  let tag = workspaceTag conf i
  case W.tagMember tag wset of
    False -> viewEmpty conf sid i wset
    True -> do
      wset' <- hide conf i wset
      let mws = W.lookupWorkspace sid wset'
      
      return $ case mws of
        Just ws -> W.view tag . W.view ws $ wset'
        Nothing -> wset'

create' :: DynamicWorkspaceConfig -> ScreenId -> Int -> WindowSet -> X WindowSet
create' conf sid i wset = do
  layout <- asks $ layoutHook . config
  let tag = workspaceTag conf i
      workspace = W.Workspace tag layout Nothing
  
  return $ if W.tagMember tag wset then wset else wset { W.hidden = workspace:(W.hidden wset) }

viewEmpty :: DynamicWorkspaceConfig -> ScreenId -> Int -> WindowSet -> X WindowSet
viewEmpty conf sid i = create' conf sid i >=> viewOn' conf sid i

nextEmpty' ::DynamicWorkspaceConfig -> Int -> ScreenId -> WindowSet -> Int
nextEmpty' conf n sid wset = fst $ head $ dropWhile ((`W.tagMember` wset) . snd) $ map (id &&& workspaceTag conf) $ defaultScreenWorkspaces conf n sid

nextEmpty ::  DynamicWorkspaceConfig -> ScreenId -> WindowSet -> Int
nextEmpty conf sid wset = nextEmpty' conf (length $ W.screens wset) sid wset

viewNextEmpty :: DynamicWorkspaceConfig -> ScreenId -> WindowSet -> X WindowSet
viewNextEmpty conf sid wset = viewEmpty conf sid (nextEmpty conf sid wset) wset

windows' :: (WindowSet -> X WindowSet) -> X ()
windows' f = do
  wset <- gets windowset
  wset' <- f wset
  windows . const $ wset'

create :: DynamicWorkspaceConfig -> Int -> X ()
create conf i = do
  sid <- workspaceScreen conf i
  windows' $ create' conf sid i

cleanup :: X ()
cleanup = windows cleanup'

focusWindow :: DynamicWorkspaceConfig -> Window -> X ()
focusWindow conf w = do
  wset <- gets windowset
  let mws = W.findTag w wset
  whenJust mws $ \ws -> do
    let i = head $ filter ((== ws) . workspaceTag conf) [0..]
    sid <- workspaceScreen conf i
    windows' $ liftM (cleanup' . W.focusWindow w) . viewOn' conf sid i

view :: DynamicWorkspaceConfig -> Int -> X ()
view conf i = do
  sid <- workspaceScreen conf i
  windows' $ liftM cleanup' . viewOn' conf sid i

viewOn :: DynamicWorkspaceConfig -> ScreenId -> Int -> X ()
viewOn conf sid i = do
  updateStoreData $ \(DWSD m) -> DWSD $ M.insert i sid m
  windows' $ liftM cleanup' . viewOn' conf sid i

viewOnCurrent :: DynamicWorkspaceConfig -> Int -> X ()
viewOnCurrent conf i = do
  sid <- gets (W.screen . W.current . windowset)
  viewOn conf sid i