-- Import

import Data.Monoid
import Data.Char (toLower)
import Control.Monad (forM_, liftM2, when)
import Data.Maybe (isJust)
import XMonad
import XMonad.Actions.OnScreen (greedyViewOnScreen)
import XMonad.Actions.PhysicalScreens
import XMonad.Hooks.EwmhDesktops
import XMonad.Hooks.ManageDocks
import XMonad.Hooks.ManageHelpers
import XMonad.Hooks.Rescreen
import XMonad.Hooks.StatusBar
import XMonad.Hooks.StatusBar.PP
import XMonad.Layout.Grid
import XMonad.Layout.Magnifier
import XMonad.Layout.NoBorders
import XMonad.Layout.PerWorkspace
import XMonad.Layout.Renamed
import XMonad.Layout.Spacing
import XMonad.Layout.Spiral
import XMonad.Layout.Tabbed
import XMonad.Layout.ThreeColumns
import qualified XMonad.StackSet as W
import XMonad.Util.ClickableWorkspaces
import XMonad.Util.EZConfig
import qualified XMonad.Util.Hacks as Hacks
import XMonad.Util.Loggers
import XMonad.Util.NamedActions
import XMonad.Util.NamedScratchpad
import XMonad.Util.Run (runProcessWithInput, spawnPipe)
import XMonad.Util.SpawnOnce
import System.IO (hClose, hPutStr)

-- Main
main :: IO ()
main =
  xmonad
    . ewmhFullscreen
    . setEwmhSwitchDesktopHook safeView
    . ewmh
    . docks
    . rescreenHook def {afterRescreenHook = restartStatusBars, rescreenWorkspacesHook = Last (Just safeRescreen)}
    . addRandrChangeHook (spawn myScreenHotplug)
    $ myConfig

-- Variables
myModMask :: KeyMask
myModMask = mod4Mask

myTerminal :: String
myTerminal = "kitty"

myBrowser :: String
myBrowser = "zen-browser"

myNormBorderColor :: String
myNormBorderColor = "#1f2335"

myFocusBorderColor :: String
myFocusBorderColor = "#41a6b5"

myKeyboardLayoutIT :: String
myKeyboardLayoutIT = " -layout it"

myKeyboardLayoutUS :: String
myKeyboardLayoutUS = " -layout us"

myKeyboardLayoutSK :: String
myKeyboardLayoutSK = " -layout sk -variant qwerty"

myScreenshotDir :: String
myScreenshotDir = "/home/fuxino/Pictures/Screenshots/"

myScreenshotName :: String
myScreenshotName = "Screenshot-%Y-%m-%d-%H%M%S.png"

-- Sets up the outputs, the wallpaper and trayer; shared with ~/.local/bin/movie
myScreenLayout :: String
myScreenLayout = "/home/fuxino/.local/bin/screen-layout"

-- Safe to re-run (e.g. spawnOnce after this string changes and xmonad restarts); keeps the last wallpapers across logins
myScreenSetup :: String
myScreenSetup =
  "sleep 1; " ++ myScreenLayout ++ " --keep-wallpaper; pgrep -x picom > /dev/null || picom -b -c"

-- Run when a monitor is (dis)connected; restarts trayer so it moves to the right screen
myScreenHotplug :: String
myScreenHotplug = myScreenLayout ++ " --keep-wallpaper"

