{-# LANGUAGE OverloadedStrings #-}

-- | Command-line fuzzy selection of arbitrary values.
--     Shells out to fzf for the actual selection.
module Unison.CommandLine.FuzzySelect
  ( fuzzySelect,
    isFZFInstalled,
    fzfPathEnvVar,
    Options (..),
    FuzzySelections (..),
    defaultOptions,
  )
where

import Control.Monad.Except (throwError)
import Control.Monad.Trans.Except
import Data.Set qualified as Set
import Data.Text qualified as Text
import Data.Text.IO qualified as Text
import GHC.IO.Handle (hDuplicateTo)
import System.Environment (lookupEnv)
import System.IO (BufferMode (NoBuffering), hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)
import Unison.Prelude
import Unison.Util.Monoid qualified as Monoid
import UnliftIO qualified
import UnliftIO.Directory (findExecutable)
import UnliftIO.Exception (bracket)
import UnliftIO.IO (hGetBuffering, hSetBuffering, stdin)
import UnliftIO.Process qualified as Proc

-- | An environment variable that can be set to override the default fzf executable.
fzfPathEnvVar :: String
fzfPathEnvVar :: String
fzfPathEnvVar = String
"UNISON_FZF_PATH"

fzfExecutable :: IO (Maybe FilePath)
fzfExecutable :: IO (Maybe String)
fzfExecutable = do
  envPath <- String -> IO (Maybe String)
lookupEnv String
fzfPathEnvVar
  case (envPath) of
    Just String
path
      | Text -> Text
Text.toUpper (String -> Text
Text.pack String
path) Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"NONE" -> Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe String
forall a. Maybe a
Nothing
      | Bool
otherwise -> Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Maybe String
forall a. a -> Maybe a
Just String
path)
    Maybe String
Nothing -> String -> IO (Maybe String)
forall (m :: * -> *). MonadIO m => String -> m (Maybe String)
findExecutable String
"fzf"

isFZFInstalled :: Bool
isFZFInstalled :: Bool
isFZFInstalled =
  IO Bool -> Bool
