module Rahm.Desktop.Theater
  ( getCurrentTheaterName,
    getTheaters,
    saveTheater,
    saveCurrentTheater,
    restoreTheater,
    isStackSetEmpty,
    updateStateWithWindows,
  )
where

-- A "Theater" is the state of a stackset. One can save the current stackset as
-- a "theater" and then restore it later. If new windows were added, those new
-- windows are put into the hidden workspace.

import Data.Default (Default (..))
import Data.List (sortOn, (\\))
import Data.Map (Map)
import qualified Data.Map as Map (delete, insert, lookup, toList)
import Data.Maybe (isNothing)
import XMonad (X)
import qualified XMonad as X
  ( ExtensionClass (..),
    MonadState (..),
    ScreenDetail,
    ScreenId,
    StateExtension (..),
    Window,
    WindowSet,
    WorkspaceId,
    XState (windowset),
    readsLayout,
    windows,
    withWindowSet,
    rescreen
  )
import qualified XMonad.StackSet as W
  ( Screen (..),
    StackSet (..),
    Workspace (..),
    allWindows,
    differentiate,
    filter,
    integrate',
    mapLayout,
    mapWorkspace,
    tagMember,
    workspaces,
  )
import qualified XMonad.Util.ExtensibleState as XS (get, modify)

type WinSet = W.StackSet X.WorkspaceId String X.Window X.ScreenId X.ScreenDetail

newtype Theater = Theater WinSet
  deriving (Read, Show)

type TheaterName = Maybe String

data Theaters = Theaters
  { currentTheater :: TheaterName,
    theaters :: Map TheaterName Theater
  }
  deriving (Read, Show)

instance Default Theaters where
  def = Theaters Nothing mempty

instance X.ExtensionClass Theaters where
  initialValue = def
  extensionType = X.PersistentExtension

getCurrentTheaterName :: X TheaterName
getCurrentTheaterName = currentTheater <$> XS.get

getTheaters :: X [(TheaterName, Theater, Bool)]
getTheaters = do
  (Theaters cur theaters) <- XS.get
  return $
    sortOn (\(s, _, _) -> s) $
      map (\(name, theater) -> (name, theater, name == cur)) $
        filter (\(name, Theater ws) -> not (isStackSetEmpty "*" ws) || name == cur) $
          Map.toList theaters

saveTheater :: TheaterName -> X.WindowSet -> X ()
saveTheater name ws = do
  XS.modify $ \(Theaters cur map) ->
    Theaters cur $ do
      if isStackSetEmpty "*" ws && name /= cur
        then Map.delete name map
        else Map.insert name (Theater $ unboxLayout ws) map
  where
    unboxLayout = W.mapLayout show

saveCurrentTheater :: TheaterName -> X ()
saveCurrentTheater = X.withWindowSet . saveTheater

restoreTheater :: TheaterName -> X ()
restoreTheater name = do
  currentWindowset <- X.windowset <$> X.get
  currentTheaterName <- currentTheater <$> XS.get
  currentLayout <- W.layout . W.workspace . W.current . X.windowset <$> X.get

  saveTheater currentTheaterName currentWindowset

  newStackSet' <- do
    (Theaters cur mp) <- XS.get
    case Map.lookup name mp of
      Nothing -> do
        ws <- X.windowset <$> X.get
        return $
          W.mapWorkspace
            ( \(W.Workspace i l _) -> W.Workspace i l Nothing
            )
            ws
      Just (Theater ws) ->
        return $
          W.mapLayout
            ( \serialized -> case [x | (x, "") <- X.readsLayout currentLayout serialized] of
                [x] -> x
                [] -> currentLayout
                (_ : _) -> currentLayout
            )
            ws

  let newStackSet = updateStateWithWindows (W.allWindows currentWindowset) "*" newStackSet'

  saveTheater Nothing currentWindowset

  XS.modify $ \(Theaters _ m) -> Theaters name m
  saveTheater name newStackSet
  X.windows $ const newStackSet

  X.rescreen

isStackSetEmpty ::
  (Eq i, Eq a) => i -> W.StackSet i l a si sd -> Bool
isStackSetEmpty hiddenWorkspace =
  all
    ( \(W.Workspace t l s) -> isNothing s || t == hiddenWorkspace
    )
    . W.workspaces

updateStateWithWindows ::
  (Eq i, Eq a) => [a] -> i -> W.StackSet i l a si sd -> W.StackSet i l a si sd
updateStateWithWindows allWindows hiddenWorkspace ss =
  let missingWindows = allWindows \\ W.allWindows ss
      layout = W.layout $ W.workspace $ W.current ss
   in if null missingWindows
        then ss
        else
          if not (W.tagMember hiddenWorkspace ss)
            then
              ss
                { W.hidden =
                    W.Workspace hiddenWorkspace layout (W.differentiate missingWindows) :
                    W.hidden ss
                }
            else
              W.mapWorkspace
                ( \(W.Workspace t l s') ->
                    let s = W.filter (`elem` allWindows) =<< s'
                     in if t == hiddenWorkspace
                          then W.Workspace t l (W.differentiate $ W.integrate' s ++ missingWindows)
                          else W.Workspace t l s
                )
                ss