myScratchpads :: [NamedScratchpad]
myScratchpads =
  [ NS "Riseup-vpn" "riseup-vpn" (className =? "RiseupVPN") (customFloating $ W.RationalRect 0.15 0.4 0.15 0.4),
    NS "Terminal" spawnTerminal (className =? "scratch-terminal") (customFloating $ W.RationalRect 0.2 0.2 0.6 0.6),
    NS "Music" spawnMusic (className =? "scratch-music") (customFloating $ W.RationalRect 0.2 0.2 0.6 0.6),
    NS "Mail" spawnMail (className =? "scratch-mail") (customFloating $ W.RationalRect 0.2 0.2 0.6 0.6),
    NS "Cal" spawnCal (className =? "scratch-cal") (customFloating $ W.RationalRect 0.2 0.2 0.6 0.6),
    NS "Btop" spawnBtop (className =? "scratch-btop") (customFloating $ W.RationalRect 0.2 0.2 0.6 0.6)
  ]
  where
    spawnTerminal = myTerminal ++ " --class scratch-terminal -T kitty-float"
    spawnMusic = myTerminal ++ " --class scratch-music -T ncmpcpp ncmpcpp"
    spawnMail = myTerminal ++ " --class scratch-mail -T mutt bash -c 'pushd ~/Downloads > /dev/null; mutt; popd > /dev/null'"
    spawnCal = myTerminal ++ " --class scratch-cal -T calcurse calcurse"
    spawnBtop = myTerminal ++ " --class scratch-btop -T btop btop"

-- Config
myConfig =
  addDescrKeys' ((myModMask, xK_F1), showKeys) myKeys $ def
    { modMask = myModMask,
      terminal = myTerminal,
      workspaces = myWorkspaces,
      focusedBorderColor = myFocusBorderColor,
      normalBorderColor = myNormBorderColor,
      layoutHook = myLayoutHook,
      startupHook = myStartupHook,
      manageHook = myManageHook,
      handleEventHook = myHandleEventHook,
      logHook = myLogHook
    }

-- Keys
-- xmonad's described defaults minus the ones overridden here, so M-<F1> lists each key once
myKeys :: XConfig Layout -> [((KeyMask, KeySym), NamedAction)]
myKeys c = dropEmptyGroups (filter ((`notElem` overridden) . fst) (defaultKeysDescr c)) ++ custom
  where
    custom = myCustomKeys c
    -- M-r/M-S-r: there's no third screen
    overridden = (myModMask, xK_r) : (myModMask .|. shiftMask, xK_r) : filter (/= (0, 0)) (map fst custom)
    -- (0, 0) is a subtitle or separator; drop it when nothing follows before the next one
    dropEmptyGroups (((0, 0), _) : rest@(((0, 0), _) : _)) = dropEmptyGroups rest
    dropEmptyGroups [((0, 0), _)] = []
    dropEmptyGroups (x : rest) = x : dropEmptyGroups rest
    dropEmptyGroups [] = []

