module Rahm.Desktop.PinWindow where

import Control.Monad.RWS (Endo (Endo, appEndo), forM_)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Rahm.Desktop.Common (setBorderColor)
import Rahm.Desktop.Hooks.WindowChange (StackChangeHook (..))
import qualified Rahm.Desktop.StackSet as W
import XMonad (Default (..), ExtensionClass (..), StateExtension (PersistentExtension), Window, X, XState (windowset), getScreenSaver)
import qualified XMonad as X
import qualified XMonad.Util.ExtensibleState as XS

newtype PinWindowState = PinWindowState (Set Window)
  deriving (Read, Show)

instance Default PinWindowState where
  def = PinWindowState mempty

instance ExtensionClass PinWindowState where
  initialValue = def
  extensionType = PersistentExtension

focPinColor = "#40ff40"

pinColor = "#80a080"

togglePinWindow :: Window -> X ()
togglePinWindow w = do
  foc <- X.withWindowSet (return . W.peek)
  (fc, nc) <-
    (,)
      <$> X.asks (X.focusedBorderColor . X.config)
      <*> X.asks (X.normalBorderColor . X.config)

  (PinWindowState s) <- XS.get
  let (del, s') =
          ( if Set.member w s
              then (True, Set.delete w s)
              else (False, Set.insert w s)
          )

      bc = case (Just w == foc, del) of
        (True, True) -> fc
        (False, True) -> nc
        (False, False) -> pinColor
        (True, False) -> focPinColor

  _ <- setBorderColor bc [w]
  XS.put (PinWindowState s')

pinWindowBorderUpdate :: X ()
pinWindowBorderUpdate = do
  (PinWindowState s) <- XS.get
  foc <- X.withWindowSet (return . W.peek)
  forM_ (Set.toList s) $ \w -> do
    let bc = if Just w == foc then focPinColor else pinColor
    _ <- setBorderColor bc [w]
    return ()

pinWindowChangeHook :: StackChangeHook
pinWindowChangeHook = StackChangeHook $ \last current -> do
  (PinWindowState s) <- XS.get
  let lastVisible = visiblePinnedWindows s last
      currentVisible = visiblePinnedWindows s current
      diff = Map.difference lastVisible currentVisible
      scrUpdates = Map.elems $ screenUpdates last current

  X.windows $
    foldl
      ( \fn (ow, nw) ->
          let l =
                map (Endo . W.shiftWinNoFocus nw) $
                  filter (`Set.member` s) (W.windowsOnWorkspace ow last)
           in appEndo (mconcat l) . fn
      )
      id
      scrUpdates
  where
    screenUpdates :: (Ord si, Eq i) => W.StackSet i l a si sd -> W.StackSet i l a si sd -> Map si (i, i)
    screenUpdates (W.StackSet c1 v1 _ _) (W.StackSet c2 v2 _ _) =
      let makeMap = Map.fromList . map (\s -> (W.screen s, W.tag $ W.workspace s))
          m1 = makeMap (c1 : v1)
          m2 = makeMap (c2 : v2)
       in foldl
            ( \m' (k, w1) ->
                case Map.lookup k m2 of
                  Just w2 | w1 == w2 -> m'
                  Just w2 -> Map.insert k (w1, w2) m'
                  _ -> Map.insert k (w1, w1) m'
            )
            mempty
            (Map.assocs m1)

    visiblePinnedWindows :: (Ord a) => Set a -> W.StackSet i l a si sd -> Map a si
    visiblePinnedWindows pinned (W.StackSet cur vis _ _) =
      Map.fromList $
        flip concatMap (cur : vis) $ \scr -> do
          let onScr = filter (`Set.member` pinned) $ W.integrate' $ W.stack $ W.workspace scr
           in map (,W.screen scr) onScr