module Rahm.Desktop.Common
  ( focusLocation,
    masterWindow,
    windowsInWorkspace,
    getString,
    askWindowId,
    windowJump,
    withBorderColor,
    withBorderWidth,
    gotoWorkspace,
    moveLocationToWorkspace,
    getCurrentWorkspace,
    getCurrentLocation,
    runMaybeT_,
    Location (..),
  )
where

import Control.Monad (forM_, void, when)
import Control.Monad.Trans.Maybe (MaybeT (..))
import Data.Char (toLower)
import Data.List (concatMap, head, isInfixOf, map, (++))
import qualified Data.Map as Map (fromListWith)
import Data.Maybe (Maybe (..))
import Rahm.Desktop.DMenu (runDMenuPromptWithMap)
import qualified Rahm.Desktop.StackSet as S
  ( Screen (Screen, workspace),
    StackSet (StackSet, current),
    Workspace (Workspace, stack, tag),
    allWindows,
    focusWindow,
    greedyView,
    integrate',
    peek,
    shiftWin,
    workspaces,
  )
import Text.Printf (printf)
import XMonad
  ( Window,
    WorkspaceId,
    X,
    XConf (config, display),
    XConfig (focusedBorderColor, normalBorderColor),
    appName,
    asks,
    focus,
    io,
    refresh,
    runQuery,
    setWindowBorderWidth,
    setWindowBorderWithFallback,
    title,
    windows,
    withFocused,
    withWindowSet,
  )
import XMonad.Prompt (XPrompt (commandToComplete, showXPrompt))
import XMonad.Util.XUtils (pixelToString, stringToPixel)

-- A location is a workspace and maybe a window with that workspace.
data Location = Location
  { locationWorkspace :: WorkspaceId,
    locationWindow :: Maybe Window
  }
  deriving (Read, Show, Eq, Ord)

focusLocation :: Location -> X ()
focusLocation (Location ws Nothing) = windows $ S.greedyView ws
focusLocation (Location _ (Just win)) = windows $ S.focusWindow win

masterWindow :: MaybeT X Window
masterWindow = MaybeT $
  withWindowSet $ \ss ->
    let windows = (S.integrate' . S.stack . S.workspace . S.current) ss
     in case windows of
          (a : _) -> return $ Just a
          _ -> return Nothing

windowsInWorkspace :: WorkspaceId -> X [Location]
windowsInWorkspace wid =
  withWindowSet $
    return
      . concatMap
        ( \ws ->
            if S.tag ws == wid
              then map (Location wid . Just) $ S.integrate' (S.stack ws)
              else []
        )
      . S.workspaces

data WinPrompt = WinPrompt

instance XPrompt WinPrompt where
  showXPrompt _ = "[Window] "
  commandToComplete _ = id

getString :: Window -> X String
getString = runQuery $ do
  t <- title
  a <- appName
  return $
    if map toLower a `isInfixOf` map toLower t
      then t
      else printf "%s - %s" t a

askWindowId :: X (Maybe [Window])
askWindowId = do
  windowTitlesToWinId <- withWindowSet $ \ss ->
    Map.fromListWith (++) <$> mapM (\wid -> (,) <$> getString wid <*> return [wid]) (S.allWindows ss)

  runDMenuPromptWithMap "Window" (Just "#f542f5") windowTitlesToWinId

windowJump :: X ()
windowJump = mapM_ (focus . head) =<< askWindowId

-- Temporarily set the border color of the given windows.
withBorderColor :: String -> [Window] -> X a -> X a
withBorderColor color wins fn = do
  d <- asks display
  px <- stringToPixel d color
  oPx <- stringToPixel d =<< asks (normalBorderColor . config)
  fPx <- stringToPixel d =<< asks (focusedBorderColor . config)

  colorName <- io (pixelToString d px)
  oColorName <- io (pixelToString d oPx)
  fColorName <- io (pixelToString d fPx)

  forM_ wins $ \w ->
    setWindowBorderWithFallback d w colorName px

  ret <- fn

  withFocused $ \fw -> do
    forM_ wins $ \w ->
      when (w /= fw) $
        setWindowBorderWithFallback d w oColorName oPx

    setWindowBorderWithFallback d fw fColorName fPx

  return ret

withBorderWidth :: Int -> [Window] -> X a -> X a
withBorderWidth width ws fn = do
  d <- asks display

  forM_ ws $ \window ->
    io $ setWindowBorderWidth d window $ fromIntegral width

  ret <- fn

  forM_ ws $ \window ->
    io $ setWindowBorderWidth d window 2

  refresh

  return ret

gotoWorkspace :: WorkspaceId -> X ()
gotoWorkspace wid = windows $ S.greedyView wid

moveLocationToWorkspace :: Location -> WorkspaceId -> X ()
moveLocationToWorkspace (Location _ (Just win)) wid =
  windows $ S.shiftWin wid win
moveLocationToWorkspace _ _ = return ()

getCurrentWorkspace :: X WorkspaceId
getCurrentWorkspace = withWindowSet $
  \(S.StackSet (S.Screen (S.Workspace t _ _) _ _) _ _ _) -> do
    return t

getCurrentLocation :: X Location
getCurrentLocation = do
  ws <- getCurrentWorkspace
  win <- withWindowSet (return . S.peek)
  return (Location ws win)

runMaybeT_ :: (Monad m) => MaybeT m a -> m ()
runMaybeT_ = void . runMaybeT