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

import Data.Default
import Data.List (find)
import Data.List.Safe (head)
import qualified Data.Map as Map
import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
import XMonad.StackSet as W hiding (greedyView, shiftWin)
import qualified XMonad.StackSet
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