-- Wml: Window Management Language.
--
-- Parser for WML objects
--
-- Some examples of WML objects are:
--
-- a          // The workspace or window (context dependent) tagged 'a'
-- @a         // All windows on workspace 'a' or the workspace with window 'a'
-- ,.         // The workspace to to the right of the current one.
-- @,.        // All windows on the workspace to the right of the current one.
-- @,^        // All the windows on the screen second from the left
-- &z!~@,,^   // The window tagged with z and The last window on the screen third from the left
-- @@s        // All the windows that share a workspace with the window tagged s
-- \%@s       // All windows except those on workspace 's'
module Rahm.Desktop.Keys.Wml
  ( readWorkspaceMacro,
    readWindowsetMacro,
    readNextWorkspace,
    readNextLocationSet,
    readNextLocationSet',
    moveLocationToWorkspace,
    moveWindowToWorkspaceFn,
    getAndResetWindowSelection,
    gotoWorkspaceFn,
    toggleWindowInSelection,
    addWindowToSelection,
    clearWindowSelection,
    removeWindowFromSelection,
    readMacroString,
    justWorkspace,
    justWorkspaceWithPreferredWindow,
    blackHoleWorkspace,
    alternateWorkspace,
    floatWorkspace,
    joinMaybe,
    workspaceForKeys,
    workspaceForStringT,
    workspaceForString,
    locationSetForKeysT,
    locationSetForKeys,
    readNextWorkspaceName,
    workspaceName,
  )
where

import Control.Monad (forM_, join, void, when)
-- getMostRecentLocationInHistory,

-- pastHistory,

import Control.Monad.Reader (ReaderT (runReaderT))
import Control.Monad.Trans.Class (MonadTrans (lift))
import Control.Monad.Trans.Maybe (MaybeT (..), mapMaybeT)
import Control.Monad.Trans.State as S
  ( StateT (StateT, runStateT),
    evalState,
    evalStateT,
    get,
    gets,
    modify',
    put,
  )
import Data.Char (isAlpha, isAlphaNum, isDigit, ord)
import Data.List (find, intercalate, sortOn)
import Data.List.Safe (head, last)
import Data.Map (Map)
import qualified Data.Map as Map
  ( delete,
    elems,
    empty,
    insert,
    keys,
    lookup,
    member,
  )
import Data.Maybe (catMaybes, fromJust, fromMaybe)
import Data.Ord (Down (..))
import Data.Typeable (cast)
import Data.Void (Void, absurd)
import Rahm.Desktop.BorderColors (BorderColor (BorderColor), setBorderColor)
import Rahm.Desktop.Common
  ( Location (..),
    askWindowId,
    getCurrentLocation,
    getCurrentWorkspace,
    gotoWorkspace,
    moveLocationToWorkspace,
    windowsInWorkspace,
  )
import Rahm.Desktop.History
  ( lastLocation,
    nextLocation,
  )
import Rahm.Desktop.Keys.KeyFeed
import Rahm.Desktop.Layout.PinWindow (pinnedWindows)
import Rahm.Desktop.Logger (LogLevel (Info, Trace), logs)
import Rahm.Desktop.Marking
  ( farLeftWindow,
    farRightWindow,
    getAlternateWindows,
    getAlternateWorkspace,
    getMarkedLocations,
    windowLocation,
  )
import qualified Rahm.Desktop.StackSet as W
import Rahm.Desktop.Submap (mapNextStringWithKeysym)
import Rahm.Desktop.Workspaces
  ( accompaningWorkspace,
    adjacentWorkspace,
    adjacentWorkspaceNotVisible,
    getHorizontallyOrderedScreens,
    next,
    prev,
    workspaceWithWindow,
  )
import Rahm.Desktop.XMobarLog.PendingBuffer
  ( addStringToPendingBuffer,
    setPendingBuffer,
  )
