module Rahm.Desktop.Marking (
  historyNext, historyPrev,
  markCurrentWindow, pushHistory,
  jumpToMark, jumpToLast, swapWithLastMark,
  swapWithMark
  ) where

import Rahm.Desktop.Windows (mapWindows, findWindow, getLocationWorkspace)
import XMonad
import XMonad.StackSet hiding (focus)
import Data.IORef
import Data.Map (Map)
import Control.Monad (when)

import System.FilePath
import System.IO
import Control.Exception
import System.Environment
import qualified Data.Sequence as Seq
import Data.Sequence (Seq(..))

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

historySize = 100 -- max number of history elements the tail.

data History a = History [a] (Seq a)
  deriving (Read, Show)

instance Default (History a) where
  
  def = History [] Seq.empty

seqPush :: a -> Seq a -> Seq a
seqPush elem s@(seq :|> _) | Seq.length s >= historySize = elem :<| seq
seqPush elem s = elem :<| s

historyForward :: History a -> History a
historyForward (History (a:as) tail) = History as (seqPush a tail)
historyForward z = z

historyBackward :: History a -> History a
historyBackward (History head (a :<| as)) = History (a : head) as
historyBackward z = z

historyCurrent :: History a -> Maybe a
historyCurrent (History (a:_) _) = Just a
historyCurrent _ = Nothing

historyPush :: (Eq a) => a -> History a -> History a
historyPush a h@(History (w : _) _) | a == w = h
historyPush a (History (w : _) tail) = History [a] (seqPush w tail)
historyPush a (History _ tail) = History [a] tail

historySwap :: History a -> History a
historySwap (History (a:as) (t :<| ts)) = History (t : as) (seqPush a ts)
historySwap z = z

historyLast :: History a -> Maybe a
historyLast (History _ (t :<| _)) = Just t
historyLast _ = Nothing

data Spot =
        WindowSpot Window | -- Focus is on a window.
        TagSpot String -- Focus is on an (empty) tag
        deriving (Read, Show, Eq, Ord)

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

data MarkState =
  MarkState {
    markStateMap :: Map Mark Window
  , windowHistory :: History Spot
  } deriving (Read, Show)


instance ExtensionClass MarkState where
  initialValue = MarkState Map.empty def
  extensionType = PersistentExtension

changeHistory :: (History Spot -> History Spot) -> (MarkState -> MarkState)
changeHistory fn ms = ms { windowHistory = fn (windowHistory ms)}

withMaybeFocused :: (Maybe Window -> X ()) -> X ()
withMaybeFocused f = withWindowSet $ f . peek

normalizeWindows :: X ()
normalizeWindows = do
  MarkState { windowHistory = h } <- XS.get
  mapM_ greedyFocus (historyCurrent h)

-- 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
        }

pushHistory :: X () -> X ()
pushHistory fn = do
  withMaybeFocused $ \maybeWindowBefore -> do
    case maybeWindowBefore of
      (Just windowBefore) ->
        XS.modify $ changeHistory (historyPush (WindowSpot windowBefore))
      Nothing -> 
        withWindowSet $ \ws ->
          XS.modify $ changeHistory (historyPush (TagSpot (currentTag ws)))

    fn

    withMaybeFocused $ \maybeWindowAfter ->
      case maybeWindowAfter of
        Just windowAfter ->
          XS.modify $ changeHistory (historyPush $ WindowSpot windowAfter)
        Nothing ->
          withWindowSet $ \ws ->
            XS.modify $ changeHistory (historyPush $ TagSpot $ currentTag ws)

withHistory :: (History Spot -> X ()) -> X ()
withHistory fn = do
  MarkState { windowHistory = w } <- XS.get
  fn w

jumpToLast :: X ()
jumpToLast = do
  XS.modify (changeHistory historySwap)
  normalizeWindows

jumpToMark :: Mark -> X ()
jumpToMark mark = do
  MarkState {markStateMap = m} <- XS.get
  case Map.lookup mark m of
    Nothing -> return ()
    Just w -> pushHistory $
      greedyFocus (WindowSpot 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 = pushHistory $ withHistory $ \hist -> do

  case historyLast hist of
    Just (WindowSpot win) ->
      windows $ swapWithFocused win
    Nothing -> return ()

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

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

historyPrev :: X ()
historyPrev = do
  XS.modify $ changeHistory historyBackward
  normalizeWindows

historyNext :: X ()
historyNext = do
  XS.modify $ changeHistory historyForward
  normalizeWindows