CLI.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoFieldSelectors #-}

module Skeletest.Internal.CLI (
  Flag (..),
  flag,
  IsFlag (..),
  FlagSpec (..),
  getFlag,
  loadCliArgs,

  -- * Internal
  parseCliArgs,
  CLIParseResult (..),
  CLIFlagStore,
) where

import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class qualified as Trans
import Control.Monad.Trans.Except qualified as Trans
import Control.Monad.Trans.State qualified as Trans
import Data.Bifunctor (first, second)
import Data.Dynamic (Dynamic, fromDynamic, toDyn)
import Data.Foldable1 qualified as Foldable1
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Proxy (Proxy (..))
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.IO qualified as Text
import Data.Typeable (TypeRep, Typeable, typeOf, typeRep)
import Skeletest.Internal.Error (SkeletestError (..), invariantViolation)
import Skeletest.Internal.TestTargets (TestTargets, parseTestTargets)
import System.Environment (getArgs)
import System.Exit (exitFailure, exitSuccess)
import System.IO (stderr)
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO.Exception (throwIO)

-- | Register a CLI flag.
--
-- Usage:
--
-- @
-- {- MyFixture.hs -}
-- import Skeletest
--
-- newtype MyFlag = MyFlag String
-- instance IsFlag MyFlag where
--   flagName = "my-flag"
--   flagHelp = "The value for MyFixture"
--   flagSpec =
--     OptionalFlag
--       { flagDefault = "foo"
--       , flagParse = \case
--           "illegal" -> Left "invalid flag value"
--           s -> Right (MyFlag s)
--       }
--
-- instance Fixture MyFixture where
--   fixtureAction = do
--     MyFlag val <- getFlag
--     ...
--
-- {- Main.hs -}
-- import MyFixture
--
-- cliFlags =
--   [ flag @MyFlag
--   ]
-- @
data Flag = forall a. (IsFlag a) => Flag (Proxy a)

flag :: forall a. (IsFlag a) => Flag
flag = Flag (Proxy @a)

class (Typeable a) => IsFlag a where
  flagName :: String

  flagShort :: Maybe Char
  flagShort = Nothing

  -- | The placeholder for the flag to show in the help text, if
  -- the flag takes an argument.
  flagMetaVar :: String
  flagMetaVar = "VAR"

  flagHelp :: String

  flagSpec :: FlagSpec a

data FlagSpec a
  = SwitchFlag
      { flagFromBool :: Bool -> a
      }
  | RequiredFlag
      { flagParse :: String -> Either String a
      }
  | OptionalFlag
      { flagDefault :: a
      , flagParse :: String -> Either String a
      }

getFlag :: forall a m. (MonadIO m, IsFlag a) => m a
getFlag =
  liftIO $
    lookupCliFlag rep >>= \case
      Just dyn ->
        case fromDynamic dyn of
          Just a -> pure a
          Nothing ->
            invariantViolation . unwords $
              [ "CLI flag store contained incorrect types."
              , "Expected: " <> show rep <> "."
              , "Got: " <> show dyn
              ]
      Nothing -> throwIO $ CliFlagNotFound (Text.pack $ flagName @a)
 where
  rep = typeRep (Proxy @a)

{----- Load CLI arguments -----}

-- | Parse the CLI arguments using the given user-defined flags, then
-- stores the flags in the global state and returns the positional
-- arguments.
loadCliArgs :: [Flag] -> [Flag] -> IO TestTargets
loadCliArgs builtinFlags flags = do
  args0 <- getArgs
  case parseCliArgs (builtinFlags <> flags) args0 of
    CLISetupFailure msg -> do
      Text.hPutStrLn stderr $ "ERROR: " <> msg
      exitFailure
    CLIHelpRequested -> do
      Text.putStrLn helpText
      exitSuccess
    CLIParseFailure msg -> do
      Text.hPutStrLn stderr $ msg <> "\n\n" <> helpText
      exitFailure
    CLIParseSuccess{testTargets, flagStore} -> do
      setCliFlagStore flagStore
      pure testTargets
 where
  helpText = getHelpText builtinFlags flags