myCustomKeys :: XConfig Layout -> [((KeyMask, KeySym), NamedAction)]
myCustomKeys c =
  concatMap
    (\(t, ks) -> subtitle t : mkNamedKeymap c ks)
    [ ( "launchers",
        [ ("M-d", addName "Run a program (dmenu)" $ spawn "dmenu_run"),
          ("M-p", addName "Passwords (passmenu)" $ spawn "passmenu -i"),
          ("<XF86HomePage>", addName "Browser" $ spawn myBrowser),
          ("<XF86Calculator>", addName "Calculator" $ spawn "qalculate-gtk")
        ]
      ),
      ( "scratchpads",
        [ ("M-S-s", addName "Terminal" $ namedScratchpadAction myScratchpads "Terminal"),
          ("M-S-a", addName "Calendar (calcurse)" $ namedScratchpadAction myScratchpads "Cal"),
          ("M-o", addName "btop" $ namedScratchpadAction myScratchpads "Btop"),
          ("M-S-v", addName "RiseupVPN" $ namedScratchpadAction myScratchpads "Riseup-vpn"),
          ("<XF86Mail>", addName "Mail (mutt)" $ namedScratchpadAction myScratchpads "Mail"),
          ("<XF86Tools>", addName "Music (ncmpcpp)" $ namedScratchpadAction myScratchpads "Music")
        ]
      ),
      ( "session and layout",
        [ ("M-S-l", addName "Lock the screen" $ spawn "loginctl lock-session"),
          ("M-S-n", addName "Show older notification" $ spawn "dunstctl history-pop"),
          ("M-C-n", addName "Close all notifications" $ spawn "dunstctl close-all"),
          ("M-b", addName "Toggle the status bar gap" $ sendMessage ToggleStruts),
          ("M-f", addName "Tabbed layout" $ sendMessage $ JumpToLayout "Tabbed"),
          ("M-i", addName "Keyboard: Italian" $ spawn $ "setxkbmap" ++ myKeyboardLayoutIT),
          ("M-u", addName "Keyboard: US" $ spawn $ "setxkbmap" ++ myKeyboardLayoutUS),
          ("M-s", addName "Keyboard: Slovak" $ spawn $ "setxkbmap" ++ myKeyboardLayoutSK),
          ("M-x", addName "Connect earbuds" $ spawn "bluetoothctl connect F4:9D:8A:B9:8C:EE"),
          ("M-S-x", addName "Disconnect earbuds" $ spawn "bluetoothctl disconnect F4:9D:8A:B9:8C:EE")
        ]
      ),
      ( "screenshots",
        [ ("<Print>", addName "Whole screen" $ unGrab *> spawn ("scrot " ++ myScreenshotDir ++ myScreenshotName)),
          ("S-<Print>", addName "Select an area" $ unGrab *> spawn ("scrot -s " ++ myScreenshotDir ++ myScreenshotName)),
          ("M-S-<Print>", addName "Focused window" $ unGrab *> spawn ("scrot -u " ++ myScreenshotDir ++ myScreenshotName))
        ]
      ),
      ( "media keys",
        [ ("<XF86AudioMute>", addName "Mute" $ spawn "pactl set-sink-mute $(pactl get-default-sink) toggle"),
          ("<XF86AudioLowerVolume>", addName "Volume down" $ spawn "pactl set-sink-volume $(pactl get-default-sink) -10%"),
          ("<XF86AudioRaiseVolume>", addName "Volume up" $ spawn "wpctl set-volume -l 1.0 @DEFAULT_AUDIO_SINK@ 10%+"),
          ("<XF86AudioPlay>", addName "Play/pause" $ spawn "playerctl play-pause"),
          ("<XF86AudioPause>", addName "Pause" $ spawn "playerctl pause"),
          ("<XF86AudioStop>", addName "Stop" $ spawn "playerctl stop"),
          ("<XF86AudioNext>", addName "Next track" $ spawn "playerctl next"),
          ("<XF86AudioPrev>", addName "Previous track" $ spawn "playerctl previous")
        ]
      ),
      ( "workspaces and screens",
        [ (m ++ show n, addName (d ++ show n ++ " (" ++ name ++ ")") $ windows $ f ws)
          | (n, ws, name) <- zip3 [1 :: Int ..] myWorkspaces ["terminal", "web", "games", "chat", "media", "docs", "VM", "free", "free"],
            (m, d, f) <- [("M-", "Workspace ", safeView), ("M-S-", "Send window to workspace ", W.shift)]
        ]
          ++ [ (m ++ k, addName (d ++ side) $ f def sc)
               | (k, sc, side) <- [("w", 0, "left screen"), ("e", 1, "right screen")],
                 (m, d, f) <- [("M-", "Focus ", viewScreen), ("M-S-", "Send window to ", sendToScreen)]
             ]
      )
    ]

