content: -- Copyright (c) Facebook, Inc. and its affiliates.
  --
  -- This source code is licensed under the MIT license found in the
  -- LICENSE file in the root directory of this source tree.
  --
  {-# LANGUAGE ApplicativeDo #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE NumericUnderscores #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TupleSections #-}
  module Retrie.Options
    ( -- * Options
      Options
    , Options_(..)
    , ExecutionMode(..)
    , defaultOptions
    , parseOptions
      -- * Internal
    , buildGrepChain
    , forFn
    , getOptionsParser
    , getTargetFiles
    , parseRewritesInternal
    , parseVerbosity
    , ProtoOptions
    , resolveOptions
    , GrepCommands(..)
    ) where
  
  import Control.Concurrent.Async (mapConcurrently)
  import Control.Monad (when, foldM)
  import Data.Bool
  import Data.Char (isAlphaNum, isSpace)
  import Data.Default as D
  import Data.Foldable (toList)
  import Data.Functor.Identity
  import Data.List
  import Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import Data.Traversable
  import Options.Applicative
  import System.Directory
  import System.FilePath
  import System.Process
  import System.Random.Shuffle
  
  import Retrie.CPP
  import Retrie.Debug
  import Retrie.Elaborate
  import Retrie.ExactPrint
  import Retrie.Fixity
  import Retrie.GroundTerms
  import Retrie.GHC
  import Retrie.Pretty
  import Retrie.Rewrites
  import Retrie.Types
  import Retrie.Universe
  import Retrie.Util
  
  -- | Command-line options.
  type Options = Options_ [Rewrite Universe] AnnotatedImports
  
  -- | Parse options using the given 'FixityEnv'.
  parseOptions :: LibDir -> FixityEnv -> IO Options
  parseOptions libdir fixityEnv = do
    p <- getOptionsParser fixityEnv
    opts <- execParser (info (p <**> helper) fullDesc)
    resolveOptions libdir opts
  
  -- | Create 'Rewrite's from string specifications of rewrites.
  -- We expose this from "Retrie" with a nicer type signature as
  -- 'Retrie.Options.parseRewrites'. We have it here so we can use it with
  -- 'ProtoOptions'.
  parseRewritesInternal :: LibDir -> Options_ a b -> [RewriteSpec] -> IO [Rewrite Universe]
  parseRewritesInternal libdir Options{..} = parseRewriteSpecs libdir parser fixityEnv
    where
      parser fp = parseCPPFile (parseContent libdir fixityEnv) (targetDir </> fp)
  
  -- | Controls the ultimate action taken by 'apply'. The default action is
  -- 'ExecRewrite'.
  data ExecutionMode
    = ExecDryRun -- ^ Pretend to do rewrites, show diff.
    | ExecRewrite -- ^ Perform rewrites.
    | ExecExtract -- ^ Print the resulting expression for each match.
    | ExecSearch -- ^ Print the matched expressions.
    deriving (Show)
  
  data Options_ rewrites imports = Options
    { additionalImports :: imports
      -- ^ Imports specified by the command-line flag '--import'.
    , colorise :: ColoriseFun
      -- ^ Function used to colorize results of certain execution modes.
    , elaborations :: rewrites
      -- ^ Rewrites which are applied to the left-hand side of the actual rewrites.
    , executionMode :: ExecutionMode
      -- ^ Controls behavior of 'apply'. See 'ExecutionMode'.
    , extraIgnores :: [FilePath]
      -- ^ Specific files that should be ignored. Paths should be relative to
      -- 'targetDir'.
    , fixityEnv :: FixityEnv
      -- ^ Fixity information for operators used during parsing (of rewrites and
      -- target modules). Defaults to base fixities.
    , iterateN :: Int
      -- ^ Iterate the given rewrites or 'Retrie' computation up to this many
      -- times. Iteration may stop before the limit if no changes are made during
      -- a given iteration.
    , noDefaultElaborations :: Bool
      -- ^ Do not apply any of the built in elaborations in 'defaultElaborations'.
    , randomOrder :: Bool
      -- ^ Whether to randomize the order of target modules before rewriting them.
    , rewrites :: rewrites
      -- ^ Rewrites specified by command-line flags such as '--adhoc'.
    , roundtrips :: [RoundTrip]
      -- ^ Paths that should be roundtripped through ghc-exactprint to debug.
      -- Specified by the '--roundtrip' command-line flag.
    , singleThreaded :: Bool
      -- ^ Whether to concurrently rewrite target modules.
      -- Mostly useful for viewing debugging output without interleaving it.
    , targetDir :: FilePath
      -- ^ Directory that contains the code being targeted for rewriting.
    , targetFiles :: [FilePath]
      -- ^ Instead of targeting all Haskell files in 'targetDir', only target
      -- specific files. Paths should be relative to 'targetDir'.
    , verbosity :: Verbosity
      -- ^ How much should be output on 'stdout'.
    }
  
  -- | Construct default options for the given target directory.
  defaultOptions
    :: (Default rewrites, Default imports)
    => FilePath -> Options_ rewrites imports
  defaultOptions fp = Options
    { additionalImports = D.def
    , colorise = noColor
    , elaborations = D.def
    , executionMode = ExecRewrite
    , extraIgnores = []
    , fixityEnv = mempty
    , iterateN = 1
    , noDefaultElaborations = False
    , randomOrder = False
    , rewrites = D.def
    , roundtrips = []
    , singleThreaded = False
    , targetDir = fp
    , targetFiles = []
    , verbosity = Normal
    }
  
  -- | Get the options parser. The returned 'ProtoOptions' should be passed
  -- to 'resolveOptions' to get final 'Options'.
  getOptionsParser :: FixityEnv -> IO (Parser ProtoOptions)
  getOptionsParser fEnv = do
    dOpts <- defaultOptions <$> getCurrentDirectory
    return $ buildParser dOpts { fixityEnv = fEnv }
  
  buildParser :: ProtoOptions -> Parser ProtoOptions
  buildParser dOpts = do
    singleThreaded <- switch $ mconcat
      [ long "single-threaded"
      , showDefault
      , help "Don't try to parallelize things (for debugging)."
      ]
    targetDir <- option str $ mconcat
      [ long "target"
      , short 't'
      , metavar "PATH"
      , action "directory" -- complete with directory
      , value (targetDir dOpts)
      , showDefault
      , help "Path to target with rewrites."
      ]
    targetFiles <- many $ option str $ mconcat
      [ long "target-file"
      , metavar "PATH"
      , action "file" -- complete with filenames
      , help "Target specific file for rewriting."
      ]
    verbosity <- parseVerbosity (verbosity dOpts)
    additionalImports <- many $ option str $ mconcat
      [ long "import"
      , metavar "IMPORT"
      , help
          "Add given import statement to modules that are modified by a rewrite."
      ]
    extraIgnores <- many $ option str $ mconcat
      [ long "ignore"
      , metavar "PATH"
      , action "file" -- complete with filenames
      , help "Ignore specific file while rewriting."
      ]
    colorise <- fmap (bool noColor addColor) $ switch $ mconcat
      [ long "color"
      , help "Highlight matches with color."
      ]
    noDefaultElaborations <- switch $ mconcat
      [ long "no-default-elaborations"
      , showDefault
      , help "Don't apply any of the default elaborations to rewrites."
      ]
    randomOrder <- switch $ mconcat
      [ long "random-order"
      , help "Randomize the order of targeted modules."
      ]
    iterateN <- option auto $ mconcat
      [ long "iterate"
      , short 'i'
      , metavar "N"
      , value 1
      , help "Iterate rewrites up to N times."
      ]
  
    executionMode <- parseMode
    rewrites <- parseRewriteSpecOptions
    elaborations <- parseElaborations
    roundtrips <- parseRoundtrips
    return Options{ fixityEnv = fixityEnv dOpts, ..}
  
  parseElaborations :: Parser [RewriteSpec]
  parseElaborations = concat <$> traverse many
    [ fmap Adhoc $ option str $ mconcat
      [ long "elaborate"
      , metavar "EQUATION"
      , help "Elaborate the left-hand side of rewrites using the given equation."
      ]
    , fmap AdhocType $ option str $ mconcat
      [ long "elaborate-type"
      , metavar "EQUATION"
      , help "Elaborate the left-hand side of rewrites using the given equation."
      ]
    , fmap AdhocPattern $ option str $ mconcat
      [ long "elaborate-pattern"
      , metavar "EQUATION"
      , help "Elaborate the left-hand side of rewrites using the given equation."
      ]
    ]
  
  parseRewriteSpecOptions :: Parser [RewriteSpec]
  parseRewriteSpecOptions = concat <$> traverse many
    [ fmap Unfold $ option str $ mconcat
        [ long "unfold"
        , short 'u'
        , metavar "NAME"
        , help "Unfold given fully-qualified name."
        ]
    , fmap Fold $ option str $ mconcat
        [ long "fold"
        , short 'f'
        , metavar "NAME"
        , help "Fold given fully-qualified name."
        ]
    , fmap RuleForward $ option str $ mconcat
        [ long "rule-forward"
        , short 'l'
        , metavar "NAME"
        , help "Apply fully-qualified RULE name left-to-right."
        ]
    , fmap RuleBackward $ option str $ mconcat
        [ long "rule-backward"
        , short 'r'
        , metavar "NAME"
        , help "Apply fully-qualified RULE name right-to-left."
        ]
    , fmap TypeForward $ option str $ mconcat
        [ long "type-forward"
        , metavar "NAME"
        , help "Apply fully-qualified type synonym name left-to-right."
        ]
    , fmap TypeBackward $ option str $ mconcat
        [ long "type-backward"
        , metavar "NAME"
        , help "Apply fully-qualified type synonym name right-to-left."
        ]
    , fmap Adhoc $ option str $ mconcat
        [ long "adhoc"
        , metavar "EQUATION"
        , help "Apply an adhoc equation of the form: forall vs. lhs = rhs"
        ]
    , fmap AdhocType $ option str $ mconcat
        [ long "adhoc-type"
        , metavar "EQUATION"
        , help "Apply an adhoc type equation of the form: lhs = rhs"
        ]
    , fmap PatternForward $ option str $ mconcat
        [ long "pattern-forward"
        , metavar "NAME"
        , help "Apply fully-qualified pattern synonym name left-to-right."
        ]
    , fmap PatternBackward $ option str $ mconcat
        [ long "pattern-backward"
        , metavar "NAME"
        , help "Apply fully-qualified pattern synonym name right-to-left."
        ]
    , fmap AdhocPattern $ option str $ mconcat
        [ long "adhoc-pattern"
        , metavar "EQUATION"
        , help "Apply an adhoc pattern equation of the form: lhs = rhs"
        ]
    ]
  
  parseMode :: Parser ExecutionMode
  parseMode =
    parseDryRun <|>
    parseExtract <|>
    parseSearch <|>
    pure ExecRewrite
  
  parseDryRun :: Parser ExecutionMode
  parseDryRun = flag' ExecDryRun $ mconcat
    [ long "dry-run"
    , help "Don't overwrite files. Print rewrite results."
    ]
  
  parseExtract :: Parser ExecutionMode
  parseExtract = flag' ExecExtract $ mconcat
    [ long "extract"
    , help "Find the left-hand side, display the instantiated right-hand side."
    ]
  
  parseSearch :: Parser ExecutionMode
  parseSearch = flag' ExecSearch $ mconcat
    [ long "search"
    , help "Search for left-hand side of the rewrite and show matches."
    ]
  
  -- | Parser for 'Verbosity'.
  parseVerbosity :: Verbosity -> Parser Verbosity
  parseVerbosity defaultV = option (eitherReader verbosityReader) $ mconcat
    [ long "verbosity"
    , short 'v'
    , value defaultV
    , showDefault
    , help verbosityHelp
    ]
  
  verbosityReader :: String -> Either String Verbosity
  verbosityReader "0" = Right Silent
  verbosityReader "1" = Right Normal
  verbosityReader "2" = Right Loud
  verbosityReader _ =
    Left $ "invalid verbosity. Valid values: " ++ verbosityHelp
  
  verbosityHelp :: String
  verbosityHelp = "0: silent, 1: normal, 2: loud (implies --single-threaded)"
  
  -------------------------------------------------------------------------------
  
  -- | Options that have been parsed, but not fully resolved.
  type ProtoOptions = Options_ [RewriteSpec] [String]
  
  -- | Resolve 'ProtoOptions' into 'Options'. Parses rewrites into 'Rewrite's,
  -- parses imports, validates options, and extends 'fixityEnv' with any
  -- declared fixities in the target directory.
  resolveOptions :: LibDir -> ProtoOptions -> IO Options
  resolveOptions libdir protoOpts = do
    absoluteTargetDir <- makeAbsolute (targetDir protoOpts)
    opts@Options{..} <-
      addLocalFixities libdir protoOpts { targetDir = absoluteTargetDir }
    parsedImports <- parseImports libdir additionalImports
    debugPrint verbosity "Imports:" $
      runIdentity $ fmap astA $ transformA parsedImports $ \ imps -> do
        -- anns <- getAnnsT
        return $ map exactPrint imps
    rrs <- parseRewritesInternal libdir opts rewrites
    es <- parseRewritesInternal libdir opts $
      (if noDefaultElaborations then [] else defaultElaborations) ++
      elaborations
    elaborated <- elaborateRewritesInternal fixityEnv es rrs
    return Options
      { additionalImports = parsedImports
      , elaborations = es
      , rewrites = elaborated
      , singleThreaded = singleThreaded || verbosity == Loud
      , ..
      }
  
  -- | Find all fixity declarations in targetDir and add them to fixity env.
  addLocalFixities :: LibDir -> Options_ a b -> IO (Options_ a b)
  addLocalFixities libdir opts = do
    -- do not limit search for infix decls to only targetFiles
    let opts' = opts { targetFiles = [] }
    -- "infix" will find infixl and infixr as well
    files <- getTargetFiles opts' [HashSet.singleton "infix"]
  
    fixFns <- forFn opts files $ \ fp -> do
      ms <- toList <$> parseCPPFile (parseContentNoFixity libdir) fp
      return $ extendFixityEnv
        [ (rdrFS nm, fixity)
        | m <- ms
        , (L _ nm, fixity) <- fixityDecls (unLoc (astA m))
        ]
  
    return opts { fixityEnv = foldr ($) (fixityEnv opts) fixFns }
  
  -- | 'forM', but concurrency and input order controled by 'Options'.
  forFn :: Options_ x y -> [a] -> (a -> IO b) -> IO [b]
  forFn Options{..} c f
    | randomOrder = fn f =<< shuffleM c
    | otherwise = fn f c
    where
      fn
        | singleThreaded = mapM
        | otherwise = mapConcurrently
  
  -- | Find all files to target for rewriting.
  getTargetFiles :: Options_ a b -> [GroundTerms] -> IO [FilePath]
  -- Always include at least one set of ground terms
  -- This selects all files if the list of rewrites is empty
  getTargetFiles opts [] = getTargetFiles opts [mempty]
  getTargetFiles Options{..} gtss = do
    ignorePred <- maybe onIgnoreErr return =<< vcsIgnorePred verbosity targetDir
    let ignore fp = ignorePred fp || extraIgnorePred fp
    fpSets <- forM (dedup gtss) $ \ gts -> do
      -- See Note [Ground Terms]
      fps <- runGrepChain targetDir verbosity (buildGrepChain targetDir gts targetFiles)
  
      let
        r = filter (not . ignore)
          $ map (normalise . (targetDir </>)) fps
      debugPrint verbosity "Files:" r
      return $ HashSet.fromList r
  
    return $ HashSet.toList $ mconcat fpSets
    where
      dedup = HashSet.toList . HashSet.fromList
      extraIgnorePred =
        let fps = [ normalise (targetDir </> f) | f <- extraIgnores ]
        in \fp -> any (`isPrefixOf` fp) fps
      onIgnoreErr = do
        when (verbosity > Silent) $
          putStrLn "Reading VCS ignore failed! Continuing without ignoring."
        return $ const False
  
  -- | Return a chain of grep commands to find files with relevant groundTerms
  -- If filesGiven is empty, use all *.hs files under targetDir
  buildGrepChain
    :: FilePath
    -> HashSet String
    -> [FilePath]
    -> GrepCommands
  buildGrepChain targetDir gts filesGiven = GrepCommands {initialFileSet=filesGiven, commandChain=commands}
    where
      commands = if null filesGiven
                 then commandsWithoutFiles
                 else commandsWithFiles
  
      commandsWithFiles = case terms of
          [] -> [] -- no processing
          gs -> map normalGrep gs
      commandsWithoutFiles = case terms of
          [] -> [findCmd] -- all .hs files
          g:gs -> recursiveGrep g : map normalGrep gs -- start with recursive grep
  
      findCmd = unwords ["find", quotePath (addTrailingPathSeparator targetDir), "-iname", hsExtension]
      recursiveGrep g = unwords ["grep", "-R", "--include=" ++ hsExtension, "-l", esc g, quotePath targetDir]
      normalGrep gt = unwords ["grep", "-l", esc gt]
  
     -- Limit the number of the shell command we build by only selecting
     -- up to 10 ground terms. The goal is to filter file list down to
     -- a manageable size. It doesn't have to be exact.
      terms = take 10 $ filter p $ HashSet.toList gts
      p [] = False
      p (c:cs)
        | isSpace c = p cs
        | otherwise = isAlphaNum c
  
      hsExtension = "\"*.hs\""
  
      esc s = osquote $ intercalate "[[:space:]]\\+" (words $ escChars s)
      escChars = concatMap escChar
      escChar c
        | c `elem` magicChars = "\\" <> [c]
        | otherwise  = [c]
      magicChars :: [Char]
      magicChars = "*?[#˜=%\\"
  
      osquote s = "'" ++ concatMap escapeQuote s ++ "'"
        where escapeQuote '\'' = "'\"'\"'"
              escapeQuote c = [c]
  
  
  type CommandLine = String
  data GrepCommands = GrepCommands { initialFileSet :: [FilePath], commandChain :: [CommandLine] }
    deriving (Eq, Show)
  
  runGrepChain :: FilePath -> Verbosity -> GrepCommands -> IO [FilePath]
  runGrepChain targetDir verbosity GrepCommands{..} = foldM (commandStep targetDir verbosity)  initialFileSet commandChain
  
  -- | run a command with a list of files as quoted arguments
  commandStep :: FilePath -> Verbosity -> [FilePath]-> CommandLine -> IO [FilePath]
  commandStep targetDir verbosity files cmd =
    fmap concat $ forM (chunkifyArgs $ map quotePath files) $ \args ->
      doCmd targetDir verbosity $ cmd <> " " <> unwords args
  
  quotePath :: FilePath -> FilePath
  quotePath x = "'" <> x <> "'"
  
  -- Split arguments into chunks whose size does not exceed 'argMax'.
  -- Note that @chunkifyArgs []@ returns @[[]]@ intentionally.
  chunkifyArgs :: [String] -> [[String]]
  chunkifyArgs = go 0 []
    where
      go _ acc [] = [reverse acc]
      go n acc (arg: args)
        | n <= 0 || n + len <= argMax = go (n + len) (arg: acc) args
        | otherwise = reverse acc: go len [arg] args
        where
          len = length arg + 1
  
  -- | Hardcoded default @--max-args=max-args@ used by @xargs@.
  argMax :: Int
  argMax = 128 * 1024
  
  doCmd :: FilePath -> Verbosity -> String -> IO [FilePath]
  doCmd targetDir verbosity shellCmd = do
    debugPrint verbosity "shellCmd:" [shellCmd]
    let cmd = (shell shellCmd) { cwd = Just targetDir }
    (_ec, fps, _) <- readCreateProcessWithExitCode cmd ""
    return $ lines fps
  