getHelpText :: [Flag] -> [Flag] -> Text
getHelpText builtinFlags customFlags =
  Text.intercalate "\n\n" $
    "Usage: skeletest [OPTIONS] [--] [TARGETS]" : map (uncurry renderSection) helpSections
 where
  helpSections =
    filter (not . Text.null . snd) $
      [ ("TEST SELECTION", testSelectionDocs)
      , ("BUILTIN OPTIONS", renderFlagList builtinFlagDocs)
      , ("CUSTOM OPTIONS", renderFlagList customFlagDocs)
      ]

  testSelectionDocs =
    Text.intercalate "\n" $
      [ "Test targets may be specified as plain positional arguments, with the following syntax:"
      , "    * Tests including substring:      '[myFooFunc]'"
      , "    * Tests tagged with marker:       '@fast'"
      , "    * Tests in file, relative to CWD: 'test/MyLib/FooSpec.hs'"
      , "    * Tests matching pattern in file: 'test/MyLib/FooSpec.hs[myFooFunc]'"
      , "        * Syntax sugar for '(test/MyLib/FooSpec.hs and [myFooFunc])'"
      , "    * Tests matching both targets:    '[func1] and [func2]'"
      , "    * Tests matching either target:   '[func1] or [func2]'"
      , "    * Tests not matching target:      'not [func1]'"
      , ""
      , "More examples:"
      , "    * 'test/MySpec.hs and ([myFooFunc] or [myBarFunc]) and @fast'"
      , "    * '[myFooFunc] or test/MySpec.hs[myBarFunc]'"
      , ""
      , "When multiple targets are specified, they are joined with 'or'."
      ]

  builtinFlagDocs = ("help", Just 'h', Nothing, "Display this help text") : fromFlags builtinFlags
  customFlagDocs = fromFlags customFlags
  fromFlags flags =
    [ (Text.pack (flagName @a), flagShort @a, mMetaVar, Text.pack (flagHelp @a))
    | Flag (Proxy :: Proxy a) <- flags
    , let mMetaVar =
            case flagSpec @a of
              SwitchFlag{} -> Nothing
              RequiredFlag{} -> Just $ Text.pack (flagMetaVar @a)
              OptionalFlag{} -> Just $ Text.pack (flagMetaVar @a)
    ]

  renderSection title body =
    Text.intercalate "\n" $
      [ "===== " <> title
      , ""
      , body
      ]

  renderFlagList flagList =
    Text.intercalate "\n" . mkTabular $
      [ (shortName <> renderLongFlag longName <> metaVar, help)
      | (longName, mShortName, mMetaVar, help) <- flagList
      , let
          shortName =
            case mShortName of
              Just short -> renderShortFlag short <> ", "
              Nothing -> ""
          metaVar =
            case mMetaVar of
              Just meta -> " <" <> meta <> ">"
              Nothing -> ""
      ]

  mkTabular rows0 =
    case NonEmpty.nonEmpty rows0 of
      Nothing -> []
      Just rows ->
        let fstColWidth = Foldable1.maximum $ NonEmpty.map (Text.length . fst) rows
            margin = 2 -- space between columns
         in [ a <> Text.replicate (fstColWidth - Text.length a + margin) " " <> b
            | (a, b) <- NonEmpty.toList rows
            ]

{----- Parse args -----}

data CLIParseResult
  = CLISetupFailure Text
  | CLIHelpRequested
  | CLIParseFailure Text
  | CLIParseSuccess
      { testTargets :: TestTargets
      , flagStore :: CLIFlagStore
      }

