{-# LANGUAGE TypeOperators #-}

module Rahm.Desktop.Layout.PinWindow where

import Control.Arrow (Arrow (second))
import Control.Monad.Trans (lift)
import Control.Monad.Trans.Maybe (MaybeT (..))
import Data.Foldable (find)
import Data.List (nubBy)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe, maybe)
import Data.Semigroup (Endo (..))
import Data.Set (Set)
import qualified Data.Set as Set
import Rahm.Desktop.Common (runMaybeT_)
import Rahm.Desktop.Logger
import qualified Rahm.Desktop.StackSet as W
import XMonad
import qualified XMonad.StackSet as W (filter)
import qualified XMonad.Util.ExtensibleState as XS

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

instance Default PinWindowState where
  def = PinWindowState mempty

instance ExtensionClass PinWindowState where
  initialValue = def
  extensionType = PersistentExtension

newtype PinWindowLayout (l :: * -> *) (a :: *) = PinWindowLayout (l a)
  deriving (Show, Read)

instance (LayoutClass l a, sid ~ ScreenId, a ~ Window) => LayoutClass (PinWindowLayout l) a where
  runLayout (W.Workspace t (PinWindowLayout l) stack) rect = do
    (PinWindowState pinned) <- XS.get
    let allPinned = mconcat $ Map.elems pinned

    (W.StackSet cur vis _ floating) <- gets windowset

    (stack', winAndRect) <- case find ((== t) . (W.tag . W.workspace)) (cur : vis) of
      Just (W.Screen ws sid _) -> do
        let winsToMove = fromMaybe mempty (Map.lookup sid pinned)

        modifyWindowSet $
          appEndo $
            mconcat $ map (Endo . W.shiftWinNoFocus (W.tag ws)) (Set.toList winsToMove)

        updatedWorkspace <- withWindowSet $ return . W.findWorkspace t

        return
          ( maybe stack W.stack updatedWorkspace,
            map (second (scaleRationalRect rect)) $
              mapMaybe (\w -> (w,) <$> Map.lookup w floating) (Set.toList winsToMove)
          )
      Nothing -> return (stack, [])

    (rects, maybeNewLayout) <-
      runLayout (W.Workspace t l (W.filter (not . (`Map.member` floating)) =<< stack')) rect

    return
      ( winAndRect ++ rects,
        PinWindowLayout <$> maybeNewLayout
      )

pinWindow :: Window -> X ()
pinWindow win = runMaybeT_ $ do
  lift $ logs Debug "Pinning window %d" win

  ws@(W.StackSet cur vis _ _) <- gets windowset
  t <- hoist (W.findTag win ws)
  scr <- hoist $ find ((== t) . (W.tag . W.workspace)) (cur : vis)

  lift $
    XS.modify $ \(PinWindowState mp) ->
      PinWindowState $
        Map.alter
          (Just . maybe (Set.singleton win) (Set.insert win))
          (W.screen scr)
          mp
  where
    hoist = MaybeT . return

unpinWindow :: Window -> X ()
unpinWindow win = do
  XS.modify $ \(PinWindowState mp) ->
    PinWindowState $
      Map.map (Set.delete win) mp

toggleWindowPin :: Window -> X ()
toggleWindowPin win = do
  isPinned <- isWindowPinned win
  if isPinned
    then unpinWindow win
    else pinWindow win

isWindowPinned :: Window -> X Bool
isWindowPinned win = do
  (PinWindowState mp) <- XS.get
  return $ any (Set.member win) (Map.elems mp)