module Rahm.Desktop.StackSet
  ( masterWindow,
    allVisibleWindows,
    differentiateWithFocus,
    concatMapTiledWindows,
    windowsOnWorkspace,
    findWorkspace,
    dbgStackSet,
    ensureWorkspace,
    swapWorkspaces,
    greedyView,
    shiftWin,
    screenRotateBackward,
    screenRotateForward,
    mapWindows,
    swapWindows,
    getLocationWorkspace,
    switchWorkspaces,
    WindowLocation (..),
    windowMemberOfWorkspace,
    findWindow,
    sinkBy,
    modifyWorkspace,
    getFocusedWindow,
    windowTilePosition,
    TilePosition (..),
    module W,
  )
where

import Control.Monad.Writer
import Data.List (find, findIndex, elemIndex)
import Data.List.Safe (head)
import qualified Data.Map as Map
  ( fromList,
    keys,
    lookup,
    mapKeys, keysSet,
  )
import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
import Text.Printf (printf)
import XMonad (Rectangle (..), ScreenDetail (..), WindowSet)
import XMonad.StackSet as W hiding (greedyView, shiftWin, filter)
import qualified XMonad.StackSet (shiftWin)
import Prelude hiding (head)
import qualified Data.Set  as Set

data WindowLocation i l a s sd
  = OnScreen (Screen i l a s sd)
  | OnHiddenWorkspace (Workspace i l a)
  | Floating

data TilePosition i where
  TilePosition :: forall i. i -> Int -> TilePosition i
  deriving (Eq, Show, Ord, Read)

windowTilePosition :: (Eq a, Eq i, Ord a) => a -> StackSet i l a s sd -> Maybe (TilePosition i)
windowTilePosition win ss =
  let ks = Map.keysSet (W.floating ss) in
  case W.findTag win ss of
    Just tag
      | (Just ws) <- findWorkspace tag ss,
        (Just s) <- W.stack ws ->
        TilePosition tag <$> elemIndex win (filter (`Set.notMember`ks) $ W.integrate s)
    _ -> Nothing

getLocationWorkspace :: WindowLocation i l a s sd -> Maybe (Workspace i l a)
getLocationWorkspace (OnScreen (Screen w _ _)) = Just w
getLocationWorkspace (OnHiddenWorkspace w) = Just w
getLocationWorkspace _ = Nothing