import System.Directory (doesFileExist)
import System.Exit (ExitCode (..), exitSuccess, exitWith)
import System.FilePath ((</>))
import System.IO (readFile)
import Text.Printf (printf)
import Text.Read (readMaybe)
import XMonad
  ( Default (def),
    Directories' (dataDir),
    ExtensionClass (..),
    KeyMask,
    KeySym,
    MonadReader (ask),
    StateExtension (PersistentExtension),
    Typeable,
    Window,
    WindowSet,
    WorkspaceId,
    X,
    asks,
    clearArea,
    directories,
    io,
    killWindow,
    windows,
    withWindowSet,
    xK_Escape,
    xK_Return,
  )
import XMonad.Prompt.ConfirmPrompt (confirmPrompt)
import qualified XMonad.Util.ExtensibleState as XS (get, modify, put)
import Prelude hiding (head, last)

data MaybeMacros = NoMacros | YesMacros Macros
  deriving (Read, Show)

data Macros = Macros
  { workspaceMacros :: Map (KeyMask, KeySym) KeyString,
    windowsetMacros :: Map (KeyMask, KeySym) KeyString
  }
  deriving (Read, Show)

getMacroDataFile :: X FilePath
getMacroDataFile = asks ((</> "rde-recorded-macros") . dataDir . directories)

getMacros :: X Macros
getMacros = do
  XS.get
    >>= ( \case
            NoMacros -> do
              file <- getMacroDataFile
              exists <- io $ doesFileExist file
              if exists
                then
                  io (readFile file)
                    >>= ( \s ->
                            let macros = fromMaybe (Macros mempty mempty) (readMaybe s :: Maybe Macros)
                             in XS.put (YesMacros macros) >> return macros
                        )
                else return $ Macros mempty mempty
            YesMacros macros -> return macros
        )

saveMacros :: X ()
saveMacros = do
  dataFile <- getMacroDataFile
  macros <- getMacros
  io $ writeFile dataFile $ show macros

selColor = BorderColor "#ffff00" "#b8b880"

insertWorkspaceMacroString :: (KeyMask, KeySym) -> KeyString -> X ()
insertWorkspaceMacroString k ks = do
  macros <- getMacros
  XS.put $
    YesMacros $
      macros
        { workspaceMacros = Map.insert k ks (workspaceMacros macros)
        }
  saveMacros

insertWindowSetMacroString :: (KeyMask, KeySym) -> KeyString -> X ()
insertWindowSetMacroString k ks = do
  macros <- getMacros
  XS.put $
    YesMacros $
      macros
        { windowsetMacros = Map.insert k ks (windowsetMacros macros)
        }
  saveMacros

instance ExtensionClass MaybeMacros where
  initialValue = NoMacros

newtype WindowSelect = WindowSelect (Map Window (X ()))

instance ExtensionClass WindowSelect where
  initialValue = WindowSelect mempty

toggleWindowInSelection :: Window -> X ()
toggleWindowInSelection win = do
  (WindowSelect sel) <- XS.get
  case Map.lookup win sel of
    Nothing -> do
      cleanup <- setBorderColor selColor [win]

      XS.put $ WindowSelect $ Map.insert win cleanup sel
    (Just cleanup) -> do
      XS.put $ WindowSelect $ Map.delete win sel
      cleanup

addWindowToSelection :: Window -> X ()
addWindowToSelection win = do
  (WindowSelect sel) <- XS.get
  case Map.lookup win sel of
    Nothing -> do
      foc <- withWindowSet (return . W.peek)

      cleanup <- setBorderColor selColor [win]

      XS.put $ WindowSelect $ Map.insert win cleanup sel
    _ -> return ()

removeWindowFromSelection :: Window -> X ()
removeWindowFromSelection win = do
  (WindowSelect sel) <- XS.get
  case Map.lookup win sel of
    (Just cleanup) -> do
      XS.put $ WindowSelect $ Map.delete win sel
      cleanup
    _ -> return ()

clearWindowSelection :: X ()
clearWindowSelection = void getAndResetWindowSelection

