Internal.hs

module YMonad.Prompt.Internal (
  choosePrompt,
  getCommands,
  getSshHosts,
  getMans,
  getCommandOutput,
  runInConfiguredTerminal,
  split,
) where

import Control.Concurrent (forkIO)
import Data.Char (isDigit, isSpace)
import Data.List (dropWhileEnd)
import Data.String qualified as String
import System.Directory (doesDirectoryExist, doesFileExist, executable, findExecutable, getDirectoryContents, getHomeDirectory, getPermissions, listDirectory)
import System.Exit (ExitCode (..))
import System.FilePath (dropExtensions, takeFileName, (</>))
import System.IO qualified as SIO
import System.IO.Error (catchIOError)
import System.Process (createProcess, proc, readCreateProcessWithExitCode, shell, waitForProcess)

import YMonad.Config (YConfig (..))
import YMonad.Prompt (YPConfig (..))

choosePrompt :: YPConfig -> String -> [String] -> IO (Maybe String)
choosePrompt xpConfig promptLabel rawChoices = do
  resolved <- resolvePromptProgram xpConfig promptLabel
  case resolved of
    Nothing -> pure Nothing
    Just (programPath, args) -> do
      let choices = ordNub rawChoices
          stdinPayload = if null choices then "\n" else String.unlines choices
      (exitCode, stdoutText, _stderrText) <- readCreateProcessWithExitCode (proc programPath args) stdinPayload
      case exitCode of
        ExitSuccess -> pure (nonEmptySelection stdoutText)
        ExitFailure _ -> pure Nothing

getCommands :: IO [String]
getCommands = do
  pathValue <- fromMaybe "" <$> lookupEnv "PATH"
  let directories = filter (not . null) (split ':' pathValue)
  (fmap (ordNub . concat) . forM directories) $ \dir -> do
    exists <- doesDirectoryExist dir
    if not exists
      then pure []
      else do
        names <- catchIOError (listDirectory dir) (const (pure []))
        filterM (isExecutableFile . (dir </>)) names

getSshHosts :: IO [String]
getSshHosts = do
  localHosts <- getLocalSshHosts
  globalHosts <- getGlobalSshHosts
  pure (ordNub (localHosts ++ globalHosts))

getMans :: IO [String]
getMans = do
  pathsText <- do
    p1 <- getCommandOutput "manpath -g 2>/dev/null"
    p2 <- getCommandOutput "manpath 2>/dev/null"
    pure (intercalate ":" (String.lines (p1 ++ p2)))
  let sections = ["man" <> show n | n <- [(1 :: Int) .. 9]]
      directories = [baseDir </> section | baseDir <- split ':' pathsText, not (null baseDir), section <- sections]
  nested <- forM (ordNub directories) $ \dir -> do
    exists <- doesDirectoryExist dir
    if exists
      then map dropExtensions <$> getDirectoryContents dir
      else pure []
  pure (filter validEntry (ordNub (concat nested)))

getCommandOutput :: String -> IO String
getCommandOutput cmd = do
  (exitCode, stdoutText, _stderrText) <- readCreateProcessWithExitCode (shell cmd) ""
  pure $ case exitCode of
    ExitSuccess -> stdoutText
    ExitFailure _ -> ""

runInConfiguredTerminal :: YConfig -> YPConfig -> String -> IO ()
runInConfiguredTerminal yConfig xpConfig commandText = do
  let terminalCommand = fromMaybe (terminal yConfig) (promptTerminal xpConfig)
      fullCommand = terminalCommand <> " -e sh -lc " <> shellQuote commandText
  (_, _, _, processHandle) <- createProcess (shell fullCommand)
  void (forkIO (void (waitForProcess processHandle)))

resolvePromptProgram :: YPConfig -> String -> IO (Maybe (FilePath, [String]))
resolvePromptProgram xpConfig promptLabel = do
  resolvedProgram <- case promptProgram xpConfig of
    Just explicit -> pure (Just explicit)
    Nothing -> firstJustM findExecutable ["wmenu", "bemenu", "fuzzel", "tofi", "rofi", "dmenu"]
  pure $ fmap (\programPath -> (programPath, promptArgs (takeFileName programPath) xpConfig promptLabel)) resolvedProgram

