content: {-# 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 System.Environment (getArgs)
  import System.Exit (exitFailure, exitSuccess)
  import System.IO (stderr)
  import System.IO.Unsafe (unsafePerformIO)
  import UnliftIO.Exception (throwIO)
  
  import Skeletest.Internal.Error (SkeletestError (..), invariantViolation)
  import Skeletest.Internal.TestTargets (TestTargets, parseTestTargets)
  
  -- | 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
  