{-# LANGUAGE ScopedTypeVariables #-}
module Internal.Marking where

import Internal.Windows (mapWindows, findWindow, getLocationWorkspace)
import XMonad
import XMonad.StackSet hiding (focus)
import Data.IORef
import Data.Map (Map)

import System.FilePath
import System.IO
import Control.Exception
import System.Environment

import qualified XMonad.Util.ExtensibleState as XS

import qualified Data.Map as Map

{- Submodule that handles marking windows so they can be jumped back to. -}

type Mark = Char

data MarkState =
  MarkState {
    markStateMap :: Map Mark Window
  , markLast :: Maybe Window
  } deriving (Read, Show)


instance ExtensionClass MarkState where
  initialValue = MarkState Map.empty Nothing

greedyFocus :: Window -> X ()
greedyFocus win = do
    ws <- withWindowSet $ \ss ->
            return $ getLocationWorkspace =<< findWindow ss win
    mapM_ (windows . greedyView . tag) ws
    focus win

markCurrentWindow :: Mark -> X ()
markCurrentWindow mark = do
  withFocused $ \win ->
    XS.modify $ \state@MarkState {markStateMap = ms} ->
        state {
          markStateMap = Map.insert mark win ms
        }

saveLastMark :: X ()
saveLastMark =
  withFocused $ \win ->
      XS.modify $ \state -> state { markLast = Just win }

jumpToLast :: X ()
jumpToLast = do
  m <- markLast <$> XS.get
  saveLastMark
  mapM_ greedyFocus m

jumpToMark :: Mark -> X ()
jumpToMark mark = do
  MarkState {markStateMap = m} <- XS.get
  case Map.lookup mark m of
    Nothing -> return ()
    Just w -> do
      saveLastMark
      greedyFocus w

setFocusedWindow :: a -> StackSet i l a s sd -> StackSet i l a s sd
setFocusedWindow
  window
  (StackSet (Screen (Workspace t l stack) a b) vis hid float) =
    let newStack =
          case stack of
            Nothing -> Nothing
            Just (Stack _ up down) -> Just (Stack window up down) in
    StackSet (Screen (Workspace t l newStack) a b) vis hid float

swapWithFocused :: (Ord a) => a -> StackSet i l a s sd -> StackSet i l a s sd
swapWithFocused winToSwap stackSet =
  case peek stackSet of
    Nothing -> stackSet
    Just focused -> do
       setFocusedWindow winToSwap $
          mapWindows (
            \w -> if w == winToSwap then focused else w) stackSet

swapWithLastMark :: X ()
swapWithLastMark = do
  m <- markLast <$> XS.get
  saveLastMark

  case m of
    Nothing -> return ()
    Just win -> windows $ swapWithFocused win

swapWithMark :: Mark -> X ()
swapWithMark mark = do
  MarkState {markStateMap = m} <- XS.get
  saveLastMark

  case Map.lookup mark m of
    Nothing -> return ()
    Just winToSwap ->
      windows $ swapWithFocused winToSwap