promptArgs :: FilePath -> YPConfig -> String -> [String]
promptArgs programName xpConfig promptLabel =
  case programName of
    "wmenu" -> promptProgramArgs xpConfig ++ ignoreCaseArg ++ listLinesArg ++ ["-p", promptLabel]
    "bemenu" -> promptProgramArgs xpConfig ++ ignoreCaseArg ++ ["-p", promptLabel]
    "fuzzel" -> promptProgramArgs xpConfig ++ ["--dmenu", "--prompt", promptLabel]
    "tofi" -> promptProgramArgs xpConfig ++ ["--prompt-text", promptLabel]
    "rofi" -> promptProgramArgs xpConfig ++ ["-dmenu", "-p", promptLabel] ++ ignoreCaseArg
    _ -> promptProgramArgs xpConfig ++ ["-p", promptLabel] ++ ignoreCaseArg
  where
    ignoreCaseArg = ["-i" | not (promptCaseSensitive xpConfig) && programName `elem` ["wmenu", "bemenu", "rofi", "dmenu"]]
    listLinesArg = maybe [] (\lineCount -> ["-l", show lineCount]) (promptListLines xpConfig)

nonEmptySelection :: String -> Maybe String
nonEmptySelection output =
  let stripped = strip output
   in if null stripped then Nothing else Just stripped

{-
ordNub :: (Ord a) => [a] -> [a]
ordNub = go S.empty
  where
    go _seen [] = []
    go seen (x : xs)
      | S.member x seen = go seen xs
      | otherwise = x : go (S.insert x seen) xs
-}

split :: (Eq a) => a -> [a] -> [[a]]
split delimiter = go
  where
    go [] = [[]]
    go (x : xs)
      | x == delimiter = [] : go xs
      | otherwise =
          case go xs of
            [] -> [[x]]
            (segment : segments) -> (x : segment) : segments

strip :: String -> String
strip = dropWhileEnd isSpace . dropWhile isSpace

shellQuote :: String -> String
shellQuote value = "'" <> concatMap escapeChar value <> "'"
  where
    escapeChar '\'' = "'\\''"
    escapeChar ch = [ch]

validEntry :: String -> Bool
validEntry entry =
  not (null entry)
    && entry /= "."
    && entry /= ".."

isExecutableFile :: FilePath -> IO Bool
isExecutableFile path = do
  exists <- doesFileExist path
  if not exists
    then pure False
    else do
      permissions <- getPermissions path
      pure (executable permissions)

getLocalSshHosts :: IO [String]
getLocalSshHosts = do
  homeDir <- getHomeDirectory
  knownHosts <- sshHostsFromKnownHosts (homeDir </> ".ssh" </> "known_hosts")
  configHosts <- sshHostsFromConfig (homeDir </> ".ssh" </> "config")
  pure (knownHosts ++ configHosts)

getGlobalSshHosts :: IO [String]
getGlobalSshHosts = do
  envKnownHosts <- lookupEnv "SSH_KNOWN_HOSTS"
  let candidates = maybeToList envKnownHosts ++ ["/usr/local/etc/ssh/ssh_known_hosts", "/usr/local/etc/ssh_known_hosts", "/etc/ssh/ssh_known_hosts", "/etc/ssh_known_hosts"]
  existing <- filterM doesFileExist candidates
  case existing of
    [] -> pure []
    firstKnownHosts : _rest -> sshHostsFromKnownHosts firstKnownHosts

sshHostsFromKnownHosts :: FilePath -> IO [String]
sshHostsFromKnownHosts path = do
  exists <- doesFileExist path
  if not exists
    then pure []
    else do
      contents <- catchIOError (SIO.readFile path) (const (pure ""))
      pure
        [ getWithPort host
        | line <- String.lines contents
        , nonComment line
        , firstField <- take 1 (String.words line)
        , host <- split ',' firstField
        , not (null host)
        , not ("|" `isPrefixOf` host)
        ]

sshHostsFromConfig :: FilePath -> IO [String]
sshHostsFromConfig path = do
  exists <- doesFileExist path
  if not exists
    then pure []
    else do
      contents <- catchIOError (SIO.readFile path) (const (pure ""))
      pure
        [ candidate
        | ws <- map String.words (String.lines contents)
        , isHostDirective ws
        , candidate <- drop 1 ws
        , isConcreteHost candidate
        ]

isHostDirective :: [String] -> Bool
isHostDirective (directive : _values) = directive == "Host"
isHostDirective _ = False

isConcreteHost :: String -> Bool
isConcreteHost value = not (null value) && not (any (`elem` ['*', '?']) value) && value /= "!"

getWithPort :: String -> String
getWithPort host =
  case break (== ':') host of
    (baseHost, []) -> baseHost
    (baseHost, suffix) ->
      case suffix of
        ':' : portDigits | all isDigit portDigits -> baseHost
        _other -> host

nonComment :: String -> Bool
nonComment line =
  case dropWhile isSpace line of
    [] -> False
    ('#' : _) -> False
    _ -> True

firstJustM :: (Monad m) => (a -> m (Maybe b)) -> [a] -> m (Maybe b)
firstJustM _ [] = pure Nothing
firstJustM f (x : xs) = do
  result <- f x
  case result of
    Just value -> pure (Just value)
    Nothing -> firstJustM f xs