module Rahm.Desktop.BorderColors where

import Control.Monad (when)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.Except (ExceptT (..), catchE, runExceptT, throwE)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
import Data.Foldable (forM_)
import Data.Map (Map)
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Map as Map
import Data.Void (absurd)
import Rahm.Desktop.Common (runMaybeT_)
import Rahm.Desktop.Hooks.WindowChange (StackChangeHook (StackChangeHook))
import qualified Rahm.Desktop.StackSet as W
import XMonad
import qualified XMonad.Util.ExtensibleState as XS
import XMonad.Util.Font (pixelToString, stringToPixel)

data BorderColor = BorderColor
  { focusColor :: String,
    normalColor :: String
  }
  deriving (Read, Show, Ord, Eq)

newtype BorderColorsState = BorderColorsState (Map Window BorderColor)
  deriving (Read, Show)

instance ExtensionClass BorderColorsState where
  initialValue = BorderColorsState mempty
  extensionType = PersistentExtension

stackChangeHook :: StackChangeHook
stackChangeHook =
  StackChangeHook
    ( \_ _ -> do
        mp <- gets mapped
        (BorderColorsState s) <- XS.get
        XS.put $
         BorderColorsState $
           Map.filterWithKey (\k _ -> k `Set.member` mp) s
        (BorderColorsState s) <- XS.get
        updateBorderColors $ Map.keys s
    )

updateBorderColors :: [Window] -> X ()
updateBorderColors windows = do
  (BorderColorsState mp) <- XS.get
  foc <- withWindowSet $ return . W.peek

  forM_ windows $ \win -> do
    (BorderColorsState m) <- XS.get

    dnc <- asks (normalBorderColor . config)
    dfc <- asks (focusedBorderColor . config)

    case Map.lookup win m of
      Just (BorderColor fc nc) ->
        let bc = if Just win == foc then fc else nc
         in setBorderColorRaw bc win
      Nothing ->
        let bc = if Just win == foc then dfc else dnc
         in setBorderColorRaw bc win

-- 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 :: BorderColor -> [Window] -> ExceptT e X a -> ExceptT e X a
withBorderColorE color wins fn = do
  cleanup <- lift (setBorderColor color wins)
  myFinallyE fn (lift cleanup)

-- Set the border color raw.
setBorderColorRaw :: String -> Window -> X ()
setBorderColorRaw color w = do
  d <- asks display
  px <- stringToPixel d color
  colorName <- io $ pixelToString d px
  setWindowBorderWithFallback d w colorName px

-- Set the border color for the given windows. This function returns another
-- function that should be used to clean up the border changes.
setBorderColor :: BorderColor -> [Window] -> X (X ())
setBorderColor border wins = do
  (BorderColorsState oldMap) <- XS.get

  XS.put $
    BorderColorsState $
      foldl (\m' win -> Map.insert win border m') oldMap wins
  updateBorderColors wins

  return $ do
    XS.modify $ \(BorderColorsState cur) ->
      BorderColorsState $
        foldl
          (flip $ Map.updateWithKey (\w _ -> Map.lookup w oldMap))
          cur
          wins
    updateBorderColors wins

resetBorderColor :: [Window] -> X ()
resetBorderColor wins = do
  XS.modify $ \(BorderColorsState mp) ->
    BorderColorsState $
      foldl (flip Map.delete) mp wins

  updateBorderColors  wins

withBorderColorM :: BorderColor -> [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 :: BorderColor -> [Window] -> X a -> X a
withBorderColor s ws fn =
  either absurd id <$> runExceptT (withBorderColorE s ws (lift fn))