module Rahm.Desktop.BorderColors where

import Control.Monad (when)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.Except (ExceptT (..), catchE, runExceptT, throwE)
import Data.Foldable (forM_)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Void (absurd)
import Rahm.Desktop.Common (Xish (..), runMaybeT_)
import Rahm.Desktop.Hooks.WindowChange (StackChangeHook (StackChangeHook))
import Rahm.Desktop.Logger (LogLevel (Debug), logs)
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

withBorderColor :: (Xish x, Monad x) => BorderColor -> [Window] -> x a -> x a
withBorderColor s ws fn = do
  cleanup <- liftFromX (setBorderColor s ws)
  finallyX fn $ liftFromX cleanup