parseCliArgs :: [Flag] -> [String] -> CLIParseResult
parseCliArgs flags args = either id id $ do
  longFlags <- extractLongFlags
  shortFlags <- extractShortFlags

  -- quick sweep for --help/-h after flag validation; skip parsing flags if so
  when (any (`elem` ["--help", "-h"]) args) $ Left CLIHelpRequested

  (args', flagStore) <- first CLIParseFailure $ parseCliArgsWith longFlags shortFlags args
  testTargets <- first CLIParseFailure $ parseTestTargets args'
  flagStore' <- first CLIParseFailure $ resolveFlags flags flagStore
  pure CLIParseSuccess{testTargets, flagStore = flagStore'}
 where
  extractLongFlags =
    toFlagMap renderLongFlag $
      [ (Text.pack $ flagName @a, f)
      | f@(Flag (Proxy :: Proxy a)) <- flags
      ]

  extractShortFlags =
    toFlagMap renderShortFlag $
      [ (shortFlag, f)
      | f@(Flag (Proxy :: Proxy a)) <- flags
      , Just shortFlag <- pure $ flagShort @a
      ]

  toFlagMap :: (Ord name) => (name -> Text) -> [(name, a)] -> Either CLIParseResult (Map name a)
  toFlagMap renderFlag vals =
    let go seen = \case
          [] -> Right $ Map.fromList vals
          (name, _) : xs
            | name `Set.member` seen -> Left . CLISetupFailure $ "Flag registered multiple times: " <> renderFlag name
            | otherwise -> go (Set.insert name seen) xs
     in go Set.empty vals

type ArgParserM = Trans.StateT ([Text], CLIFlagStore) (Trans.Except Text)

parseCliArgsWith :: Map Text Flag -> Map Char Flag -> [String] -> Either Text ([Text], CLIFlagStore)
parseCliArgsWith longFlags shortFlags = Trans.runExcept . flip Trans.execStateT ([], Map.empty) . parseArgs
 where
  parseArgs = \case
    [] -> pure ()
    "--" : rest -> addArgs rest
    curr : rest
      | Just longFlag <- Text.stripPrefix "--" (Text.pack curr) -> parseLongFlag longFlag rest
      | Just chars <- Text.stripPrefix "-" (Text.pack curr) ->
          case Text.unpack chars of
            [] -> argError "Invalid flag: -"
            [shortFlag] -> parseShortFlag shortFlag rest
            _ -> argError $ "Invalid flag: -" <> chars
      | otherwise -> addArgs [curr] >> parseArgs rest

  parseLongFlag name args =
    let (name', args') =
          case Text.breakOn "=" name of
            (_, "") -> (name, args)
            (n, post) -> (n, (drop 1 . Text.unpack) post : args)
     in parseFlag renderLongFlag longFlags name' args'
  parseShortFlag = parseFlag renderShortFlag shortFlags

  parseFlag :: (Ord name) => (name -> Text) -> Map name Flag -> name -> [String] -> ArgParserM ()
  parseFlag renderFlag flagMap name args = do
    Flag (Proxy :: Proxy a) <-
      case Map.lookup name flagMap of
        Nothing -> argError $ "Unknown flag: " <> renderFlag name
        Just f -> pure f
    let parseFlagArg parseArg =
          case args of
            [] -> argError $ "Flag requires argument: " <> renderFlag name
            curr : rest -> parseArg curr >>= addFlagStore >> parseArgs rest
    case flagSpec @a of
      SwitchFlag{flagFromBool} -> addFlagStore (flagFromBool True) >> parseArgs args
      RequiredFlag{flagParse} -> parseFlagArg (Trans.lift . Trans.except . first Text.pack . flagParse)
      OptionalFlag{flagParse} -> parseFlagArg (Trans.lift . Trans.except . first Text.pack . flagParse)

  argError = Trans.lift . Trans.throwE

  addArgs :: [String] -> ArgParserM ()
  addArgs args = Trans.modify (first (<> map Text.pack args))

  addFlagStore :: (Typeable a) => a -> ArgParserM ()
  addFlagStore x = Trans.modify (second (insertFlagStore x))

resolveFlags :: [Flag] -> CLIFlagStore -> Either Text CLIFlagStore
resolveFlags = flip (foldlM go)
 where
  go flagStore (Flag (Proxy :: Proxy a)) = do
    let rep = typeRep (Proxy @a)
    case flagSpec @a of
      SwitchFlag{flagFromBool} ->
        pure $
          if rep `Map.member` flagStore
            then flagStore
            else insertFlagStore (flagFromBool False) flagStore
      RequiredFlag{} ->
        if rep `Map.member` flagStore
          then pure flagStore
          else Left $ "Required flag not set: " <> renderLongFlag (Text.pack $ flagName @a)
      OptionalFlag{flagDefault} ->
        pure $
          if rep `Map.member` flagStore
            then flagStore
            else insertFlagStore flagDefault flagStore

  foldlM f z = \case
    [] -> pure z
    x : xs -> do
      z' <- f z x
      foldlM f z' xs

renderLongFlag :: Text -> Text
renderLongFlag = ("--" <>)

renderShortFlag :: Char -> Text
renderShortFlag c = Text.pack ['-', c]

{----- CLIFlagStore -----}

type CLIFlagStore = Map TypeRep Dynamic

insertFlagStore :: (Typeable a) => a -> CLIFlagStore -> CLIFlagStore
insertFlagStore x = Map.insert (typeOf x) (toDyn x)

cliFlagStoreRef :: IORef CLIFlagStore
cliFlagStoreRef = unsafePerformIO $ newIORef Map.empty
{-# NOINLINE cliFlagStoreRef #-}

setCliFlagStore :: CLIFlagStore -> IO ()
setCliFlagStore = writeIORef cliFlagStoreRef

lookupCliFlag :: TypeRep -> IO (Maybe Dynamic)
lookupCliFlag rep = Map.lookup rep <$> readIORef cliFlagStoreRef