{-# LANGUAGE DeriveDataTypeable #-}

module DynamicPerScreenWorkspaces ( WorkspaceScreens
                                  , DynamicWorkspaceConfig(..)
                                  , dynamicRescreenHook
                                  , 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

import System.IO


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
  io $ do
    file <- openFile "/tmp/log" AppendMode
    hPrint file "viewOn'"
    hPrint file wset
    hPrint file i
    hClose file
  
  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'
      
      io $ do
        file <- openFile "/tmp/log" AppendMode
        hPrint file "viewOn'+"
        hPrint file wset
        hPrint file mws
        hPrint file tag
        hClose file

      return $ case mws of
        Just ws -> W.view tag . W.view ws $ wset'
        Nothing -> wset'

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

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 . cleanup $ wset'

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

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

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