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

import Control.Applicative ((<*))
import Control.Monad (forM_, void, when)
import Control.Monad.Trans.Class
import Control.Monad.Trans.Except (ExceptT (..), catchE, runExceptT, throwE)
import Control.Monad.Trans.Identity (IdentityT (..))
import Control.Monad.Trans.Maybe (MaybeT (..))
import Data.Char (toLower)
import Data.Either (either)
import Data.List (concatMap, head, isInfixOf, map, (++))
import qualified Data.Map as Map (fromListWith)
import Data.Maybe (Maybe (..), maybe)
import Data.Void (Void (..), absurd)
import Rahm.Desktop.DMenu (runDMenuPromptWithMap)
import Rahm.Desktop.Logger
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
  ( ScreenId,
    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

-- Have to add a definition because Stack uses an ancient version of
-- transformers for some reason.
myFinallyE :: (Monad m) => ExceptT e m a -> ExceptT e m () -> ExceptT e m a
myFinallyE m closer = catchE (m <* closer) (\e -> closer >> throwE e)

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

  (colorName, oColorName, fColorName) <-
    lift $
      (,,)
        <$> io (pixelToString d px)
        <*> io (pixelToString d oPx)
        <*> io (pixelToString d fPx)

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

  myFinallyE fn $
    lift $ do
      forM_ wins $ \w ->
        setWindowBorderWithFallback d w oColorName oPx
      withFocused $ \fw ->
        setWindowBorderWithFallback d fw fColorName fPx

withBorderColorM :: String -> [Window] -> MaybeT X a -> MaybeT X a
withBorderColorM s ws fn = toMaybeT $ withBorderColorE s ws (toExceptT fn)
  where
    toExceptT (MaybeT fn) = ExceptT $ maybe (Left ()) Right <$> fn
    toMaybeT (ExceptT fn) = MaybeT $ either (const Nothing) Just <$> fn

withBorderColor :: String -> [Window] -> X a -> X a
withBorderColor s ws fn =
  either absurd id <$> runExceptT (withBorderColorE s ws (lift fn))

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

getCurrentScreen :: X ScreenId
getCurrentScreen = withWindowSet $
  \(S.StackSet (S.Screen _ sid _) _ _ _) -> return sid

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