Unix.hs

{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE TypeFamilies      #-}
{-# LANGUAGE ViewPatterns      #-}

{-|
Module      : Mangrove.Scheme.Unix
Copyright   : (c) Quytelda Kahja, 2026
License     : BSD-3-Clause

A parsing scheme for Unix-style command line arguments.
-}
module Mangrove.Scheme.Unix
  ( -- * Describing Commands & Options
    Flag(..)
  , OptionInfo(..)
  , CommandInfo(..)

    -- * Unix Scheme
  , UnixParser
  , UnixScheme(..)
  , Token(..)
  , UnixRequest(..)
  , UnixRequest'
  , helpRequest
  , versionRequest

    -- * Help
  , addHelpOptions
  , renderHelp
  ) where

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Except
import           Data.Foldable
import           Data.List.NonEmpty     (NonEmpty (..), nonEmpty)
import qualified Data.List.NonEmpty     as NonEmpty
import           Data.Map.Strict        (Map)
import qualified Data.Map.Strict        as Map
import           Data.Maybe
import           Data.String
import           Data.Text              (Text)
import qualified Data.Text              as T
import qualified Data.Text.Lazy.Builder as TLB
import           Data.Version
import           Data.Void
import           GHC.Generics

import           Mangrove.Parser
import           Mangrove.ParseTree
import           Mangrove.Render
import           Mangrove.Resolve
import           Mangrove.Scheme
import           Mangrove.Scheme.Common
import           Mangrove.Scheme.Sub    (SubScheme)
import qualified Mangrove.Scheme.Sub    as Sub
import           Mangrove.Stream
import           Mangrove.TextParser
import           Mangrove.Token
import           Mangrove.Valency

--------------------------------------------------------------------------------
-- User Interface Descriptions

-- | A flag is a special argument that identifies a named option to
-- the parser. Flags can have two forms: long flags start with a
-- double dash (e.g. "--example") followed by a string while short
-- flags start with only a single dash (e.g. "-e") and are identified
-- by a single character.
--
-- For convenience, 'Flag' is an instance of 'Data.String.IsString'.
-- Thus, you can write @"--flop"@ instead of @LongFlag "flop"@ and
-- @"-c"@ instead of @ShortFlag \'c\'@.
data Flag
  = ShortFlag !Char
  | LongFlag !Text
  deriving (Eq, Generic, Ord, Show)

instance IsString Flag where
  fromString ('-':'-':name)
    | not (null name) = LongFlag $ T.pack name
  fromString ['-', c]
    | c /= '-' = ShortFlag c
  fromString s = error $ "not a valid flag: " <> s

instance Render Flag where
  render (LongFlag s)  = "--" <> render s
  render (ShortFlag c) = "-" <> render c

-- | A description of a CLI option.
data OptionInfo = OptionInfo
  { optFlags :: !(NonEmpty Flag) -- ^ A list of flags that trigger this option.
  , optHelp  :: !Text -- ^ A description displayed in help output.
  } deriving (Eq, Ord, Show)

-- | Get a representative flag for this option (e.g. the first one).
optHead :: OptionInfo -> Flag
optHead = NonEmpty.head . optFlags

-- | A description of a CLI command.
data CommandInfo = CommandInfo
  { cmdNames :: !(NonEmpty Text) -- ^ Command Names
  , cmdHelp  :: !Text -- ^ A description displayed in help output.
  } deriving (Eq, Ord, Show)

-- | Get a representative command name for this command (e.g. the
-- first one).
cmdHead :: CommandInfo -> Text
cmdHead = NonEmpty.head . cmdNames

-- | Requests supported by UNIX-style parsers.
data UnixRequest
  = VersionRequest -- ^ A request for version information
  | HelpRequest [Text] -- ^ A request for help and usage information
  deriving (Eq, Generic, Show)

-- | Context-aware requests
type UnixRequest' = [Token UnixScheme] -> UnixRequest

-- | A request for help information within the current command context
helpRequest :: UnixRequest'
helpRequest context = HelpRequest [cmd | UnixCommand cmd <- context]

-- | A request for version information
versionRequest :: UnixRequest'
versionRequest _ = VersionRequest

-- | Generate a response to a help request.
respondHelpRequest
  :: [Text]
  -> ParseTree UnixScheme r
  -> ProgramInfo
  -> Text
respondHelpRequest cmds tree info = renderText
  $ "Usage(s):\n"
  <> formatUsages (programName info) usages <> "\n"
  <> render (programDesc info) <> "\n"
  <> renderHelp tree cmds
  where
    usages = decomposeTree tree cmds

-- | Generate a response to a version request.
respondVersionRequest
  :: ProgramInfo
  -> Text
respondVersionRequest info = renderText
  $ render (programName info)
  <> " version "
  <> renderVersion (programVersion info)
  <> "\n"
  where
    renderVersion = TLB.fromString . showVersion

-- | A parsing scheme for Unix-style command line syntax.
data UnixScheme r
  -- | A freeform positional parameter
  = Parameter (TextParser r)
  -- | A subcommand with its own parse tree
  | Command !CommandInfo (ParseTree UnixScheme r)
  -- | A named option that might support suboptions
  | Option !OptionInfo (ParseTree SubScheme r)
  -- | A special option that raises a request for information
  | RequestOption !OptionInfo !UnixRequest'
  deriving (Functor)

instance Show (UnixScheme r) where
  showsPrec p (Parameter tp) =
    showParen (p >= 10)
    $ showString "Parameter "
    . showsTextParser tp
  showsPrec p (Option info subtree) =
    showParen (p >= 10)
    $ showString "Option "
    . showsPrec 11 info
    . showString " "
    . showsPrec 11 subtree
  showsPrec p (Command info subtree) =
    showParen (p >= 10)
    $ showString "Command "
    . showsPrec 11 info
    . showString " "
    . showsPrec 11 subtree
  showsPrec p (RequestOption info _) =
    showParen (p >= 10)
    $ showString "RequestOption "
    . showsPrec 11 info
    . showString " _"

instance Valency UnixScheme where
  valency (Parameter _)       = Just 1
  valency (Command _ subtree) = fmap (+1) (valency subtree)
  valency (Option _ subtree)  = fmap (max 2) (valency subtree)
  valency (RequestOption {})  = Just 1

instance Resolve UnixScheme where
  resolve (Parameter (TextParser hint _)) =
    ExpectedError [render hint]
  resolve (Option info _) =
    ExpectedError [render $ optHead info]
  resolve (RequestOption info _) =
    ExpectedError [render $ optHead info]
  resolve (Command info _) =
    ExpectedError [render $ cmdHead info]

-- | A parser for interpreting options. An option always begins with a
-- flag, followed optionally by an "=" sign and a bound argument. The
-- strings "--" and "-" are not treated as options.
parseUnixOption :: Alternative f => Text -> f (Flag, Maybe Text)
parseUnixOption (T.stripPrefix "--" -> Just s)
  | not (T.null s) =
    case keyEqualsValue s of
      Just (k, v) -> pure (LongFlag k, Just v)
      Nothing     -> pure (LongFlag s, Nothing)
parseUnixOption (T.stripPrefix "-" >=> T.uncons -> Just (k,v))
  | k /= '-' =
    pure (ShortFlag k, munless (T.null v) $ Just v)
parseUnixOption _ = empty

-- | Does this text look like a flag? We check whether it starts with
-- "-" followed by any other character.
isMarked :: Text -> Bool
isMarked "-" = False
isMarked s   = "-" `T.isPrefixOf` s

instance HasTokens UnixScheme where
  data Token UnixScheme
    -- | A freeform positional argument that is not an option or command
    = UnixArgument Text
    -- | A recognized subcommand
    | UnixCommand Text
    -- | A named option with optional bound argument
    | UnixOption Flag (Maybe Text)
    deriving (Eq, Generic, Show)

  delimiter _ = ' '

instance Scheme UnixScheme where
  type Request UnixScheme = UnixRequest

  respond (HelpRequest cmds) tree info = respondHelpRequest cmds tree info
  respond VersionRequest _ info        = respondVersionRequest info

  parseSpecials = do
    peekMaybe >>= \case
      Just "--" -> pop_ *> setEscaped True
      _         -> pure ()

  activate (Parameter tp) = do
    next <- peek

    -- Arguments that begin with a dash should never be treated as
    -- unbound subarguments. However, the string "-" is always
    -- accepted since this is commonly used to represent stdin.
    escaped <- getEscaped
    guard $ escaped || not (isMarked next)

    withContext (UnixArgument next) $
      pop_ *> runTextParser tp next

  activate (Option info subtree) = do
    -- Arguments should never be interpreted as options when escaped.
    getEscaped >>= guard . not

    (flag, mbound) <- peek >>= parseUnixOption
    guard $ flag `elem` optFlags info
    pop_

    -- We need to convert whatever argument string we have (if any)
    -- into a list of subarguments as input for the subparser. If the
    -- subtree accepts multiple arguments, we split the input by
    -- comma. Otherwise, we can just pass a singleton list containing
    -- the argument string.
    --
    -- If the subtree contains no suboptions, we enable escaping to
    -- prevent arguments containing an "=" sign from being interpreted
    -- as suboptions. This is necessary because individual
    -- subparameter parsers have no way to determine that such an
    -- argument won't be consumed by a subsequent suboption parser.
    -- Escaping forces subparameter parsers to consume the argument,
    -- regardless of its form.
    let splitArgs s = if multary subtree
                      then T.split (== ',') s
                      else [s]
        initState args = StreamState
          { streamContent = args
          , streamContext = []
          , streamEscaped = not $ Sub.hasSubOptions subtree
          }
        parseSubargs args =
          case runArgumentParser' subtree (initState args) of
            Success leftover result -> pure (leftover, result)
            Failure err             -> throwError $ render err

    withContext (UnixOption flag mbound) $ do
      -- If a bound argument (e.g. --floop=blah) is provided, we
      -- expect it to be consumed by the subparser. If it isn't fully
      -- consumed, we have nothing to do with the leftovers, so we
      -- throw an error.
      --
      -- If there's no bound argument but the next regular argument
      -- doesn't look like an option, then we try running the
      -- subparser using that as input. If it is fully consumed, we
      -- pop it from the front of the stream. If nothing is consumed,
      -- we leave it at the head of the stream. However, if it is
      -- partially consumed, then something has gone wrong, and we
      -- throw an error.
      mnext <- peekMaybe
      case (mbound, mnext) of
        (Just argString, _) -> do
          (leftover, result) <- parseSubargs (splitArgs argString)
          forM_ leftover $ \arg ->
            throwError $ "unrecognized subargument: " <> render arg
          pure result
        (_, Just argString)
          | not (isMarked argString) -> do
              let args = splitArgs argString
              (leftover, result) <- parseSubargs args
              when (length args /= length leftover) $ do
                forM_ leftover $ \arg ->
                  throwError $ "unrecognized subargument: " <> render arg
                pop_
              pure result
        _ -> do
          (_, result) <- parseSubargs []
          pure result

  activate (RequestOption info mkRequest) = do
    -- Arguments should never be interpreted as options when escaped.
    getEscaped >>= guard . not

    (flag, mbound) <- peek >>= parseUnixOption
    guard $ flag `elem` optFlags info
    pop_

    withContext (UnixOption flag mbound) $
      getContext >>= request . mkRequest

  activate (Command info subtree) = do
    -- Arguments should never be interpreted as commands when escaped.
    getEscaped >>= guard . not

    next <- peek
    guard $ next `elem` cmdNames info
      && not ("-" `T.isPrefixOf` next) -- not sure if this check is necessary?
    pop_

    withContext (UnixCommand next) $ do
      satiate subtree
      >>= resolveLifted

prec_bind :: Int
prec_bind = 12

instance Usage UnixScheme where
  usagePrec _ (Parameter tp) =
    usage tp
  usagePrec _ (RequestOption info _) =
    render $ optHead info
  usagePrec d (Option info subtree)
    | nullary subtree = render flag
    | otherwise =
        bracesIf (d > prec_bind)
        $ render flag
        <> separator
        <> usagePrec prec_prod subtree
    where
      flag = optHead info
      separator = case flag of
        LongFlag _ -> "="
        _          -> ""
  usagePrec d (Command info subtree) =
    bracesIf (d > prec_bind)
    $ render (cmdHead info) <> params
    where
      params =
        munless (nullary subtree)
        $ " " <> usagePrec prec_prod subtree

instance Render (UnixScheme r) where
  render = usage

instance Render (Token UnixScheme) where
  render (UnixArgument s)                      = render s
  render (UnixCommand s)                       = render s
  render (UnixOption f Nothing)                = render f
  render (UnixOption f@(LongFlag _) (Just v))  = render f <> "=" <> render v
  render (UnixOption f@(ShortFlag _) (Just v)) = render f <> render v

-- | Convenient type alias for Unix-flavored parse trees.
type UnixParser = ParseTree UnixScheme

--------------------------------------------------------------------------------
-- Decomposing Trees

data Branch r
  = ReqBranch (ParseTree UnixScheme Void) -- ^ Request branches
  | CmdBranch (ParseTree UnixScheme r)    -- ^ Command branches
  deriving (Functor, Show)

unwrapBranch :: Branch r -> ParseTree UnixScheme r
unwrapBranch (ReqBranch b) = vacuous b
unwrapBranch (CmdBranch b) = b

-- | Alter the shape of a 'Branch'.
mapBranch
  :: (forall a. ParseTree UnixScheme a -> ParseTree UnixScheme a)
  -> Branch r
  -> Branch r
mapBranch f (ReqBranch b) = ReqBranch $ f b
mapBranch f (CmdBranch b) = CmdBranch $ f b

-- | A factored group of subtrees representing different usage modes.
data Usages a = Usages
  (Maybe (ParseTree UnixScheme a)) -- ^ Trunk
  [Branch a]                       -- ^ Branches

usagesToList :: Usages r -> [ParseTree UnixScheme r]
usagesToList (Usages trunk branches) =
  maybeToList trunk <> fmap unwrapBranch branches

-- | Factor a 'ParseTree' into several independant subtrees (a trunk
-- and branches), potentially filtered to specific commands.
--
-- Each branch can be thought of as corresponding to one particular
-- mode of operation, in that it contains at least one command or
-- option that conflicts with commands or options in other branches.
--
-- We can select only branches that correspond to a particular
-- subcommand by passing the components of that subcommand as a list:
--
-- > decomposeTree tree [] -- No filtering
-- > decomposeTree tree ["stash", "list"] -- Select "stash list" command
decomposeTree :: ParseTree UnixScheme r -> [Text] -> Usages r
decomposeTree (ParseNode (RequestOption info mkRequest)) commands =
  -- If we're currently searching for a specific command, then
  -- this request option is irrelevant.
  let node = ParseNode (RequestOption info mkRequest)
  in Usages Nothing [ReqBranch node | null commands]

decomposeTree (ParseNode (Command info subtree)) commands
  | commandMismatch =
    -- We are looking for a specific command and it's not this
    -- one, so don't return any trees.
    Usages Nothing []
  | otherwise =
    -- Either this is the command we're looking for, or we're not
    -- looking for a command.
    let Usages trunk branches = decomposeTree subtree (drop 1 commands)
    in Usages Nothing
       $ mapBranch (ParseNode . Command info)
       <$> maybeToList (fmap CmdBranch trunk) <> branches
  where
    commandMismatch =
      case commands of
        (command : _) -> command `notElem` cmdNames info
        []            -> False

decomposeTree (SumNode l r) commands =
  let Usages trunkL branchLs = decomposeTree l commands
      Usages trunkR branchRs = decomposeTree r commands

      -- If both subtrees have trunks, we want to sum them normally.
      -- However, if only one subtree yields a trunk, we can replace
      -- the sum with it.
      trunk = liftA2 SumNode trunkL trunkR
             <|> trunkL
             <|> trunkR
  in Usages trunk (branchLs <> branchRs)

decomposeTree (ProdNode f l r) commands =
  let Usages trunkL branchLs = decomposeTree l commands
      Usages trunkR branchRs = decomposeTree r commands

      -- Requests prevent any further parsing, so if one of the
      -- subtrees yields request branches, the other subtree is
      -- irrelevant. If somehow both subtrees yield request
      -- branches, then a product node behaves effectively like a
      -- sum node because we could never actually trigger both
      -- requests.
      graft with _trunk _branches = _branches >>= \case
        ReqBranch b -> pure $ ReqBranch b
        CmdBranch b -> CmdBranch . with b <$> maybeToList _trunk

      prod = ProdNode f
      trunk = liftA2 prod trunkL trunkR
      branches = graft (flip prod) trunkL branchRs <>
                 graft prod        trunkR branchLs
  in Usages trunk branches

decomposeTree tree _ = Usages (Just tree) []

formatUsages :: Text -> Usages r -> Builder
formatUsages progName =
  foldMap (\tree -> "  " <> render progName <> " " <> usage tree <> "\n")
  . usagesToList

--------------------------------------------------------------------------------
-- Help

-- | Automatically insert a help option at the top level of the tree
-- and every subcommand tree.
addHelpOptions
  :: NonEmpty Flag
  -> Text
  -> ParseTree UnixScheme r
  -> ParseTree UnixScheme r
addHelpOptions flags desc tree = ParseNode helpOption <|> go tree
  where
    helpOption :: UnixScheme a
    helpOption = RequestOption (OptionInfo flags desc) helpRequest

    go :: ParseTree UnixScheme a -> ParseTree UnixScheme a
    go (ParseNode (Command info subtree)) =
      ParseNode
      $ Command info
      $ ParseNode helpOption <|> go subtree
    go (ProdNode f l r) = ProdNode f (go l) (go r)
    go (SumNode l r) = SumNode (go l) (go r)
    go (ManyNode require p) = ManyNode require (go p)
    go node = node

-- | Pretty-print a list of flags.
displayFlagList :: NonEmpty Flag -> Builder
displayFlagList =
  fold
  . NonEmpty.intersperse ", "
  . fmap render
  . NonEmpty.sort

-- | Pretty-print a description of an option with the given
-- t'OptionInfo' and subtree.
displayOption :: OptionInfo -> ParseTree SubScheme r -> Builder
displayOption info subtree =
  let flagList = displayFlagList (optFlags info)
      paramInfo = if nullary subtree
                  then mempty
                  else " " <> render subtree
  in "  " <> flagList <> paramInfo <> "\n" <>
     "         " <> render (optHelp info) <> "\n"

-- | Enumerate descriptive information for all options available in a
-- parse tree, indexed by the set of commands under which they exist.
collectOptions :: ParseTree UnixScheme r -> Map [CommandInfo] [Builder]
collectOptions tree = go tree (Map.singleton [] [])
  where
    go :: ParseTree UnixScheme r
       -> Map [CommandInfo] [Builder]
       -> Map [CommandInfo] [Builder]
    go (ParseNode (Option info subtree)) =
      Map.adjust (displayOption info subtree :) []
    go (ParseNode (RequestOption info _)) =
      Map.adjust (displayOption info empty :) []
    go (ParseNode (Command info subtree)) =
      Map.union $ Map.mapKeys (info :) $ collectOptions subtree
    go (ProdNode _ l r) = go r . go l
    go (SumNode l r)    = go r . go l
    go (ManyNode _ p)   = go p
    go _                = id

formatCommandList :: NonEmpty CommandInfo -> Text
formatCommandList cmdList =
  T.intercalate ", "
  $ T.unwords
  <$> qualifiedCmds
  where
    currentNames = NonEmpty.toList $ cmdNames $ NonEmpty.last cmdList
    parentsNames = fmap cmdHead $ NonEmpty.init cmdList
    qualifiedCmds = [parentsNames <> [name] | name <- currentNames]

renderCmdHeader :: NonEmpty CommandInfo -> Builder
renderCmdHeader infos =
  render (formatCommandList infos) <> "\n"
  <> "  " <> (render . cmdHelp . NonEmpty.last) infos <> "\n\n"

-- | Format an index of commands and options for help output display.
renderTables :: Map [CommandInfo] [Builder] -> Builder
renderTables m =
  case Map.findWithDefault [] [] m of
    [] -> mempty
    gs -> "\nGLOBAL OPTIONS\n\n"
          <> mconcat gs
  <>
  munless (all null (Map.keys m)) "\nSUBCOMMANDS\n"
  <>
  Map.foldrWithKey
  (\cmds bs acc ->
     case nonEmpty cmds of
       Nothing -> acc
       Just cmds' ->
         "\n" <> renderCmdHeader cmds' <> mconcat bs <> acc
  ) mempty m

-- | Select only the options tables which exist under a particular
-- command sequence.
selectSubtable
  :: [Text]
  -> Map [CommandInfo] [a]
  -> Map [CommandInfo] [a]
selectSubtable cmds =
  Map.filterWithKey (\infos _ -> isParentCommand cmds infos)

isParentCommand :: [Text] -> [CommandInfo] -> Bool
isParentCommand cmds =
  and . zipWith (\cmd info -> cmd `elem` cmdNames info) cmds

-- | Render formatted help information for all commands and options
-- that exist underneath the current command context.
renderHelp
  :: ParseTree UnixScheme r
  -> [Text] -- ^ Command Context
  -> Builder
renderHelp tree cmds =
  renderTables
  $ selectSubtable cmds
  $ collectOptions tree