module Rahm.Desktop.StackSet
  ( masterWindow,
    findWorkspace,
    ensureWorkspace,
    swapWorkspaces,
    greedyView,
    shiftWin,
    screenRotateBackward,
    screenRotateForward,
    mapWindows,
    swapWindows,
    getLocationWorkspace,
    WindowLocation (..),
    windowMemberOfWorkspace,
    findWindow,
    module W,
  )
where

import Data.List (find)
import Data.List.Safe (head)
import qualified Data.Map as Map
  ( fromList,
    keys,
    lookup,
    mapKeys,
  )
import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
import XMonad.StackSet as W
  ( RationalRect (..),
    Screen (..),
    Stack (..),
    StackSet (..),
    Workspace (..),
    abort,
    allWindows,
    currentTag,
    delete,
    delete',
    differentiate,
    ensureTags,
    filter,
    findTag,
    float,
    focusDown,
    focusDown',
    focusMaster,
    focusUp,
    focusUp',
    focusWindow,
    index,
    insertUp,
    integrate,
    integrate',
    lookupWorkspace,
    mapLayout,
    mapWorkspace,
    member,
    modify,
    modify',
    new,
    peek,
    renameTag,
    screens,
    shift,
    shiftMaster,
    sink,
    swapDown,
    swapMaster,
    swapUp,
    tagMember,
    view,
    workspaces,
  )
import qualified XMonad.StackSet (shiftWin)
import Prelude hiding (head)

data WindowLocation i l a s sd
  = OnScreen (Screen i l a s sd)
  | OnHiddenWorkspace (Workspace i l a)
  | Floating

getLocationWorkspace :: WindowLocation i l a s sd -> Maybe (Workspace i l a)
getLocationWorkspace (OnScreen (Screen w _ _)) = Just w
getLocationWorkspace (OnHiddenWorkspace w) = Just w
getLocationWorkspace _ = Nothing

mapWindows :: (Ord a, Ord b) => (a -> b) -> StackSet i l a s sd -> StackSet i l b s sd
mapWindows fn (StackSet cur vis hid float) =
  StackSet
    (mapWindowsScreen cur)
    (map mapWindowsScreen vis)
    (map mapWindowsWorkspace hid)
    (Map.mapKeys fn float)
  where
    mapWindowsScreen (Screen work a b) = Screen (mapWindowsWorkspace work) a b
    mapWindowsWorkspace (Workspace t l stack) =
      Workspace t l (fmap (fmap fn) stack)

swapWindows :: (Ord a) => [(a, a)] -> StackSet i l a s d -> StackSet i l a s d
swapWindows toSwap = mapWindows $ \w ->
  fromMaybe w (Map.lookup w toSwapM)
  where
    toSwapM = Map.fromList (toSwap ++ map (\(a, b) -> (b, a)) toSwap)

masterWindow :: StackSet i l a s sd -> Maybe a
masterWindow = head . integrate' . stack . workspace . current

findWorkspace ::
  (Eq i) =>
  i ->
  StackSet i l a s sd ->
  Maybe (Workspace i l a)
findWorkspace wid = find ((== wid) . tag) . workspaces

ensureWorkspace ::
  (Eq i) =>
  i ->
  StackSet i l a s sd ->
  (StackSet i l a s sd, Workspace i l a)
ensureWorkspace t ss =
  case findWorkspace t ss of
    Nothing ->
      let ws = Workspace t (layout . workspace . current $ ss) Nothing
       in (ss {hidden = ws : hidden ss}, ws)
    Just ws -> (ss, ws)

swapWorkspaces ::
  (Eq i) =>
  i ->
  i ->
  StackSet i l a s sd ->
  StackSet i l a s sd
swapWorkspaces wid1 wid2 ss =
  let (ss', workspace1) = ensureWorkspace wid1 ss
      (ss'', workspace2) = ensureWorkspace wid2 ss'
   in mapWorkspace
        ( \w ->
            case () of
              _ | tag w == wid1 -> workspace2
              _ | tag w == wid2 -> workspace1
              _ -> w
        )
        ss''

greedyView :: (Eq i) => i -> StackSet i l a s sd -> StackSet i l a s sd
greedyView wid ss = swapWorkspaces (tag . workspace . current $ ss) wid ss

shiftWin :: (Ord a, Eq s, Eq i) => i -> a -> StackSet i l a s sd -> StackSet i l a s sd
shiftWin wid a = XMonad.StackSet.shiftWin wid a . fst . ensureWorkspace wid

screenRotateBackward :: W.StackSet i l a sid sd -> W.StackSet i l a sid sd
screenRotateBackward (W.StackSet current visible others floating) = do
  let screens = current : visible
      workspaces = tail $ cycle $ map W.workspace screens
      (current' : visible') = zipWith (\s w -> s {workspace = w}) screens workspaces
   in W.StackSet current' visible' others floating

screenRotateForward :: W.StackSet i l a sid sd -> W.StackSet i l a sid sd
screenRotateForward (W.StackSet current visible others floating) = do
  let screens = current : visible
      workspaces = rcycle $ map W.workspace screens
      (current' : visible') = zipWith (\s w -> s {workspace = w}) screens workspaces
   in W.StackSet current' visible' others floating
  where
    rcycle l = last l : l

{- Finds a Window and returns the screen its on and the workspace its on.
 - Returns nothing if the window doesn't exist.
 -
 - If the window is not a screen Just (Nothing, workspace) is returned.
 - If the window is a floating window Just (Nothing, Nothing) is returned. -}
findWindow ::
  (Eq a) => StackSet i l a s sd -> a -> Maybe (WindowLocation i l a s sd)
findWindow (StackSet cur vis hid float) win =
  listToMaybe . catMaybes $
    map findWindowScreen (cur : vis)
      ++ map findWindowWorkspace hid
      ++ [findWindowFloat]
  where
    findWindowScreen s@(Screen ws _ _) =
      if windowMemberOfWorkspace ws win
        then Just (OnScreen s)
        else Nothing

    findWindowWorkspace w =
      if windowMemberOfWorkspace w win
        then Just (OnHiddenWorkspace w)
        else Nothing

    findWindowFloat =
      if win `elem` Map.keys float
        then Just Floating
        else Nothing

windowMemberOfWorkspace :: (Eq a) => Workspace i l a -> a -> Bool
windowMemberOfWorkspace (Workspace _ _ s) w = w `elem` integrate' s