{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE NamedFieldPuns #-}

module XMonad.Hooks.EwmhDesktops.Extra where

import Control.Monad (when)
import Data.Maybe (fromMaybe)
import Data.Monoid (All)
import Data.Tuple.Extra (both)
import Graphics.X11.EWMH (getDesktopNames, setDesktopNames)
import Graphics.X11.EWMH.Atom (_NET_CURRENT_DESKTOP, _NET_DESKTOP_NAMES)
import Graphics.X11.Xlib.Display.Extra (withDefaultDisplay)
import XMonad hiding (workspaces)
import XMonad.Prelude ((!?)) -- this is part of >=base-4.19's Data.List
import XMonad.Actions.DynamicWorkspaces (addHiddenWorkspace, removeEmptyWorkspaceByTag)
import XMonad.StackSet (currentTag, greedyView, mapWorkspace, tag, workspaces)
import XMonad.Util.WorkspaceCompare (getSortByIndex)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified XMonad


ewmhExtra :: XConfig a -> IO (XConfig a)
ewmhExtra c = do
    -- XMonad.Hooks.EwmhDesktops.setDesktopViewport uses _NET_DESKTOP_VIEWPORT
    -- only if it exists.  This seems to be a harmless issue, but by creating
    -- the atom here, we suppress the error message:
    --
    --    xmonad: X11 error: BadAtom (invalid Atom parameter),
    --    request code=18, error code=5
    --
    _ <-
      withDefaultDisplay $ \dpy -> internAtom dpy "_NET_DESKTOP_VIEWPORT" False

    initialWorkspaces <-
      Data.Maybe.fromMaybe (XMonad.workspaces def)
        <$> withDefaultDisplay getDesktopNames

    return
      c { handleEventHook = ewmhDesktopsExtraEventHook <> handleEventHook c
        , rootMask = rootMask c .|. propertyChangeMask
        , XMonad.workspaces = initialWorkspaces
        }

ewmhDesktopsExtraEventHook :: Event -> X All
ewmhDesktopsExtraEventHook = \case

    -- Replace default non-greedy implementation of
    -- XMonad.Hooks.EwmhDesktops.ewmhDesktopsEventHook'
    ClientMessageEvent{ev_window, ev_message_type, ev_data = n : _}
          | ev_message_type == _NET_CURRENT_DESKTOP-> do
      r <- asks theRoot
      when (ev_window == r) $
        withWindowSet $ \s -> do
          sort <- getSortByIndex
          case (map tag . sort . workspaces $ s) !? fromIntegral n of
            Just t ->
              when (currentTag s /= t) $
                windows $ greedyView t
            _ ->
              trace $ "Bad _NET_CURRENT_DESKTOP with data=" <> show n
      mempty

    PropertyEvent{ev_window, ev_atom} -> do
      r <- asks theRoot
      when (ev_window == r && ev_atom == _NET_DESKTOP_NAMES) $
        withDisplay $ \dpy -> do
          sort <- getSortByIndex

          oldNames <- gets $ map tag . sort . workspaces . windowset
          newNames <- fromMaybe oldNames <$> io (getDesktopNames dpy)

          let
            (renamesFrom, renamesTo) = both Set.fromList $ unzip renames

            renames = go oldNames newNames where
              go old@(headOld : tailOld) new@(headNew : tailNew) = do
                let
                  deleteOld = Set.member headOld deleteNameSet
                  createNew = Set.member headNew createNameSet

                if
                  | headOld == headNew ->
                    -- assert (not deleteOld && not createNew)
                    go tailOld tailNew

                  | deleteOld && createNew ->
                    (headOld, headNew) :
                    go tailOld tailNew

                  | deleteOld ->
                    go tailOld new

                  | createNew ->
                    go old tailNew

                  | otherwise ->
                    -- assert (headOld == headNew)
                    go tailOld tailNew

              go _ _ = []

            oldNameSet = Set.fromList oldNames
            newNameSet = Set.fromList newNames
            deleteNameSet = Set.difference oldNameSet newNameSet
            createNameSet = Set.difference newNameSet oldNameSet

            deleteNames = Set.toAscList $
                            Set.difference deleteNameSet renamesFrom
            createNames = Set.toAscList $
                            Set.difference createNameSet renamesTo

          mapM_ addHiddenWorkspace createNames
          mapM_ removeEmptyWorkspaceByTag deleteNames
          when (not (null renames)) $ do
            let
              renameMap = Map.fromList renames
              rename w =
                case Map.lookup (tag w) renameMap of
                  Just newName -> w { tag = newName }
                  Nothing -> w

            modifyWindowSet $ mapWorkspace rename

          names <- gets $ map tag . sort . workspaces . windowset

          when (names /= newNames) $ do
            trace $ "setDesktopNames " <> show names
            io (setDesktopNames names dpy)

      mempty

    _ ->
      mempty