forall a. IO a -> a
unsafePerformIO (Maybe String -> Bool
forall a. Maybe a -> Bool
isJust (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Maybe String)
fzfExecutable)
{-# NOINLINE isFZFInstalled #-}

-- | Fuzzy Selection options
data Options = Options
  { Options -> Bool
allowMultiSelect :: Bool
  }

-- | Default 'Options'
defaultOptions :: Options
defaultOptions :: Options
defaultOptions =
  Options
    { allowMultiSelect :: Bool
allowMultiSelect = Bool
True
    }

-- | Convert options into command-line args for fzf
optsToArgs :: Options -> Bool -> [String]
optsToArgs :: Options -> Bool -> [String]
optsToArgs Options
opts Bool
useNumberings =
  [String]
defaultArgs [String] -> [String] -> [String]
forall a. Semigroup a => a -> a -> a
<> case Options
opts of
    Options {allowMultiSelect :: Options -> Bool
allowMultiSelect = Bool
True} -> [String
"-m"]
    Options
_ -> []
  where
    defaultArgs :: [String]
defaultArgs =
      -- When using numberings, don't show or match on the first column of input.
      -- This allows us to prepend each line with a number, and use that number to determine
      -- which values from the input list were selected.
      Bool -> [String] -> [String]
forall a. Monoid a => Bool -> a -> a
Monoid.whenM
        Bool
useNumberings
        [String
"--with-nth", String
"2.."]
        [String] -> [String] -> [String]
forall a. Semigroup a => a -> a -> a
<> [ -- Use only half the screen (it's nice to see what you were working on when searching)
             String
"--height=50%",
             -- But if 50% of the screen is too small, ensure show at least 10 results.
             String
"--min-height=10"
           ]

data FuzzySelections a where
  SelectFromChoices :: (a -> Text) -> [a] -> FuzzySelections a
  SelectFiles :: FuzzySelections Text

-- | Allows prompting the user to interactively fuzzy-select a result from a list of options, currently shells out to `fzf` under the hood.
-- If fzf is missing, or an error (other than ctrl-c) occurred, returns Nothing.
fuzzySelect :: forall a. Options -> FuzzySelections a -> IO (Maybe [a])
fuzzySelect :: forall a. Options -> FuzzySelections a -> IO (Maybe [a])
fuzzySelect Options
opts FuzzySelections a
selections =
  (SomeException -> IO (Maybe [a]))
-> IO (Maybe [a]) -> IO (Maybe [a])
forall (m :: * -> *) a.
MonadUnliftIO m =>
(SomeException -> m a) -> m a -> m a
UnliftIO.handleAny SomeException -> IO (Maybe [a])
handleException
    (IO (Maybe [a]) -> IO (Maybe [a]))
-> (ExceptT Text IO (Maybe [a]) -> IO (Maybe [a]))
-> ExceptT Text IO (Maybe [a])
-> IO (Maybe [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Either Text (Maybe [a])) -> IO (Maybe [a])
handleError
    (IO (Either Text (Maybe [a])) -> IO (Maybe [a]))
-> (ExceptT Text IO (Maybe [a]) -> IO (Either Text (Maybe [a])))
-> ExceptT Text IO (Maybe [a])
-> IO (Maybe [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Either Text (Maybe [a])) -> IO (Either Text (Maybe [a]))
forall c. IO c -> IO c
restoreBuffering
    (IO (Either Text (Maybe [a])) -> IO (Either Text (Maybe [a])))
-> (ExceptT Text IO (Maybe [a]) -> IO (Either Text (Maybe [a])))
-> ExceptT Text IO (Maybe [a])
-> IO (Either Text (Maybe [a]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExceptT Text IO (Maybe [a]) -> IO (Either Text (Maybe [a]))
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT
    (ExceptT Text IO (Maybe [a]) -> IO (Maybe [a]))
-> ExceptT Text IO (Maybe [a]) -> IO (Maybe [a])
forall a b. (a -> b) -> a -> b
$ do
      fzfPath <-
        IO (Maybe String) -> ExceptT Text IO (Maybe String)
forall a. IO a -> ExceptT Text IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (Maybe String)
fzfExecutable ExceptT Text IO (Maybe String)
-> (Maybe String -> ExceptT Text IO String)
-> ExceptT Text IO String
forall a b.
ExceptT Text IO a -> (a -> ExceptT Text IO b) -> ExceptT Text IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe String
Nothing -> Text -> ExceptT Text IO String
forall a. Text -> ExceptT Text IO a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError Text
"I couldn't find the `fzf` executable on your path, consider installing `fzf` to enable fuzzy searching."
          Just String
fzfPath -> String -> ExceptT Text IO String
forall a. a -> ExceptT Text IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
fzfPath
      case selections of
        SelectFromChoices a -> Text
intoSearchText [a]
choices -> do
          let [String]
fzfArgs :: [String] =
                Options -> Bool -> [String]
optsToArgs Options
opts Bool
True
          let [(Int, a)]
numberedChoices :: [(Int, a)] =
                [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [a]
choices
          let [Text]
searchTexts :: [Text] =
                (\(Int
n, a
ch) -> Int -> Text
forall a. Show a => a -> Text
tShow (Int
n) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> a -> Text
intoSearchText a
ch) ((Int, a) -> Text) -> [(Int, a)] -> [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, a)]
numberedChoices

          result <- String
-> [String]
-> [Text]
-> ExceptT Text IO (Either SomeException [Text])
fzfWithChoices String
fzfPath [String]
fzfArgs [Text]
searchTexts
          -- Since we prefixed every search term with its number earlier, we know each result
          -- is prefixed with a number, we need to parse it and use it to select the matching
          -- value from our input list.
          pure $ case result of
            Left SomeException
_ -> Maybe [a]
forall a. Maybe a
Nothing
            Right [Text]
selections ->
              [Text]
selections
                [Text] -> ([Text] -> [Int]) -> [Int]
forall a b. a -> (a -> b) -> b
& (Text -> Maybe Int) -> [Text] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
forall (f :: * -> *) a b.
Filterable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (forall a. Read a => String -> Maybe a
readMaybe @Int (String -> Maybe Int) -> (Text -> String) -> Text -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
Text.unpack (Text -> String) -> (Text -> Text) -> Text -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Text -> Text
Text.takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
' '))
                [Int] -> ([Int] -> Set Int) -> Set Int
forall a b. a -> (a -> b) -> b
& [Int] -> Set Int
forall a. Ord a => [a] -> Set a
Set.fromList
                Set Int -> (Set Int -> [a]) -> [a]
forall a b. a -> (a -> b) -> b
& ( \Set Int
selectedNumbers ->
                      [(Int, a)]
numberedChoices
                        [(Int, a)] -> ([(Int, a)] -> [a]) -> [a]
forall a b. a -> (a -> b) -> b
& ((Int, a) -> Maybe a) -> [(Int, a)] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
forall (f :: * -> *) a b.
Filterable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (\(Int
n, a
a) -> if Int
n Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Int
selectedNumbers then a -> Maybe a
forall a. a -> Maybe a
Just a
a else Maybe a
forall a. Maybe a
Nothing)
                  )
                [a] -> ([a] -> Maybe [a]) -> Maybe [a]
forall a b. a -> (a -> b) -> b
& [a] -> Maybe [a]
forall a. a -> Maybe a
Just
        FuzzySelections a
SelectFiles -> do
          let [String]
fzfArgs :: [String] = Options -> Bool -> [String]
optsToArgs Options
opts Bool
False
          Either SomeException [a] -> Maybe [a]
forall a b. Either a b -> Maybe b
eitherToMaybe (Either SomeException [a] -> Maybe [a])
-> ExceptT Text IO (Either SomeException [a])
-> ExceptT Text IO (Maybe [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> [String] -> ExceptT Text IO (Either SomeException [Text])
fzfFileSelector String
fzfPath [String]
fzfArgs
  where
    fzfWithChoices :: FilePath -> [String] -> [Text] -> ExceptT Text IO (Either SomeException [Text])
    fzfWithChoices :: String
-> [String]
-> [Text]
-> ExceptT Text IO (Either SomeException [Text])
fzfWithChoices String
fzfPath [String]
fzfArgs [Text]
searchTexts = do
      let CreateProcess
fzfProc :: Proc.CreateProcess =
            (String -> [String] -> CreateProcess
Proc.proc String
fzfPath [String]
fzfArgs)
              { Proc.std_in = Proc.CreatePipe,
                Proc.std_out = Proc.CreatePipe,
                Proc.delegate_ctlc = True
              }
      (Just stdin', Just stdout', _, procHandle) <- CreateProcess
-> ExceptT
     Text IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
forall (m :: * -> *).
MonadIO m =>
CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
Proc.createProcess CreateProcess
fzfProc
      -- Generally no-buffering is helpful for highly interactive processes.
      hSetBuffering stdin NoBuffering
      hSetBuffering stdin' NoBuffering
      liftIO . UnliftIO.tryAny $ do
        -- Dump the search terms into fzf's stdin
        traverse_ (Text.hPutStrLn stdin') searchTexts
        -- Wire up the interactive terminal to fzf now that the inputs have been loaded.
        hDuplicateTo stdin stdin'
        void $ Proc.waitForProcess procHandle
        Text.lines <$> liftIO (Text.hGetContents stdout')
    fzfFileSelector :: FilePath -> [String] -> ExceptT Text IO (Either SomeException [Text])
    fzfFileSelector :: String -> [String] -> ExceptT Text IO (Either SomeException [Text])
fzfFileSelector String
fzfPath [String]
fzfArgs = do
      let CreateProcess
fzfProc :: Proc.CreateProcess =
            (String -> [String] -> CreateProcess
Proc.proc String
fzfPath [String]
fzfArgs)
              { Proc.std_in = Proc.Inherit,
                Proc.std_out = Proc.CreatePipe,
                Proc.delegate_ctlc = True
              }
      (_stdin', Just stdout', _, procHandle) <- CreateProcess
-> ExceptT
     Text IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
forall (m :: * -> *).
MonadIO m =>
CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
Proc.createProcess CreateProcess
fzfProc
      -- Generally no-buffering is helpful for highly interactive processes.
      hSetBuffering stdin NoBuffering
      liftIO . UnliftIO.tryAny $ do
        void $ Proc.waitForProcess procHandle
        Text.lines <$> liftIO (Text.hGetContents stdout')
    handleException :: SomeException -> IO (Maybe [a])
    handleException :: SomeException -> IO (Maybe [a])
handleException SomeException
err = SomeException -> IO ()
forall a (f :: * -> *). (Show a, Applicative f) => a -> f ()
traceShowM SomeException
err IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Handle -> String -> IO ()
hPutStrLn Handle
stderr String
"Oops, something went wrong. No input selected." IO () -> IO (Maybe [a]) -> IO (Maybe [a])
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Maybe [a] -> IO (Maybe [a])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [a]
forall a. Maybe a
Nothing
    handleError :: IO (Either Text (Maybe [a])) -> IO (Maybe [a])
    handleError :: IO (Either Text (Maybe [a])) -> IO (Maybe [a])
handleError IO (Either Text (Maybe [a]))
m =
      IO (Either Text (Maybe [a]))
m IO (Either Text (Maybe [a]))
-> (Either Text (Maybe [a]) -> IO (Maybe [a])) -> IO (Maybe [a])
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Left Text
err -> Handle -> Text -> IO ()
Text.hPutStrLn Handle
stderr Text
err IO () -> IO (Maybe [a]) -> IO (Maybe [a])
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Maybe [a] -> IO (Maybe [a])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [a]
forall a. Maybe a
Nothing
        Right Maybe [a]
as -> Maybe [a] -> IO (Maybe [a])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [a]
as
    restoreBuffering :: IO c -> IO c
    restoreBuffering :: forall c. IO c -> IO c
restoreBuffering IO c
action =
      IO BufferMode
-> (BufferMode -> IO ()) -> (BufferMode -> IO c) -> IO c
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (Handle -> IO BufferMode
forall (m :: * -> *). MonadIO m => Handle -> m BufferMode
hGetBuffering Handle
stdin) (Handle -> BufferMode -> IO ()
forall (m :: * -> *). MonadIO m => Handle -> BufferMode -> m ()
hSetBuffering Handle
stdin) (IO c -> BufferMode -> IO c
forall a b. a -> b -> a
const IO c
action)