-- Shows the key list in dmenu (type to filter; picking a line does nothing).
-- spawnPipe rather than runProcessWithInput, which would block xmonad while dmenu is open.
showKeys :: [((KeyMask, KeySym), NamedAction)] -> NamedAction
showKeys ks = addName "Show keybindings" $ do
  h <- spawnPipe "dmenu -i -l 30 -p Keys"
  io $ hPutStr h (unlines $ map tidy $ filter (not . null) $ lines $ unlines $ showKm ks) >> hClose h
  where
    -- showKm writes the mod key as M4- (it doesn't know modMask); re-pad after shortening it
    tidy l@('>' : _) = l
    tidy l = let (k, d) = break (== ' ') l in take 24 (fixMod k ++ repeat ' ') ++ dropWhile (== ' ') d
    fixMod ('M' : '4' : '-' : r) = "M-" ++ fixMod r
    fixMod (x : r) = x : fixMod r
    fixMod [] = []

-- greedyView that keeps wsGames on screen 0 (HDMI): Wine games lose the mouse when moved between screens
safeView :: WorkspaceId -> WindowSet -> WindowSet
safeView i ws
  | i == wsGames = greedyViewOnScreen 0 i ws
  | W.currentTag ws == wsGames && i `elem` map (W.tag . W.workspace) (W.visible ws) = W.view i ws
  | otherwise = W.greedyView i ws

-- xmonad's rescreen, except that a hidden wsGames is never put on a secondary screen
safeRescreen :: X ()
safeRescreen = do
  screens <- withDisplay getCleanedScreenInfo
  case screens of
    [] -> return ()
    xinesc : xinescs ->
      windows $ \ws ->
        let notGames = (/= wsGames) . W.tag
            pool = map W.workspace (W.visible ws) ++ filter notGames (W.hidden ws)
            (xs, _) = splitAt (length xinescs) pool
            shown = map W.tag (W.workspace (W.current ws) : xs)
         in ws
              { W.current = W.Screen (W.workspace (W.current ws)) 0 (SD xinesc),
                W.visible = zipWith3 W.Screen xs [1 ..] (map SD xinescs),
                W.hidden = filter ((`notElem` shown) . W.tag) (map W.workspace (W.visible ws) ++ W.hidden ws)
              }

-- Workspaces
myWorkspaces :: [String]
myWorkspaces = [wsTerm, wsWeb, wsGames, wsChat, wsMedia, wsDocs, wsVM] ++ map show [8, 9]

-- Icons spill past their advance in the non-Mono Nerd Font: keep them ~1000-1160 units wide or gaps look uneven
wsTerm, wsWeb, wsGames, wsChat, wsMedia, wsDocs, wsVM :: String
wsTerm = "1:\xf489 "
wsWeb = "2:\xf0ac "
wsGames = "3:\xf11b "
wsChat = "4:\xf086 "
wsMedia = "5:\xefab "
wsDocs = "6:\xede2 "
wsVM = "7:\xf108 "

-- Layout
myLayoutHook = onWorkspace wsWeb myWebLayout $ onWorkspace wsGames myGamesLayout $ onWorkspace wsMedia myVideoLayout myDefaultLayout
  where
    myWebLayout = avoidStruts $ renamed [CutWordsLeft 1] $ spacingWithEdge 3 $ myTabbed ||| tiled ||| Mirror tiled ||| threeCol ||| Grid ||| spiral 0.856
    myGamesLayout = avoidStruts $ smartBorders Full
    myVideoLayout = avoidStruts $ smartBorders Full
    myDefaultLayout = avoidStruts $ renamed [CutWordsLeft 1] $ spacingWithEdge 3 $ tiled ||| Mirror tiled ||| myTabbed ||| threeCol ||| Grid ||| spiral 0.856
    threeCol =
      renamed [Replace "ThreeCol"] $
        magnifiercz' 1.3 $
          ThreeColMid nmaster delta ratio
    myTabbed =
      renamed [Replace "Tabbed"] $
        tabbed shrinkText myTabConfig
    myTabConfig =
      def
        { activeColor = "#24283b",
          inactiveColor = "#737aa2"
        }
    tiled = Tall nmaster delta ratio
    nmaster = 1
    ratio = 1 / 2
    delta = 3 / 100

-- Xmobar
-- The laptop config leaves room for trayer, so it goes on the laptop screen,
-- which is screen 1 with HDMI and screen 0 without.
myStatusBar :: Int -> Int -> String
myStatusBar nScreens i = "xmobar -x " ++ show i ++ " ~/.config/xmobar/" ++ rc
  where
    rc = if nScreens > 1 && i /= 1 then "xmobarrc_hdmi" else "xmobarrc_laptop"

-- One bar per screen, all restarted on every screen change: the right config for a
-- screen id depends on the screen count, and dynamicSBs would keep the old bar.
-- Kills by name (pkill waits) because killAllStatusBars' PID list can lose bars.
restartStatusBars :: X ()
restartStatusBars = do
  _ <- runProcessWithInput "pkill" ["-x", "xmobar"] ""
  screens <- gets (W.screens . windowset)
  forM_ screens $ \s -> let S i = W.screen s in spawn (myStatusBar (length screens) i)

myXmobarPP :: PP
myXmobarPP =
  def
    { ppSep = pink " • ",
      ppTitleSanitize = xmobarStrip,
      ppCurrent = wrap " " "" . xmobarBorder "Top" "#ff79c6" 2,
      ppHidden = cultured . wrap " " "",
      ppHiddenNoWindows = outerspace . wrap " " "",
      ppUrgent = sunsetorange . wrap (keylime "!") (keylime "!"),
      -- Show "1:<icon>" as just the icon; tags stay as they are (clickablePP uses the index)
      ppRename = \tag _ -> case break (== ':') tag of
        (_, ':' : icon) -> icon
        _ -> tag,
      -- NamedScratchpad never removes NSP once created; hide it while it's empty
      ppSort = (. filter (\w -> W.tag w /= scratchpadWorkspaceTag || isJust (W.stack w))) <$> ppSort def,
      ppOrder = \[ws, l, _, wins] -> [ws, l, wins],
      ppLayout = layoutIcon,
      ppExtras = [budgetTitles]
    }
  where
    layoutIcon :: String -> String
    layoutIcon l = case l of
      "Tall" -> "\xf0bcc"
      "Mirror Tall" -> "\xf0bcb"
      "Tabbed" -> "\xf04e9"
      "ThreeCol" -> "\xf056d"
      "Grid" -> "\xf0570"
      "Spiral" -> "\xf056e"
      "Full" -> "\xf0293"
      _ -> l

    -- Titles share one width budget, so more windows get shorter titles instead of overflowing the bar
    budgetTitles :: Logger
    budgetTitles = do
      n <- gets $ length . W.integrate' . W.stack . W.workspace . W.current . windowset
      let width = max 8 . min 30 $ 60 `div` max 1 n
      logTitles (formatFocused width) (formatUnfocused width)

    formatFocused, formatUnfocused :: Int -> String -> String
    formatFocused w = wrap (cultured "[") (cultured "]") . pink . ppWindow w
    formatUnfocused w = wrap (outerspace "[") (outerspace "]") . cultured . ppWindow w

    ppWindow :: Int -> String -> String
    ppWindow w = xmobarRaw . (\t -> if null t then "untitled" else t) . shorten' "…" w

    outerspace, pink, sunsetorange, cultured, keylime :: String -> String
    pink = xmobarColor "#ff79c6" ""
    cultured = xmobarColor "#f8f8f2" ""
    keylime = xmobarColor "#f1fa8c" ""
    sunsetorange = xmobarColor "#ff5555" ""
    outerspace = xmobarColor "#44475a" ""

-- Autostart
myStartupHook :: X ()
myStartupHook = do
  restartStatusBars
  spawnOnce myScreenSetup
  spawnOnce "systemctl --user start fehbg.timer"
  spawnOnce "xsetroot -cursor_name left_ptr"
  spawnOnce "xset s 590 10"
  spawnOnce "xset dpms 0 0 900"
  spawnOnce "xss-lock -n xidle-dim -l -- screen-lock"
  spawnOnce "systemctl --user start xidle-inhibit"
  spawnOnce "redshift-gtk"
  spawnOnce "systemctl --user start udiskie"
  spawnOnce "nm-applet"
  spawnOnce "arch-audit-gtk"
  spawnOnce "calcurse --daemon"
  spawnOnce "systemctl --user start dunst"
  spawnOnce "lxpolkit"
  spawnOnce "numlockx"
  spawnOnce "blueman-applet"
  spawnOnce "adb start-server"
  spawnOnce "systemctl --user start discover-overlay-wait"

-- Manage hook
myManageHook :: ManageHook
myManageHook =
  composeAll
    [ -- Proton games are steam_app_<id>; Wine windows are named after their .exe
      className ^? "steam_app_" --> doShift wsGames,
      className $? ".exe" --> doShift wsGames,
      className =? "Gpodder" --> doShift wsMedia,
      className =? "An Anime Game Launcher" --> doShift wsGames,
      className =? "Qalculate-gtk" --> doFloat,
      title =? "Genshin Impact" --> doShift wsGames,
      classIn ["gimp", "audacity"] --> doShift wsMedia,
      classIn ["pavucontrol", "org.pulseaudio.pavucontrol", "zenity", "org.gnome.zenity", "lxappearance", "qt5ct", "qt6ct", "gsimplecal"] --> doFloat,
      className =? "Signal" --> doShift wsChat,
      className =? "Spotify" --> doShift wsMedia,
      className =? "Transmission-gtk" --> doShift wsWeb,
      className =? "Viewnior" --> doFloat,
      className =? "Vivaldi-stable" --> doShift wsWeb,
      className =? "Xmessage" --> doFloat,
      className =? "Xreader" --> doShift wsDocs,
      className =? "calibre" --> doShift wsDocs,
      className =? "discord" --> doShift wsChat,
      className =? "feh" --> doFloat,
      title =? "glxgears" --> doFloat,
      className =? "library_manager" --> doShift wsDocs,
      className =? "libreoffice-calc" --> doShift wsDocs,
      className =? "libreoffice-writer" --> doShift wsDocs,
      className =? "mpv" --> doFullFloat,
      className =? "mpv" --> viewShift wsMedia,
      className =? "freetube" --> doShift wsMedia,
      className =? myTerminal --> doShift wsTerm,
      className =? "scrcpy" --> doShift wsMedia,
      className =? "steam" --> doShift wsGames,
      className =? "steam_app_109600" --> doFloat,
      className =? "steam_app_289070" --> doFloat,
      className =? "VirtualBox Machine" --> doShift wsVM,
      className =? "VirtualBox Machine" --> doFullFloat,
      className =? "VirtualBox Manager" --> doShift wsVM,
      className =? "zen" --> doShift wsWeb,
      isDialog --> doFloat,
      -- Windows that open already fullscreen; ewmhFullscreen only handles later requests
      isFullscreen --> doFullFloat
    ]
    <+> namedScratchpadManageHook myScratchpads
  where
    viewShift = doF . liftM2 (.) safeView W.shift
    -- Case-insensitive, as GTK4 and firejailed apps capitalise their class differently
    classIn names = (`elem` names) . map toLower <$> className

-- Handle event hook
myHandleEventHook :: Event -> X All
myHandleEventHook =
  handleEventHook def
    <> Hacks.trayerAboveXmobarEventHook
    <> Hacks.trayerPaddingXmobarEventHook
    <> Hacks.fixSteamFlicker

-- Log hook
myLogHook :: X ()
myLogHook = do
  raiseDiscoverOverlay
  xmonadPropLog =<< dynamicLogString =<< clickablePP myXmobarPP

-- discover-overlay's windows are override-redirect, so xmonad doesn't manage them and
-- raising a floating window (e.g. a game) covers them; put them back on top after every change
raiseDiscoverOverlay :: X ()
raiseDiscoverOverlay = withDisplay $ \d -> do
  root <- asks theRoot
  (_, _, ws) <- io $ queryTree d root
  forM_ ws $ \w -> do
    cls <- io $ getClassHint d w
    when (resClass cls == "Discover-overlay") $ io $ raiseWindow d w