getAndResetWindowSelection :: X [Window]
getAndResetWindowSelection = do
  (WindowSelect mp) <- XS.get
  sequence_ (Map.elems mp)
  XS.put (initialValue :: WindowSelect)
  return (Map.keys mp)

data Workspace = forall a.
  (Typeable a) =>
  Workspace
  { moveWindowToWorkspaceFn :: Window -> X (WindowSet -> WindowSet),
    gotoWorkspaceFn :: X (),
    workspaceName :: Maybe String,
    extraWorkspaceData :: a
  }

readWorkspaceMacro :: MaybeT X ()
readWorkspaceMacro =
  mapNextStringWithKeysym $ \mask sym s -> do
    when (sym == xK_Escape) $
      fail ""

    lift $ setPendingBuffer $ printf "Rec %s " s
    macro <- readMacroString
    lift $ insertWorkspaceMacroString (mask, sym) macro

readWindowsetMacro :: MaybeT X ()
readWindowsetMacro =
  mapNextStringWithKeysym $ \mask sym s -> do
    when (sym == xK_Escape) $
      fail ""

    lift $ setPendingBuffer $ printf "Rec %s " s
    macro <- readMacroString
    lift $ insertWindowSetMacroString (mask, sym) macro

readMacroString :: MaybeT X KeyString
readMacroString = do
  mapNextStringWithKeysym $ \m k s -> case (m, k, s) of
    _ | k == xK_Return -> return []
    _ | k == xK_Escape -> MaybeT $ return Nothing
    r -> do
      lift $ addStringToPendingBuffer s
      ([r] ++) <$> readMacroString

justWorkspace :: String -> Workspace
justWorkspace s =
  Workspace
    { moveWindowToWorkspaceFn = return . W.shiftWin s,
      gotoWorkspaceFn = gotoWorkspace s,
      workspaceName = Just s,
      extraWorkspaceData = ()
    }

justWorkspaceWithPreferredWindow :: Window -> String -> Workspace
justWorkspaceWithPreferredWindow w s =
  Workspace
    { moveWindowToWorkspaceFn = return . W.shiftWin s,
      gotoWorkspaceFn = do
        windows $ \ws' ->
          let ws = W.greedyView s ws'
              l = W.integrate' $ W.stack $ W.workspace $ W.current ws
           in if w `elem` l
                then W.focusWindow w ws
                else ws,
      workspaceName = Just s,
      extraWorkspaceData = ()
    }

blackHoleWorkspace :: Workspace
blackHoleWorkspace =
  Workspace
    { moveWindowToWorkspaceFn = \w -> killWindow w >> return id,
      gotoWorkspaceFn =
        confirmPrompt def "Do you want to exit xmonad" $ io exitSuccess,
      workspaceName = Nothing,
      extraWorkspaceData = ()
    }

alternateWorkspace :: Workspace
alternateWorkspace =
  Workspace
    { moveWindowToWorkspaceFn = \win -> do
        alter <- getAlternateWorkspace win
        return $ \ss ->
          maybe ss (\a -> W.shiftWin a win ss) alter,
      gotoWorkspaceFn = do
        (Location _ maybeWin) <- getCurrentLocation
        case maybeWin of
          Nothing -> return ()
          Just win -> do
            mapM_ gotoWorkspace =<< getAlternateWorkspace win,
      workspaceName = Nothing,
      extraWorkspaceData = ()
    }

newtype FloatWorkspace = FloatWorkspace Workspace

