module Rahm.Desktop.Marking
  ( markCurrentWindow,
    jumpToMark,
    setAlternateWindows,
    getAlternateWindows,
    setAlternateWorkspace,
    getAlternateWorkspace,
    getMarkedLocations,
    markAllLocations,
    farLeftWindow,
    farRightWindow,
    windowLocation,
    markWindow,
    Mark,
  )
where

import Control.Monad.Trans (lift)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
import Data.List (sortOn)
import Data.List.Safe (head)
import Data.Map (Map)
import qualified Data.Map as Map
  ( empty,
    insert,
    insertWith,
    lookup,
  )
import Data.Maybe (catMaybes, fromMaybe)
import Data.Ord (Down (..))
import Rahm.Desktop.Common
  ( Location (Location),
    focusLocation,
    getCurrentWorkspace,
  )
import Rahm.Desktop.Logger (LogLevel (Debug), logs)
import Rahm.Desktop.StackSet
  ( Screen (Screen),
    Stack (Stack),
    StackSet (StackSet),
    Workspace (Workspace),
    allWindows,
    findTag,
    mapWindows,
    peek,
  )
import XMonad
  ( ExtensionClass (..),
    Rectangle (Rectangle),
    StateExtension (PersistentExtension),
    Window,
    WorkspaceId,
    X,
    catchX,
    getGeometry,
    io,
    withDisplay,
    withFocused,
    withWindowSet,
  )
import qualified XMonad.Util.ExtensibleState as XS (get, modify)
import Prelude hiding (head)

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

type Mark = String

data MarkState = MarkState
  { markStateMap :: Map Mark [Location],
    alternateWindows :: [Window],
    alternateWorkspaces :: Map Window WorkspaceId
  }
  deriving (Read, Show)

instance ExtensionClass MarkState where
  initialValue = MarkState Map.empty [] Map.empty
  extensionType = PersistentExtension

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

setAlternateWorkspace :: Window -> WorkspaceId -> X ()
setAlternateWorkspace win wid =
  XS.modify $ \m ->
    m
      { alternateWorkspaces = Map.insert win wid (alternateWorkspaces m)
      }

getAlternateWorkspace :: Window -> X (Maybe WorkspaceId)
getAlternateWorkspace window =
  Map.lookup window . alternateWorkspaces <$> XS.get

setAlternateWindows :: [Window] -> X ()
setAlternateWindows wins = XS.modify (\m -> m {alternateWindows = wins})

getAlternateWindows :: X [Window]
getAlternateWindows = alternateWindows <$> XS.get

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

markAllLocations :: Mark -> [Location] -> X ()
markAllLocations mark locs = do
  logs Debug "Marking locations %s as \"%s\"" (show locs) (show mark)

  XS.modify $ \m ->
    m
      { markStateMap = Map.insert mark locs (markStateMap m)
      }

markWindow :: Mark -> Window -> X ()
markWindow mark window = do
  logs Debug "Marking window %s as \"%s\"" (show window) (show mark)

  ws <- getCurrentWorkspace
  XS.modify $ \state@MarkState {markStateMap = ms} ->
    state
      { markStateMap = Map.insertWith (++) mark [Location ws $ Just window] ms
      }

markCurrentWindow :: Mark -> X ()
markCurrentWindow = withFocused . markWindow

jumpToMark :: Mark -> X ()
jumpToMark mark = do
  MarkState {markStateMap = m} <- XS.get
  mapM_ focusLocation $ head =<< Map.lookup mark m

getMarkedLocations :: Mark -> X [Location]
getMarkedLocations mark = do
  MarkState {markStateMap = m} <- XS.get
  return (fromMaybe [] $ Map.lookup mark m)

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

windowRect :: Window -> X (Maybe Rectangle)
windowRect win = withDisplay $ \dpy ->
  ( do
      (_, x, y, w, h, bw, _) <- io $ getGeometry dpy win
      return $ Just $ Rectangle x y (w + 2 * bw) (h + 2 * bw)
  )
    `catchX` return Nothing

getWindowsAndRects :: X [(Window, Rectangle)]
getWindowsAndRects =
  catMaybes
    <$> ( mapM (\w -> fmap (w,) <$> windowRect w)
            =<< withWindowSet (return . allWindows)
        )

windowLocation :: Window -> MaybeT X Location
windowLocation win = do
  tag <- MaybeT $ withWindowSet $ return . findTag win
  return (Location tag (Just win))

farLeftWindow :: MaybeT X Location
farLeftWindow = do
  rects <- lift $ sortOn (\(_, Rectangle x _ _ _) -> x) <$> getWindowsAndRects
  case rects of
    ((w, _) : _) -> windowLocation w
    _ -> MaybeT (return Nothing)

farRightWindow :: MaybeT X Location
farRightWindow = do
  rects <- lift $ sortOn (Down . \(_, Rectangle x _ _ _) -> x) <$> getWindowsAndRects
  case rects of
    ((w, _) : _) -> windowLocation w
    _ -> MaybeT (return Nothing)