Run.hs

{-# LANGUAGE BangPatterns #-}

module YMonad.Run (
  ymonad,
) where

import Control.Concurrent (forkIO, killThread)
import Control.Concurrent.STM (TQueue, newTQueueIO, tryReadTQueue, writeTQueue)
import Control.Exception (AsyncException (ThreadKilled), bracket, try)
import Control.Monad (foldM)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as B8
import Data.Map.Strict qualified as M
import Data.Text qualified as T
import Data.Text.IO qualified as TIO
import Network.Socket
import Network.Socket.ByteString qualified as NSB
import System.Directory (removePathForcibly)
import System.Environment (getEnvironment, getExecutablePath)
import System.FilePath ((</>))
import System.IO.Error (catchIOError)
import System.Process (createProcess, env, proc, shell)
import System.Timeout (timeout)

import YMonad.Backend.River
import YMonad.Config
import YMonad.ControlCommand
import YMonad.Core
import YMonad.Keys
import YMonad.Prompt.Internal (choosePrompt, getCommands, getMans, getSshHosts, runInConfiguredTerminal)

ymonad :: YConfig -> IO ()
ymonad conf = do
  controlPath <- getControlPath
  queue <- newTQueueIO
  withControlServer controlPath queue $ do
    logTrace ("control socket ready: " <> toText controlPath)
    runtimeResult <- loadRiverRuntime
    riverRt0 <- case runtimeResult of
      Left err -> die ("Failed to load river protocol: " <> err)
      Right rt -> pure rt
    socketResult <- connectRiver Nothing
    sock <- case socketResult of
      Left fault -> die ("Failed to connect to river: " <> toString (faultMessage fault))
      Right connected -> pure connected
    bracket
      (pure sock)
      close
      ( \activeSock -> do
          startupResult <- acquireRiverManager activeSock riverRt0
          (runtimeSt0, riverRt1) <- case startupResult of
            Left fault -> die ("Failed to acquire river window manager: " <> toString (faultMessage fault))
            Right (startupEvents, riverRt') -> foldM (\(st, rt) event -> applyEvent st rt event) (runtimeStateFromConfig conf, riverRt') startupEvents
          riverRt2 <- maybeInstallKeyBindings conf riverRt1
          loop conf activeSock queue runtimeSt0 riverRt2
      )

runtimeStateFromConfig :: YConfig -> RuntimeState
runtimeStateFromConfig conf =
  runtimeSt0
    { rtWindowSet =
        windowSet0
          { current = (current windowSet0) {screenWorkspace = seedLayout (screenWorkspace (current windowSet0))}
          , hidden = map seedLayout (hidden windowSet0)
          }
    , rtTheme =
        RuntimeTheme
          { themeBorderWidth = borderWidth conf
          , themeNormalBorderColor = fromMaybe (opaqueRgb8 0xdd 0xdd 0xdd) (parseColorString (normalBorderColor conf))
          , themeFocusedBorderColor = fromMaybe (opaqueRgb8 0xff 0x00 0x00) (parseColorString (focusedBorderColor conf))
          , themeTitleBarMode = titleBarMode conf
          }
    }
  where
    runtimeSt0 = defaultRuntimeStateWithWorkspaces (Rect 0 0 1920 1080) workspaceIds
    windowSet0 = rtWindowSet runtimeSt0

    seedLayout workspace = workspace {workspaceLayout = layoutHook conf}

    workspaceIds = case map (WorkspaceId . T.pack) (workspaces conf) of
      [] -> WorkspaceId "1" :| map (WorkspaceId . show) [(2 :: Int) .. 9]
      firstWorkspace : remainingWorkspaces -> firstWorkspace :| remainingWorkspaces

data RiverPoll = RiverPoll
  { riverPollRuntime :: RiverRuntime
  , riverPollCallbacks :: [RiverCallback]
  , riverPollSteps :: [RiverPollStep]
  }

data RiverPollStep
  = RiverPollAction YAction
  | RiverPollEvent YEvent

loop :: YConfig -> Socket -> TQueue ControlCommand -> RuntimeState -> RiverRuntime -> IO ()
loop conf sock queue runtimeSt riverRt = do
  polled <- pollRiverEvents sock riverRt
  case polled of
    Left fault -> stopAfterFault runtimeSt riverRt fault
    Right poll -> do
      logCallbacks (riverPollCallbacks poll)
      (runtimeSt0, riverRt0) <- foldM (applyPollStep conf) (runtimeSt, riverPollRuntime poll) (riverPollSteps poll)

      pendingCommands <- atomically (drainQueue queue)
      (runtimeSt1, riverRt1) <- foldM applyControlCommand (runtimeSt0, riverRt0) pendingCommands
      riverRt2 <- maybeInstallKeyBindings conf riverRt1

      flushed <- flushTransportLogged sock riverRt2
      case flushed of
        Left fault -> stopAfterFault runtimeSt1 riverRt2 fault
        Right riverRt3 -> loop conf sock queue runtimeSt1 riverRt3
  where
    stopAfterFault runtimeSt' riverRt' fault = do
      _ <- applyEvent runtimeSt' riverRt' (EvBackendFault fault)
      logTrace "stopping river event loop after transport failure"

applyPollStep :: YConfig -> (RuntimeState, RiverRuntime) -> RiverPollStep -> IO (RuntimeState, RiverRuntime)
applyPollStep conf state' step =
  case step of
    RiverPollAction action -> applyAction conf state' action
    RiverPollEvent event -> let (runtimeSt, riverRt) = state' in applyEvent runtimeSt riverRt event

applyEvent :: RuntimeState -> RiverRuntime -> YEvent -> IO (RuntimeState, RiverRuntime)
applyEvent runtimeSt riverRt event = do
  case event of
    EvBackendFault fault -> logBackendFault fault
    _otherEvent -> pass
  let transition = reduce runtimeSt event
      requests = trRequests transition
      !requestCount = length requests
      !runtimeSt' = trState transition
      !riverRt' = pumpPending (stageRequests riverRt requests)
  requestCount `seq` pass
  pure (runtimeSt', riverRt')

applyAction :: YConfig -> (RuntimeState, RiverRuntime) -> YAction -> IO (RuntimeState, RiverRuntime)
applyAction conf (runtimeSt, riverRt) action =
  case action of
    Spawn cmd -> do
      _ <- createProcess (shell cmd)
      pure (runtimeSt, riverRt)
    ToggleFullscreen ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdToggleFullscreen)
    ViewWorkspace workspaceName ->
      applyEvent runtimeSt riverRt (EvUserCommand (CmdViewWorkspace (WorkspaceId (T.pack workspaceName))))
    MoveFocusedTo workspaceName ->
      applyEvent runtimeSt riverRt (EvUserCommand (CmdMoveFocusedTo (WorkspaceId (T.pack workspaceName))))
    FocusNext ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdFocusNext)
    FocusPrev ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdFocusPrev)
    SwapNext ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdSwapNext)
    SwapPrev ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdSwapPrev)
    KillFocused ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdKillFocused)
    CycleLayout ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdCycleLayout)
    ShrinkMaster ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdShrink)
    ExpandMaster ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdExpand)
    IncMasterNBy delta ->
      applyEvent runtimeSt riverRt (EvUserCommand (CmdIncMasterN delta))
    ExitSession ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdExitSession)
    RecompileRestart -> do
      restartYMonad
      pure (runtimeSt, riverRt)
    ShellPrompt xpConfig -> do
      commands <- getCommands
      selection <- choosePrompt xpConfig "Run: " commands
      case selection of
        Nothing -> pure (runtimeSt, riverRt)
        Just commandText -> do
          _ <- createProcess (shell commandText)
          pure (runtimeSt, riverRt)
    SafePrompt commandPath xpConfig -> do
      selection <- choosePrompt xpConfig "Run: " []
      case selection of
        Nothing -> pure (runtimeSt, riverRt)
        Just argument -> do
          _ <- createProcess (proc commandPath [argument])
          pure (runtimeSt, riverRt)
    UnsafePrompt commandPrefix xpConfig -> do
      selection <- choosePrompt xpConfig "Run: " []
      case selection of
        Nothing -> pure (runtimeSt, riverRt)
        Just argument -> do
          _ <- createProcess (shell (commandPrefix <> " " <> argument))
          pure (runtimeSt, riverRt)
    SshPrompt xpConfig -> do
      hosts <- getSshHosts
      selection <- choosePrompt xpConfig "SSH to: " hosts
      case selection of
        Nothing -> pure (runtimeSt, riverRt)
        Just host -> do
          runInConfiguredTerminal conf xpConfig ("ssh " <> host)
          pure (runtimeSt, riverRt)
    ManPrompt xpConfig -> do
      pages <- getMans
      selection <- choosePrompt xpConfig "Manual page: " pages
      case selection of
        Nothing -> pure (runtimeSt, riverRt)
        Just page -> do
          runInConfiguredTerminal conf xpConfig ("man " <> page)
          pure (runtimeSt, riverRt)

