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

import Control.Monad.Writer
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 Text.Printf (printf)
import XMonad (Rectangle (..), ScreenDetail (..), WindowSet)
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)

ensureWorkspaces ::
  (Eq i) =>
  [i] ->
  StackSet i l a s sd ->
  (StackSet i l a s sd, [Workspace i l a])
ensureWorkspaces (t : ts) ss =
  let (ss', w) = ensureWorkspace t ss
      (ss'', ws) = ensureWorkspaces ts ss'
   in (ss'', w : ws)
ensureWorkspaces [] ss = (ss, [])

swapWorkspaces ::
  (Eq i) =>
  i ->
  i ->
  StackSet i l a s sd ->
  StackSet i l a s sd
swapWorkspaces tag1 tag2 =
  W.mapWorkspace
    ( \(W.Workspace t a b) ->
        W.Workspace
          ( case (t == tag1, t == tag2) of
              (True, False) -> tag2
              (False, True) -> tag1
              _ -> t
          )
          a
          b
    )

switchWorkspaces ::
  (Eq i) =>
  i ->
  i ->
  StackSet i l a s sd ->
  StackSet i l a s sd
switchWorkspaces t1 t2 (ensureWorkspaces [t1, t2] -> (ss, [w1, w2])) =
  W.mapWorkspace
    ( \case
        (Workspace t _ _) | t == t1 -> w2
        (Workspace t _ _) | t == t2 -> w1
        w -> w
    )
    ss

dbgStackSet :: WindowSet -> String
dbgStackSet ws@(W.StackSet cur vis hidden _) = execWriter $ do
  tell "* " >> logScreen cur >> tell "\n"
  mapM_ (\s -> tell "  " >> logScreen s >> tell "\n") vis

  mapM_ logWorkspace (W.workspaces ws)
  where
    logWorkspace (Workspace tag _ st) = do
      tell $ printf "WS %s\n" tag
      forM_ st $ \(Stack foc up down) -> do
        mapM_ (tell . printf "    %d\n") up
        tell $ printf "  * %d\n" foc
        mapM_ (tell . printf "    %d\n") down

    logScreen (Screen ws sid (SD (Rectangle _ _ w h))) = do
      tell (printf "id=%s (%sx%s) - [%s]" (show sid) (show w) (show h) (W.tag ws))

greedyView :: (Eq i) => i -> StackSet i l a s sd -> StackSet i l a s sd
greedyView wid ss = switchWorkspaces (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

windowsOnWorkspace :: (Eq i) => i -> StackSet i l a s sd -> [a]
windowsOnWorkspace i ss = fromMaybe [] $ do
  ws <- find ((== i) . W.tag) (W.workspaces ss)
  s <- W.stack ws
  return (W.integrate s)

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