-- This module provides a wrapper around the Mosaic layout to create a more
-- consistent experience where instead of the windows being the ones it works
-- on, it instead works on the window places so things like window swapping
-- still work as expected.
module Rahm.Desktop.Layout.ConsistentMosaic
  ( doAlt,
    expandPositionAlt,
    shrinkPositionAlt,
    MosaicWrap(..)
  )
where

import qualified Data.Map as Map (fromList, lookup)
import Data.Maybe (mapMaybe)
import qualified Rahm.Desktop.StackSet as W
  ( Screen (..),
    Stack (..),
    StackSet (..),
    Workspace (..),
    integrate,
  )
import XMonad
  ( LayoutClass (description, handleMessage, runLayout),
    MonadState (..),
    Window,
    X,
    XState (..),
  )
import XMonad.Layout.MosaicAlt
  ( HandleWindowAlt,
    expandWindowAlt,
    shrinkWindowAlt,
  )

newtype MosaicWrap l a = MosaicWrap (l a) deriving (Read, Show)

doAlt :: (Window -> HandleWindowAlt) -> X HandleWindowAlt
doAlt f = do
  (W.StackSet (W.Screen (W.Workspace _ _ mStack) _ _) _ _ _) <-
    windowset <$> get

  return $
    case mStack of
      Nothing -> f 0
      Just (W.Stack _ u _) -> f (fromIntegral $ length u + 100)

expandPositionAlt :: X HandleWindowAlt
expandPositionAlt = doAlt expandWindowAlt

shrinkPositionAlt :: X HandleWindowAlt
shrinkPositionAlt = doAlt shrinkWindowAlt

instance (LayoutClass l a, Show a, Ord a, Enum a, Num a) => LayoutClass (MosaicWrap l) a where
  runLayout (W.Workspace t (MosaicWrap l) (id -> Just s)) rect = do
    let zs = zipStack [100 ..] s
        s' = fmap fst zs
        m = Map.fromList (W.integrate zs)

    (rects, maybeNewLayout) <- runLayout (W.Workspace t l (Just s')) rect
    let rects' =
          flip mapMaybe rects $ \(place, rect) ->
            (,rect) <$> Map.lookup place m

    return (rects', MosaicWrap <$> maybeNewLayout)
    where
      zipStack as (W.Stack b c d) =
        let (cz, bz : dz) = splitAt (length c) as
         in W.Stack (bz, b) (zip (reverse cz) c) (zip dz d)
  runLayout (W.Workspace t (MosaicWrap l) a) rect = do
    (rects, maybeNewLayout) <- runLayout (W.Workspace t l a) rect
    return (rects, MosaicWrap <$> maybeNewLayout)

  -- By default just pass the message to the underlying layout.
  handleMessage (MosaicWrap l) mess = do
    maybeNewLayout <- handleMessage l mess
    return (MosaicWrap <$> maybeNewLayout)

  description _ = "ConsistentMosaic"