allVisibleWindows :: StackSet i l a s sd -> [a]
allVisibleWindows =
  concatMap (W.integrate' . W.stack . W.workspace) <$> W.screens

concatMapTiledWindows :: (Ord a) => (a -> [a]) -> StackSet i l a s sd -> StackSet i l a s sd
concatMapTiledWindows fn (StackSet cur vis hid float) =
  StackSet
    (mapWindowsScreen cur)
    (map mapWindowsScreen vis)
    (map mapWindowsWorkspace hid)
    float
  where
    mapWindowsScreen (Screen work a b) = Screen (mapWindowsWorkspace work) a b
    mapWindowsWorkspace w@(Workspace t l Nothing) = w
    mapWindowsWorkspace (Workspace t l (Just (Stack foc up down))) =
      let up' = concatMap fn up
          down' = concatMap fn down
       in Workspace t l $
            case fn foc of
              [] | (h : t) <- up' -> Just $ Stack h t down
              [] | [] <- up', (h : t) <- down' -> Just $ Stack h [] t
              (h : t) -> Just $ Stack h up (t ++ down)
              _ -> Nothing

mapWindows :: (Ord a, Ord b) => (a -> b) -> StackSet i l a s sd -> StackSet i l b s sd
mapWindows fn (StackSet cur vis hid float) =
  StackSet
    (mapWindowsScreen cur)
    (map mapWindowsScreen vis)
    (map mapWindowsWorkspace hid)
    (Map.mapKeys fn float)
  where
    mapWindowsScreen (Screen work a b) = Screen (mapWindowsWorkspace work) a b
    mapWindowsWorkspace (Workspace t l stack) =
      Workspace t l (fmap (fmap fn) stack)

swapWindows :: (Ord a) => [(a, a)] -> StackSet i l a s d -> StackSet i l a s d
swapWindows toSwap = mapWindows $ \w ->
  fromMaybe w (Map.lookup w toSwapM)
  where
    toSwapM = Map.fromList (toSwap ++ map (\(a, b) -> (b, a)) toSwap)

masterWindow :: StackSet i l a s sd -> Maybe a
masterWindow = head . integrate' . stack . workspace . current

findWorkspace ::
  (Eq i) =>
  i ->
  StackSet i l a s sd ->
  Maybe (Workspace i l a)
findWorkspace wid = find ((== wid) . tag) . workspaces

ensureWorkspace ::
  (Eq i) =>
  i ->
  StackSet i l a s sd ->
  (StackSet i l a s sd, Workspace i l a)
ensureWorkspace t ss =
  case findWorkspace t ss of
    Nothing ->
      let ws = Workspace t (layout . workspace . current $ ss) Nothing
       in (ss {hidden = ws : hidden ss}, ws)
    Just ws -> (ss, ws)

ensureWorkspaces ::
  (Eq i) =>
  [i] ->
  StackSet i l a s sd ->
  (StackSet i l a s sd, [Workspace i l a])
ensureWorkspaces (t : ts) ss =
  let (ss', w) = ensureWorkspace t ss
      (ss'', ws) = ensureWorkspaces ts ss'
   in (ss'', w : ws)
ensureWorkspaces [] ss = (ss, [])

swapWorkspaces ::
  (Eq i) =>
  i ->
  i ->
  StackSet i l a s sd ->
  StackSet i l a s sd
swapWorkspaces tag1 tag2 =
  W.mapWorkspace
    ( \(W.Workspace t a b) ->
        W.Workspace
          ( case (t == tag1, t == tag2) of
              (True, False) -> tag2
              (False, True) -> tag1
              _ -> t
          )
          a
          b
    )

switchWorkspaces ::
  (Eq i) =>
  i ->
  i ->
  StackSet i l a s sd ->
  StackSet i l a s sd
switchWorkspaces t1 t2 (ensureWorkspaces [t1, t2] -> (ss, [w1, w2])) =
  W.mapWorkspace
    ( \case
        (Workspace t _ _) | t == t1 -> w2
        (Workspace t _ _) | t == t2 -> w1
        w -> w
    )
    ss

dbgStackSet :: WindowSet -> String
dbgStackSet ws@(W.StackSet cur vis hidden _) = execWriter $ do
  tell "* " >> logScreen cur >> tell "\n"
  mapM_ (\s -> tell "  " >> logScreen s >> tell "\n") vis

  mapM_ logWorkspace (W.workspaces ws)
  where
    logWorkspace (Workspace tag _ st) = do
      tell $ printf "WS %s\n" tag
      forM_ st $ \(Stack foc up down) -> do
        mapM_ (tell . printf "    %d\n") up
        tell $ printf "  * %d\n" foc
        mapM_ (tell . printf "    %d\n") down

    logScreen (Screen ws sid (SD (Rectangle _ _ w h))) = do
      tell (printf "id=%s (%sx%s) - [%s]" (show sid) (show w) (show h) (W.tag ws))

greedyView :: (Eq i) => i -> StackSet i l a s sd -> StackSet i l a s sd
greedyView wid ss = switchWorkspaces (tag . workspace . current $ ss) wid ss

shiftWin :: (Ord a, Eq s, Eq i) => i -> a -> StackSet i l a s sd -> StackSet i l a s sd
shiftWin wid a = XMonad.StackSet.shiftWin wid a . fst . ensureWorkspace wid

windowsOnWorkspace :: (Eq i) => i -> StackSet i l a s sd -> [a]
windowsOnWorkspace i ss = fromMaybe [] $ do
  ws <- find ((== i) . W.tag) (W.workspaces ss)
  s <- W.stack ws
  return (W.integrate s)

screenRotateBackward :: W.StackSet i l a sid sd -> W.StackSet i l a sid sd
screenRotateBackward (W.StackSet current visible others floating) = do
  let screens = current : visible
      workspaces = tail $ cycle $ map W.workspace screens
      (current' : visible') = zipWith (\s w -> s {workspace = w}) screens workspaces
   in W.StackSet current' visible' others floating

screenRotateForward :: W.StackSet i l a sid sd -> W.StackSet i l a sid sd
screenRotateForward (W.StackSet current visible others floating) = do
  let screens = current : visible
      workspaces = rcycle $ map W.workspace screens
      (current' : visible') = zipWith (\s w -> s {workspace = w}) screens workspaces
   in W.StackSet current' visible' others floating
  where
    rcycle l = last l : l

{- Finds a Window and returns the screen its on and the workspace its on.
 - Returns nothing if the window doesn't exist.
 -
 - If the window is not a screen Just (Nothing, workspace) is returned.
 - If the window is a floating window Just (Nothing, Nothing) is returned. -}
findWindow ::
  (Eq a) => StackSet i l a s sd -> a -> Maybe (WindowLocation i l a s sd)
findWindow (StackSet cur vis hid float) win =
  listToMaybe . catMaybes $
    map findWindowScreen (cur : vis)
      ++ map findWindowWorkspace hid
      ++ [findWindowFloat]
  where
    findWindowScreen s@(Screen ws _ _) =
      if windowMemberOfWorkspace ws win
        then Just (OnScreen s)
        else Nothing

    findWindowWorkspace w =
      if windowMemberOfWorkspace w win
        then Just (OnHiddenWorkspace w)
        else Nothing

    findWindowFloat =
      if win `elem` Map.keys float
        then Just Floating
        else Nothing

windowMemberOfWorkspace :: (Eq a) => Workspace i l a -> a -> Bool
windowMemberOfWorkspace (Workspace _ _ s) w = w `elem` integrate' s

modifyWorkspace :: (Eq i) => i -> (Workspace i l a -> Workspace i l a) -> StackSet i l a s sd -> StackSet i l a s sd
modifyWorkspace tag fn =
  mapWorkspace
    ( \ws ->
        if W.tag ws == tag
          then fn ws
          else ws
    )

differentiateWithFocus :: (Eq a) => a -> [a] -> Maybe (Stack a)
differentiateWithFocus _ [] = Nothing
differentiateWithFocus thing lst =
  case break (==thing) lst of
    (up, foc:down) -> Just $ Stack foc (reverse up) down
    _ -> differentiate lst

getFocusedWindow :: StackSet i l a s sd -> Maybe a
getFocusedWindow (StackSet cur _ _ _) = W.focus <$> (W.stack . W.workspace) cur

sinkBy :: (Eq a, Eq i, Ord a) => a -> a -> StackSet i l a s sd -> StackSet i l a s sd
sinkBy win toSinkBy ss =
  case (findTag win ss, findTag toSinkBy ss) of
    (Nothing, _) -> ss
    (Just w1, Just w2)
      | w1 == w2 ->
        modifyWorkspace
          w1
          ( \(W.Workspace t l s) ->
              W.Workspace t l (Just $ insertBy win toSinkBy s)
          )
          $ W.delete win ss
    _ -> W.sink win ss
  where
    insertBy win to Nothing = W.Stack win [] []
    insertBy win to (Just (W.Stack foc down up)) =
      case () of
        ()
          | to `elem` down ->
            W.Stack
              foc
              (concatMap (\e -> if e == to then [e, win] else [e]) down)
              up
        ()
          | to `elem` up ->
            W.Stack
              foc
              down
              (concatMap (\e -> if e == to then [e, win] else [e]) up)
        () -> W.Stack win (foc : down) up