floatWorkspace :: Workspace -> Workspace
floatWorkspace ws@Workspace {extraWorkspaceData = d} =
  Workspace
    { moveWindowToWorkspaceFn = \win -> do
        case cast d of
          Just (FloatWorkspace ws') -> do
            movefn <- moveWindowToWorkspaceFn ws' win
            return $ W.sink win . movefn
          Nothing -> do
            movefn <- moveWindowToWorkspaceFn ws win
            return $ \ss -> do
              if win `Map.member` W.floating ss
                then movefn ss
                else movefn $ W.float win (W.RationalRect (1 / 8) (1 / 8) (6 / 8) (6 / 8)) ss,
      gotoWorkspaceFn = gotoWorkspaceFn ws,
      workspaceName = workspaceName ws,
      extraWorkspaceData = FloatWorkspace ws
    }

joinMaybe :: (Monad m) => MaybeT m (Maybe a) -> MaybeT m a
joinMaybe (MaybeT ma) = MaybeT $ join <$> ma

-- Like the above, but unwrap the MaybeT
workspaceForKeysT :: KeyString -> MaybeT X Workspace
workspaceForKeysT str = runKeyFeedWithKeys str readNextWorkspace

-- Like the above, but unwrap the MaybeT
workspaceForKeys :: KeyString -> X (Maybe Workspace)
workspaceForKeys = runMaybeT . workspaceForKeysT

-- Returns a workspace for the given string. Note that workspaces which are
-- defined by special masks and symbols will not work. I.e. macros will not
-- work.
workspaceForStringT :: String -> MaybeT X Workspace
workspaceForStringT = workspaceForKeysT . map (\c -> (def, def, [c]))

workspaceForString :: String -> X (Maybe Workspace)
workspaceForString = runMaybeT . workspaceForStringT

-- Like the above, but unwrap the MaybeT
locationSetForKeysT :: KeyString -> MaybeT X [Location]
locationSetForKeysT s = runKeyFeedWithKeys s readNextLocationSet

locationSetForKeys :: KeyString -> X [Location]
locationSetForKeys s = fromMaybe [] <$> runMaybeT (locationSetForKeysT s)

readNextWorkspaceName :: KeyFeed WorkspaceId
readNextWorkspaceName = absorbMaybe $ workspaceName <$> readNextWorkspace

lift1 :: (a -> X b) -> (a -> KeyFeed b)
lift1 fn = liftXToFeed . fn

-- Returns the next workspaces associated with the next set of keystrokes, or
-- invokes the handler if the next stroke is not associated with WML command.
readNextWorkspace :: KeyFeed Workspace
readNextWorkspace =
  readNextKey $ \key -> do
    macros <- liftXToFeed $ workspaceMacros <$> getMacros

    case key of
      (mask, keysym, _)
        | (Just macro) <- Map.lookup (mask, keysym) macros -> do
            pushKeys macro
            readNextWorkspace
      (_, _, [ch])
        | isAlphaNum ch || ch == '*' ->
            return $ justWorkspace [ch]
      (_, _, "[") ->
        justWorkspace
          <$> ( lift1 (adjacentWorkspaceNotVisible prev)
                  =<< readNextWorkspaceName
              )
      -- to the non-visible workspace right of the next workspace
      (_, _, "]") ->
        justWorkspace
          <$> ( lift1 (adjacentWorkspaceNotVisible next)
                  =<< readNextWorkspaceName
              )
      -- To the left of the next workspace
      (_, _, "(") ->
        justWorkspace
          <$> ( lift1 (adjacentWorkspace prev)
                  =<< readNextWorkspaceName
              )
      -- To the right of the next workspace
      (_, _, ")") ->
        justWorkspace
          <$> ( lift1 (adjacentWorkspace next)
                  =<< readNextWorkspaceName
              )
      -- The workspace on the leftmost screen
      (_, _, "^") ->
        liftXMaybe $
          withWindowSet $ \ws ->
            return $
              ( fmap
                  ( justWorkspace . W.tag . W.workspace . snd
                  )
                  . head
              )
                (getHorizontallyOrderedScreens ws)
      -- The last workspace in history.
      (_, _, "'") ->
        justWorkspace . locationWorkspace <$> liftXMaybe lastLocation
      -- The current workspace.
      (_, _, ".") ->
        liftXToFeed $ justWorkspace <$> getCurrentWorkspace
      -- The workspace on the rightmost screen
      (_, _, "$") ->
        liftXMaybe $
          withWindowSet $ \ws ->
            return $
              (fmap (justWorkspace . W.tag . W.workspace . snd) . last)
                (getHorizontallyOrderedScreens ws)
      -- Modify the next workspace as a "floating" workspace. (Windows sent to
      -- it will float).
      (_, _, ":") ->
        floatWorkspace <$> readNextWorkspace
      -- Workspace to the next screen to the right of the next workspace.
      (_, _, ",") -> do
        ws <- readNextWorkspace
        screens <-
          liftXToFeed $
            map (W.tag . W.workspace . snd)
              <$> withWindowSet (return . getHorizontallyOrderedScreens)

        let (_, rest) = break ((== workspaceName ws) . Just) (screens ++ screens)

        justWorkspace <$> hoistMaybe (head $ tail rest)

      -- Workspace to the next screen to the left of the next workspace.
      (_, _, ";") -> do
        ws <- readNextWorkspace
        screens <-
          liftXToFeed $
            map (W.tag . W.workspace . snd)
              <$> withWindowSet (return . reverse . getHorizontallyOrderedScreens)

        let (_, rest) = break ((== workspaceName ws) . Just) (screens ++ screens)

        justWorkspace <$> hoistMaybe (head $ tail rest)

      -- The workspace with the searched for window.
      (_, _, "/") ->
        justWorkspace
          <$> ( (liftXMaybe . workspaceWithWindow) =<< liftXMaybe (head <$> askWindowId)
              )
      -- The workspace with the next read window on it.
      (_, _, "@") -> do
        loc <- readNextLocationSet
        liftXMaybe $
          withWindowSet $ \ws -> return $ do
            win <- locationWindow =<< head loc
            winLocation <- W.findWindow ws win
            justWorkspaceWithPreferredWindow win . W.tag <$> W.getLocationWorkspace winLocation

      -- The accompaning worksapce to the next read workspace.
      (_, _, "~") ->
        justWorkspace . accompaningWorkspace <$> readNextWorkspaceName
      -- The accompaning workspace to the current workspace (equivalent to ~.)
      (_, _, " ") ->
        liftXToFeed $
          justWorkspace . accompaningWorkspace <$> getCurrentWorkspace
      -- The balck hole workspace
      (_, _, "_") ->
        return blackHoleWorkspace
      -- The alternate workspace
      (_, _, "-") ->
        return alternateWorkspace
      -- If the next two read workspaces are equal, go to the third workspace
      -- otherwise go to the fourth workspace.
      (_, _, "=") -> do
        ws1 <- readNextWorkspace
        ws2 <- readNextWorkspace

        ws3 <- readNextWorkspace
        ws4 <- readNextWorkspace

        return $
          if workspaceName ws1 == workspaceName ws2
            then ws3
            else ws4
      -- If the next read location set is not empty, go to the next read
      -- workspace, otherwise go to the next-next read workspace.
      (_, _, "?") -> do
        l1 <- readNextLocationSet

        ws1 <- readNextWorkspace
        ws2 <- readNextWorkspace

        return $
          if null l1
            then ws2
            else ws1
      _ -> feedFail

nonempty :: (Monad m) => m [a] -> MaybeT m [a]
nonempty l = MaybeT $ do
  l
    >>= ( \case
            [] -> return Nothing
            a -> return (Just a)
        )

readNextLocationSet :: KeyFeed [Location]
readNextLocationSet = do
  (WindowSelect mp) <- liftXToFeed XS.get
  case Map.keys mp of
    [] -> readNextLocationSet'
    wins -> do
      liftXToFeed $ addStringToPendingBuffer "<sel> "
      mapM (hoistMaybeT . windowLocation)
        =<< liftXToFeed getAndResetWindowSelection

-- Like readNextLocationSet, but ignores the window selection.
readNextLocationSet' :: KeyFeed [Location]
readNextLocationSet' =
  readNextKey $ \key -> do
    macros <- liftXToFeed $ windowsetMacros <$> getMacros

    case key of
      -- Macros takes precedence.
      (mask, keysym, _) | (Just macro) <- Map.lookup (mask, keysym) macros -> do
        hoistMaybeT $ locationSetForKeysT macro

      -- A character is the base-case. Refers to a collection of windows.
      (_, _, [ch]) | isAlpha ch -> liftXToFeed $ getMarkedLocations [ch]

      -- Goes to the most recent location in history.
      -- (_, _, "0") -> (: []) <$> MaybeT (fromX getMostRecentLocationInHistory)
      -- A Digit goes to the past history.
      -- (_, _, [ch])
      --   | isDigit ch ->
      --     (: []) <$> MaybeT (fromX $ pastHistory (ord ch - 0x30))
      -- The current window.
      (_, _, ".") -> (: []) <$> liftXToFeed getCurrentLocation
      -- The selected windows in the selection set.
      (_, _, "#") ->
        liftXToFeed $ map (Location "*" . Just) <$> pinnedWindows
      -- The window on the far-left of the screens.
      (_, _, "^") -> (: []) <$> hoistMaybeT farLeftWindow
      -- The windows on the far-right of the screens.
      (_, _, "$") -> (: []) <$> hoistMaybeT farRightWindow
      -- The next location in history.
      (_, _, "\"") -> (: []) <$> liftXMaybe nextLocation
      -- The previous location in history.
      (_, _, "'") -> (: []) <$> liftXMaybe lastLocation
      -- All visible windows.
      (_, _, "*") -> liftXToFeed $ do
        wins <- withWindowSet $ return . W.allVisibleWindows
        catMaybes <$> mapM (runMaybeT . windowLocation) wins

      -- The last referenced windows.
      (_, _, "-") ->
        hoistMaybeT $
          mapM windowLocation =<< lift getAlternateWindows
      -- Search for the windows.
      (_, _, "/") ->
        hoistMaybeT $
          mapM windowLocation =<< nonempty askWindowId
      -- All windows.
      (_, _, "%") -> hoistMaybeT $
        mapM windowLocation =<< lift (withWindowSet (return . sortOn Down . W.allWindows))
      -- Windows in a workspace
      (_, _, s)
        | s == "\t" || s == "@" || s == "\n" ->
            (liftXToFeed . windowsInWorkspace) =<< readNextWorkspaceName
      -- The first window in the next window set.
      (_, _, "!") -> (: []) <$> absorbMaybe (head <$> readNextLocationSet)
      -- The windows except the first in a window set.
      (_, _, ",") -> tail <$> readNextLocationSet
      -- The next window set, but reversed
      (_, _, "~") -> reverse <$> readNextLocationSet
      -- All the floating windows
      (_, _, ":") ->
        liftXToFeed $
          withWindowSet $
            fmap catMaybes
              . mapM (runMaybeT . windowLocation)
              . Map.keys
              . W.floating
      -- If the next read window set is not empty, then this location
      -- otherwise the next read location.
      (_, _, "?") -> do
        l1 <- readNextLocationSet
        l2 <- readNextLocationSet
        return $ if null l1 then l2 else l1

      -- The next window set unioned with the next location set
      (_, _, "|") -> do
        l1 <- readNextLocationSet
        l2 <- readNextLocationSet
        return (l1 ++ l2)

      -- Empty window set.
      (_, _, "_") -> return []
      -- The next location set differenced with the next-next location set
      (_, _, "\\") -> do
        l1 <- readNextLocationSet
        l2 <- readNextLocationSet
        return $ filter (not . flip elem l2) l1

      -- The next location set intersected with the next-next location set
      (_, _, "&") -> do
        l1 <- readNextLocationSet
        l2 <- readNextLocationSet
        return $ filter (`elem` l2) l1

      _ -> feedFail