restartYMonad :: IO ()
restartYMonad = do
  launcherPath <- getLauncherPath
  env0 <- getEnvironment
  let env1 =
        ("YMONAD_FORCE_REBUILD", "1")
          : filter ((/= "YMONAD_FORCE_REBUILD") . fst) env0
  _ <- createProcess ((proc launcherPath []) {env = Just env1})
  logTrace ("spawned restart via " <> toText launcherPath)
  exitSuccess

getLauncherPath :: IO FilePath
getLauncherPath = do
  explicit <- lookupEnv "YMONAD_LAUNCHER_PATH"
  maybe getExecutablePath pure explicit

applyControlCommand :: (RuntimeState, RiverRuntime) -> ControlCommand -> IO (RuntimeState, RiverRuntime)
applyControlCommand (runtimeSt, riverRt) command =
  case command of
    ControlSpawn cmd -> do
      _ <- createProcess (shell cmd)
      pure (runtimeSt, riverRt)
    ControlToggleFullscreen ->
      applyEvent runtimeSt riverRt (EvUserCommand CmdToggleFullscreen)
    ControlUserCommand userCmd ->
      applyEvent runtimeSt riverRt (EvUserCommand userCmd)

pollRiverEvents :: Socket -> RiverRuntime -> IO (Either BackendFault RiverPoll)
pollRiverEvents sock riverRt = do
  recvResult <- (try (timeout 10000 (NSB.recv sock 65535)) :: IO (Either SomeException (Maybe BS.ByteString)))
  case recvResult of
    Left ex -> pure (Left (BackendFault ("Wayland recv failed: " <> toText (displayException ex))))
    Right mChunk ->
      case mChunk of
        Nothing -> pure (Right (RiverPoll riverRt [] []))
        Just chunk
          | BS.null chunk -> pure (Left (BackendFault "Wayland socket closed by compositor"))
          | otherwise ->
              pure $ do
                (riverRt', callbacks) <- ingestBytes riverRt chunk
                steps <- fmap concat (mapM callbackToSteps callbacks)
                let !callbackCount = length callbacks
                    !stepCount = length steps
                callbackCount `seq` stepCount `seq` Right (RiverPoll riverRt' callbacks steps)

pollEventsFromSteps :: [RiverPollStep] -> [YEvent]
pollEventsFromSteps steps = [event | RiverPollEvent event <- steps]

callbackToSteps :: RiverCallback -> Either BackendFault [RiverPollStep]
callbackToSteps callback =
  case callback of
    RiverActionTriggered action -> Right [RiverPollAction action]
    _otherCallback -> map RiverPollEvent <$> decodeCallback callback

acquireRiverManager :: Socket -> RiverRuntime -> IO (Either BackendFault ([YEvent], RiverRuntime))
acquireRiverManager sock riverRt0 = do
  flushed <- flushTransportLogged sock riverRt0
  case flushed of
    Left fault -> pure (Left fault)
    Right riverRt1 -> waitForRegistryRoundtrip 200 [] riverRt1
  where
    waitForRegistryRoundtrip :: Int -> [YEvent] -> RiverRuntime -> IO (Either BackendFault ([YEvent], RiverRuntime))
    waitForRegistryRoundtrip remaining acc riverRt
      | startupReady riverRt =
          case rrManagerObject riverRt of
            Nothing ->
              pure
                ( Left
                    ( BackendFault
                        "Compositor did not advertise river_window_manager_v1 during startup; YMonad requires River's window management protocol"
                    )
                )
            Just _managerObject -> do
              logTrace "river_window_manager_v1 acquired"
              pure (Right (acc, riverRt))
      | remaining <= 0 =
          pure
            ( Left
                ( BackendFault
                    "Timed out waiting for wl_registry roundtrip while acquiring river_window_manager_v1"
                )
            )
      | otherwise = do
          polled <- pollRiverEvents sock riverRt
          case polled of
            Left fault -> pure (Left fault)
            Right poll -> do
              logCallbacks (riverPollCallbacks poll)
              let pollEvents = pollEventsFromSteps (riverPollSteps poll)
              case firstBackendFault pollEvents of
                Just fault -> pure (Left fault)
                Nothing -> do
                  flushed' <- flushTransportLogged sock (riverPollRuntime poll)
                  case flushed' of
                    Left fault -> pure (Left fault)
                    Right riverRt' -> waitForRegistryRoundtrip (remaining - 1) (acc ++ pollEvents) riverRt'

    startupReady :: RiverRuntime -> Bool
    startupReady riverRt =
      rrRegistryReady riverRt
        && isJust (rrManagerObject riverRt)
        && case rrXkbBindingsObject riverRt of
          Nothing -> True
          Just _ -> isJust (rrSeat riverRt)

flushTransportLogged :: Socket -> RiverRuntime -> IO (Either BackendFault RiverRuntime)
flushTransportLogged sock riverRt = do
  let prepared = pumpPending riverRt
  logOutgoingMessages (rrOutgoing prepared)
  flushTransportIO sock riverRt

maybeInstallKeyBindings :: YConfig -> RiverRuntime -> IO RiverRuntime
maybeInstallKeyBindings conf riverRt
  | rrBindingsInstalled riverRt = pure riverRt
  | otherwise =
      case (rrXkbBindingsObject riverRt, rrSeat riverRt) of
        (Nothing, _)
          | rrRegistryReady riverRt -> do
              logTrace "river_xkb_bindings_v1 was not advertised; key bindings are unavailable"
              pure (riverRt {rrBindingsInstalled = True})
        (Just _bindingsObject, Just _seat) -> do
          let (warnings, riverRt') = stageConfiguredBindings conf riverRt
          mapM_ logTrace warnings
          logTrace "installed river xkb key bindings"
          pure riverRt'
        _notReadyYet -> pure riverRt

stageConfiguredBindings :: YConfig -> RiverRuntime -> ([Text], RiverRuntime)
stageConfiguredBindings conf riverRt0 =
  let (warnings, riverRt1) = foldl' registerBinding ([], riverRt0) (M.toList (keys conf conf))
      riverRt2 = riverRt1 {rrBindingsInstalled = True}
   in (warnings, riverRt2)
  where
    bindingsObjectId = fromMaybe 0 (rrXkbBindingsObject riverRt0)
    seatObjectId = maybe 0 riverSeatObjectId (rrSeat riverRt0)

    registerBinding :: ([Text], RiverRuntime) -> ((KeyMask, KeySym), YAction) -> ([Text], RiverRuntime)
    registerBinding (warnings, riverRt) ((mask, keySym), action) =
      case keySymToXkb keySym of
        Nothing ->
          ( warnings ++ ["skipping unsupported key symbol for river xkb bindings: " <> renderKeySym keySym]
          , riverRt
          )
        Just keysym ->
          let (bindingObjectId, riverRt1) = allocateObjectId riverRt
              createBinding =
                WireMessage
                  { wireObjectId = bindingsObjectId
                  , wireInterface = "river_xkb_bindings_v1"
                  , wireMessageName = "get_xkb_binding"
                  , wireOpcode = requestOpcode riverRt1 "river_xkb_bindings_v1" "get_xkb_binding"
                  , wireArgs =
                      [ ArgObjectValue (Just seatObjectId)
                      , ArgNewIdValue bindingObjectId
                      , ArgUintValue keysym
                      , ArgUintValue (unKeyMask mask)
                      ]
                  }
              enableBinding =
                WireMessage
                  { wireObjectId = bindingObjectId
                  , wireInterface = "river_xkb_binding_v1"
                  , wireMessageName = "enable"
                  , wireOpcode = requestOpcode riverRt1 "river_xkb_binding_v1" "enable"
                  , wireArgs = []
                  }
              riverRt2 =
                riverRt1
                  { rrBindingActions = M.insert bindingObjectId action (rrBindingActions riverRt1)
                  , rrObjectInterfaces = M.insert bindingObjectId "river_xkb_binding_v1" (rrObjectInterfaces riverRt1)
                  , rrOutgoing = rrOutgoing riverRt1 ++ [createBinding]
                  , rrPendingManage = rrPendingManage riverRt1 ++ [enableBinding]
                  }
           in (warnings, riverRt2)

requestOpcode :: RiverRuntime -> Text -> Text -> Word16
requestOpcode riverRt ifaceName name = maybe 0 messageSpecOpcode (lookupRequest (rrProtocol riverRt) ifaceName name)

logCallbacks :: [RiverCallback] -> IO ()
logCallbacks = mapM_ logCallback

logCallback :: RiverCallback -> IO ()
logCallback callback = case callback of
  RiverTrace msg -> logTrace msg
  RiverProtocolFault _msg -> pass
  RiverViewCreated vid _meta -> logTrace ("view created: " <> renderText vid)
  RiverViewDestroyed vid -> logTrace ("view destroyed: " <> renderText vid)
  RiverOutputDiscovered oid info ->
    logTrace ("output discovered: " <> renderText oid <> " " <> T.pack (show (outputRect info)))
  RiverOutputLost oid -> logTrace ("output removed: " <> renderText oid)
  RiverSeatFocus mVid -> logTrace ("seat focus: " <> maybe "none" renderText mVid)
  RiverManageStart -> pass
  RiverRenderStart -> pass
  RiverUserCommand userCommand -> logTrace ("user command: " <> renderText userCommand)
  RiverActionTriggered action -> logTrace ("binding pressed: " <> renderText action)

logOutgoingMessages :: [WireMessage] -> IO ()
logOutgoingMessages = mapM_ (logTrace . ("-> " <>) . renderWireMessage)

renderWireMessage :: WireMessage -> Text
renderWireMessage msg =
  wireInterface msg
    <> "."
    <> wireMessageName msg
    <> "(object="
    <> renderText (wireObjectId msg)
    <> ", args=["
    <> T.intercalate ", " (map renderArgValue (wireArgs msg))
    <> "])"

renderArgValue :: ArgValue -> Text
renderArgValue arg = case arg of
  ArgIntValue n -> renderText n
  ArgUintValue n -> renderText n
  ArgStringValue Nothing -> "null"
  ArgStringValue (Just txt) -> txt
  ArgObjectValue Nothing -> "object:null"
  ArgObjectValue (Just oid) -> "object:" <> renderText oid
  ArgNewIdValue oid -> "new_id:" <> renderText oid

renderText :: (Show a) => a -> Text
renderText = T.pack . show

logBackendFault :: BackendFault -> IO ()
logBackendFault fault = logTrace ("backend fault: " <> faultMessage fault)

logTrace :: Text -> IO ()
logTrace msg = TIO.hPutStrLn stderr ("[ymonad] " <> msg)

firstBackendFault :: [YEvent] -> Maybe BackendFault
firstBackendFault [] = Nothing
firstBackendFault (event : rest) = case event of
  EvBackendFault fault -> Just fault
  _otherEvent -> firstBackendFault rest

getControlPath :: IO FilePath
getControlPath = do
  explicit <- lookupEnv "YMONAD_SOCKET"
  case explicit of
    Just path -> pure path
    Nothing -> do
      runtimeDir <- lookupEnv "XDG_RUNTIME_DIR"
      case runtimeDir of
        Just dir -> pure (dir </> "ymonad.sock")
        Nothing -> pure "/tmp/ymonad.sock"

withControlServer :: FilePath -> TQueue ControlCommand -> IO a -> IO a
withControlServer path queue action = bracket acquire releaseResource (const action)
  where
    acquire = do
      removePathForcibly path `catchIOError` const pass
      server <- socket AF_UNIX Stream defaultProtocol
      bind server (SockAddrUnix path)
      listen server 5
      tid <- forkIO (acceptLoop server)
      pure (server, tid)

    releaseResource (server, tid) = do
      killThread tid
      close server
      removePathForcibly path `catchIOError` const pass

    acceptLoop server = forever $ do
      accepted <- (try (accept server) :: IO (Either SomeException (Socket, SockAddr)))
      case accepted of
        Left ex
          | isControlServerShutdown ex -> pass
          | otherwise -> logTrace ("control socket accept failed: " <> toText (displayException ex))
        Right (conn, _) -> do
          commandResult <- (try (recvAll conn) :: IO (Either SomeException BS.ByteString))
          close conn
          case commandResult of
            Left ex -> logTrace ("control socket recv failed: " <> toText (displayException ex))
            Right commandBytes ->
              case readMaybe (B8.unpack commandBytes) of
                Nothing -> logTrace ("ignored malformed control command: " <> decodeUtf8With lenientDecode commandBytes)
                Just command -> do
                  _ <- evaluateWHNF command
                  atomically (writeTQueue queue command)

isControlServerShutdown :: SomeException -> Bool
isControlServerShutdown ex =
  case fromException ex of
    Just ThreadKilled -> True
    _ -> "Bad file descriptor" `T.isInfixOf` toText (displayException ex)

recvAll :: Socket -> IO BS.ByteString
recvAll conn = go id
  where
    go chunks = do
      chunk <- NSB.recv conn 4096
      if BS.null chunk
        then pure (BS.concat (chunks []))
        else go (chunks . (chunk :))

drainQueue :: TQueue a -> STM [a]
drainQueue queue = go id
  where
    go acc = do
      next <- tryReadTQueue queue
      case next of
        Nothing -> pure (acc [])
        Just value -> go (acc . (value :))