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

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

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

shiftWinNoFocus :: (Ord a, Eq s, Eq i) => i -> a -> StackSet i l a s sd -> StackSet i l a s sd
shiftWinNoFocus n w s = case findTag w s of
  Just from | n `tagMember` s && n /= from -> go from s
  _ -> s
  where
    go from = onWorkspace n (focusDown . insertUp w) . onWorkspace from (delete' w)
    onWorkspace n f s = view (currentTag s) . f . view n $ s

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

swapUpOrMirror :: (Eq a, Eq i, Ord a) => StackSet i l a s sd -> StackSet i l a s sd
swapUpOrMirror
  ss@StackSet
    { floating = flt
    } =
    let mwin = W.peek ss
     in case mwin of
          Nothing -> ss
          Just win ->
            if win `Map.member` flt
              then ss {floating = Map.update (Just . flipVert) win flt}
              else W.swapUp ss
    where
      flipVert (RationalRect x y w h) =
        RationalRect x (abs (1 - y) - h) w h

swapDownOrMirror :: (Eq a, Eq i, Ord a) => StackSet i l a s sd -> StackSet i l a s sd
swapDownOrMirror
  ss@StackSet
    { floating = flt
    } =
    let mwin = W.peek ss
     in case mwin of
          Nothing -> ss
          Just win ->
            if win `Map.member` flt
              then ss {floating = Map.update (Just . flipHoriz) win flt}
              else W.swapDown ss
    where
      flipHoriz (RationalRect x y w h) =
        RationalRect (abs (1 - x) - w) y w h