module Wetterhorn.StackSet where

import Control.Monad.Identity
import Control.Monad.Writer (First (..), MonadWriter (tell), execWriter)
import Data.Maybe (isJust)

-- | The root datastructure for holding the state of the windows.
data StackSet s sd t l a = StackSet
  { -- | The currently selected screen.
    current :: Screen s sd t l a,
    -- | Remaining visible screens.
    visible :: [Screen s sd t l a],
    -- | Workspaces that exist, but are not on a screen.
    hidden :: [Workspace t l a]
  }
  deriving (Read, Show, Eq, Ord)

class TraverseWorkspace f where
  traverseWorkspaces ::
    (Applicative m) => (Workspace t l a -> m (Workspace t' l' a')) -> f t l a -> m (f t' l' a')

traverseWorkspaces_ :: (TraverseWorkspace f, Monad m) => (Workspace t l a -> m ()) -> f t l a -> m ()
traverseWorkspaces_ f = void . traverseWorkspaces (\w -> f w >> pure w)

foldMapWorkspaces ::
  (Monoid m, TraverseWorkspace f) => (Workspace t l a -> m) -> f t l a -> m
foldMapWorkspaces fn = execWriter . traverseWorkspaces_ (tell . fn)

mapWorkspaces ::
  (TraverseWorkspace f) =>
  (Workspace t l a -> Workspace t' l' a') ->
  f t l a ->
  f t' l' a'
mapWorkspaces fn = runIdentity . traverseWorkspaces (pure . fn)

instance TraverseWorkspace Workspace where
  traverseWorkspaces f = f

instance TraverseWorkspace (Screen s sd) where
  traverseWorkspaces f scr = (\w' -> scr {workspace = w'}) <$> f (workspace scr)

instance TraverseWorkspace (StackSet s sd) where
  traverseWorkspaces f (StackSet cur vis hid) =
    StackSet
      <$> traverseWorkspaces f cur
      <*> traverse (traverseWorkspaces f) vis
      <*> traverse (traverseWorkspaces f) hid

instance Traversable Stack where
  traverse f (Stack u d) =
    Stack <$> traverse f u <*> traverse f d

instance (TraverseWorkspace f) => Foldable (f t l) where
  foldMap fn =
    execWriter
      . traverseWorkspaces_ (\(Workspace _ _ s) -> tell (foldMap (fn . windowInSeat) s))

instance (Functor (f t l), TraverseWorkspace f) => Traversable (f t l) where
  sequenceA =
    traverseWorkspaces $
      \(Workspace t l sf) -> Workspace t l <$> traverse sequenceA sf

class HasFocus f where
  focused :: f a -> Maybe a

data Rectangle = Rectangle Int Int Int Int
  deriving (Read, Show, Eq, Ord)

instance HasFocus (StackSet s sd t l) where
  focused (StackSet c _ _) = focused c

data Screen s sd t l a = Screen
  { screenDetail :: sd,
    screenId :: s,
    workspace :: Workspace t l a
  }
  deriving (Read, Show, Eq, Ord, Functor)

instance HasFocus (Screen s sd t l) where
  focused (Screen _ _ w) = focused w

-- | Defines where a window should appear.
data WindowSeat a = Floating Rectangle a | Tiled a
  deriving (Read, Show, Eq, Ord, Functor, Foldable)

windowInSeat :: WindowSeat a -> a
windowInSeat (Floating _ a) = a
windowInSeat (Tiled a) = a

instance Traversable WindowSeat where
  sequenceA (Floating r fa) = Floating r <$> fa
  sequenceA (Tiled fa) = Tiled <$> fa

instance HasFocus WindowSeat where
  focused (Floating _ a) = Just a
  focused (Tiled a) = Just a

data Workspace t l a = Workspace
  { tag :: t,
    layout :: l,
    stack :: Stack (WindowSeat a)
  }
  deriving (Read, Show, Eq, Ord, Functor)

instance HasFocus (Workspace t l) where
  focused (Workspace _ _ s) = windowInSeat <$> focused s

data Stack a = Stack
  { -- | The elements above the focused one.
    up :: ![a],
    -- | The elements below the focused one including the focused one itself.
    down :: ![a]
  }
  deriving (Read, Show, Eq, Ord, Functor, Foldable)

instance HasFocus Stack where
  focused (Stack _ (a : _)) = Just a
  focused _ = Nothing

-- | Change the tag in a structure.
mapTag :: (TraverseWorkspace f) => (t -> t') -> f t l a -> f t' l a
mapTag fn = mapWorkspaces (\w -> w {tag = fn (tag w)})

-- | Change the layout in a structure.
mapLayout :: (TraverseWorkspace f) => (l -> l') -> f t l a -> f t l' a
mapLayout fn = mapWorkspaces (\w -> w {layout = fn (layout w)})

-- | Return all the tags in a structure.
tags :: (TraverseWorkspace f) => f t l a -> [t]
tags = foldMapWorkspaces ((: []) . tag)

-- | Insert a new window into the StackSet. The optional rectangle indicates if
-- the window should be floating or tiled.
--
-- The window is inserted just above the the currently focused window and is
-- given focus.
insert :: a -> Maybe Rectangle -> StackSet s sd t l a -> StackSet s sd t l a
insert win rect =
  runIdentity
    . onCurrentStack
      ( \(Stack u d) ->
          return $
            (\w -> Stack u (w : d)) $
              maybe (Tiled win) (`Floating` win) rect
      )

-- | Find the tag associated with a window.
findTag :: (TraverseWorkspace f, Eq a) => a -> f t l a -> Maybe t
findTag a =
  getFirst
    . foldMapWorkspaces
      ( \ws ->
          foldMap
            ( \a' ->
                First $ if a' == a then Just (tag ws) else Nothing
            )
            ws
      )

-- | Return true if the window exist in a structure
elem :: (TraverseWorkspace f, Eq a) => a -> f t l a -> Bool
elem a = isJust . findTag a

-- | Convenience function for inserting a window in stack set tiled.
insertTiled :: a -> StackSet s sd t l a -> StackSet s sd t l a
insertTiled win = insert win Nothing

applyStack ::
  (Monad m) =>
  (Stack (WindowSeat a) -> m (Stack (WindowSeat a))) ->
  Workspace t l a ->
  m (Workspace t l a)
applyStack fn (Workspace t l s) = Workspace t l <$> fn s

-- | Apply a function to the currently focused stack.
onCurrentStack ::
  (Monad m) =>
  (Stack (WindowSeat a) -> m (Stack (WindowSeat a))) ->
  StackSet s sd t l a ->
  m (StackSet s sd t l a)
onCurrentStack fn (StackSet cur vis hid) =
  StackSet <$> cur' cur <*> pure vis <*> pure hid
  where
    cur' (Screen s sd ws) = Screen s sd <$> ws' ws
    ws' (Workspace t l s) = Workspace t l <$> fn s