{-# LANGUAGE TypeOperators #-}

module Rahm.Desktop.Layout.PinWindow where

import Control.Arrow (Arrow (second))
import Control.Exception (throw)
import Control.Monad (unless, when)
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 [(Window, W.RationalRect)])
  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
    -- Clean up window id's thare should not be pinned anymore because the
    -- windows were unmanaged.
    (PinWindowState pinned') <- XS.get
    ws <- gets windowset
    XS.put $ PinWindowState (cleanupPinned pinned' ws)
    (PinWindowState pinned) <- XS.get

    let myScreen = find ((== t) . W.tag . W.workspace) (W.screens ws)
        pinnedRects =
          maybe
            []
            (map $ second (scaleRationalRect rect))
            (((`Map.lookup` pinned) . W.screen) =<< myScreen)
        allPinned = Set.fromList $ map fst $ concat $ Map.elems pinned
        pinnedOnMyScreen = map fst pinnedRects

        windowsToMove =
          filter (\win -> W.findTag win ws /= Just t) pinnedOnMyScreen

    modifyWindowSet $
      appEndo . mconcat $
        -- Move the windows that are supposed to now be on this workspace here
        map (Endo . W.shiftWinNoFocus t) windowsToMove
          ++
          -- once again, sink the windows that are supposed to be here. Make
          -- sure they don't float.
          map (Endo . W.sink) (Set.toList allPinned)

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

    return (pinnedRects ++ rects, PinWindowLayout <$> maybeNewLayout)
    where
      cleanupPinned mp ss =
        let aw = Set.fromList (W.allWindows ss)
         in Map.map (filter ((`Set.member` aw) . fst)) mp

  handleMessage (PinWindowLayout l) a = do
    maybeNewLayout <- handleMessage l a
    return (PinWindowLayout <$> maybeNewLayout)

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

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

  lift $ do
    XS.modify $ \(PinWindowState mp) ->
      PinWindowState $
        Map.alter
          (Just . maybe [(win, rect)] ((win, rect) :))
          (W.screen scr)
          mp

    -- Don't float the window anymore.
    modifyWindowSet $ W.sink win
  where
    hoist = MaybeT . return

unpinWindow :: Window -> X ()
unpinWindow win = runMaybeT_ $ do
  (PinWindowState mp) <- lift XS.get
  (win, rect) <- hoist $ find ((== win) . fst) (concat $ Map.elems mp)

  lift $ do
    XS.put $
      PinWindowState $
        Map.map (filter ((/= win) . fst)) mp

    -- refloat the window.
    modifyWindowSet $ W.float win rect
  where
    hoist = MaybeT . return

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 (any $ (== win) . fst) (Map.elems mp)

pinnedWindows :: X [Window]
pinnedWindows = do
  (PinWindowState s) <- XS.get
  return $ map fst $ concat $ Map.elems s

pinnable :: l a -> PinWindowLayout l a
pinnable = PinWindowLayout

-- Unpins the window, executes the action, then repins the window. Useful for
-- window shifts and whatnot.
withWindowsUnpinned :: [Window] -> X () -> X ()
withWindowsUnpinned wins fn =
  ( do
      mapM_ unpinWindow wins
      fn
      mapM_ pinWindow wins
  )
    `catchX` mapM_ pinWindow wins