content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Util.IO
    ( atomicIO
    , die
    , readFileUTF8
    , writeFileUTF8
    , writeFileAtomicUTF8
    , readJSON
    , eitherReadJSON
    , writeJSON
    , writeJSON'
    , removeIfExists
    , readFileIfExists
    , writeFileIfChanged
    , writeFileAtomically
    , writeTextIfChanged
    , loudlyWriteTextIfChanged
    , writeTextIfChangedWith
    , writeUtf8StringIfChanged
    , loudlyWriteUtf8StringIfChanged
    , writeUtf8StringIfChangedWith
    , Verbosity(..)
    , writeFileUTF8Text
    , readFileUTF8Text
    , getDevserverUser
    , getDevserverHostInfo
    , getHostInfo
    , getHostname
    , getUserForProcess
    , getGroupForProcess
    , getUsername
    , getUserUnixname
    , HostInfo(..)
    , copyDirectoryContents
    , copyDirectoryContents_
    , listDirectoryRecursive
    , saveStdout
    , saveStderr
    , slowIO
    , isModifiedAfter
    , safeRemovePathForcibly
    , withLazy
    ) where
  
  import Control.Concurrent
  import Control.Concurrent.Async
  import Control.Exception as Exception
  import Control.Monad
  import Control.Monad.Extra
  import Data.Aeson
  import Data.Aeson.Encode.Pretty
  import Data.Foldable
  import Data.List
  import Data.List.Split
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.IO as TextIO
  import GHC.IO.Handle
  import System.Directory
  import System.Exit hiding (die)
  import System.FilePath
  import System.IO
  import System.IO.Error
  import System.IO.Unsafe
  import System.Posix.User
  import System.Process
  import System.Environment
  import Text.Printf
  import Util.String (strip)
  import qualified Data.ByteString.Lazy as B
  
  -- | Performs some IO atomically. This is presumed to be only used
  -- to avoid output interleaving where verbose information is printed
  -- concurrently.
  atomicIO :: IO a -> IO a
  atomicIO = withMVar atomicIOLock . const
  
  {-# NOINLINE atomicIOLock #-}
  atomicIOLock :: MVar ()
  atomicIOLock = unsafePerformIO $ newMVar ()
  
  -- | Exit immediately with 'exitCode', showing a message on 'stderr'.
  die :: Int -> String -> IO a
  die exitCode msg =
    hPutStrLn stderr msg >> exitWith (ExitFailure exitCode)
  
  -- | Compare a file's content against some benchmark content to see
  -- whether it's changed. You really want to pass in some kind of lazy
  -- read function here so that (in-)equality can be determined as soon
  -- as possible during the file read. Note that if the file doesn't
  -- exist, then we return True, that is, we state that the file has
  -- changed.
  hasFileContentChanged ::
    (Eq a) => (FilePath -> IO a) -> FilePath -> a -> IO Bool
  hasFileContentChanged read filePath benchmarkContent =
    catchIOError
      (liftM (/= benchmarkContent) $ read filePath)
      (\_ -> return True)
  
  -- | Writes a file only if it differs from the content that would be
  -- written.
  writeFileIfChanged :: (Eq a) =>
    (FilePath -> IO a) -> (Handle -> a -> IO ()) -> FilePath -> a -> IO ()
  writeFileIfChanged read hWrite targetPath content = do
    different <- hasFileContentChanged read targetPath content
    when different $ writeFileAtomically hWrite targetPath content
  
  -- | Write a file "atomically" by writing it fully to a temp
  -- directory, then copying it to its final location.
  writeFileAtomically :: (Handle -> a -> IO ()) -> FilePath -> a -> IO ()
  writeFileAtomically hWrite targetPath content =
      Exception.bracketOnError mkTempFile rmTempFile commit
    where
      (targetDir, targetFile) = splitFileName targetPath
      mkTempFile = openTempFile targetDir (targetFile <.> "tmp")
      rmTempFile (tmpPath, handle) = hClose handle >> removeFile tmpPath
      commit (tmpPath, handle) = do
        hWrite handle content
        hClose handle
        renameFile tmpPath targetPath
  
  -- | Remove file if it exists. Returns True if the file existed.
  removeIfExists :: FilePath -> IO Bool
  removeIfExists filePath =
    (removeFile filePath >> return True) `catch` \e ->
      if isDoesNotExistError e
        then return False
        else throwIO e
  
  readFileIfExists :: FilePath -> IO (Maybe String)
  readFileIfExists filePath =
    (Just <$> readFile filePath) `catch` \e ->
      if isDoesNotExistError e
        then return Nothing
        else throwIO e
  
  -- | Read a FromJSON datatype from a file containing a JSON string.
  readJSON :: FromJSON a => FilePath -> IO a
  readJSON f = do
    json <- B.readFile f
    case eitherDecode json of
      Left err -> error $ printf "Failed parsing JSON file %s: %s\n" f err
      Right v  -> return v
  
  eitherReadJSON :: FromJSON a => FilePath -> IO (Either String a)
  eitherReadJSON f = (eitherDecode <$> B.readFile f)
    `catch` (\(e :: Exception.IOException) -> return $ Left (show e))
  
  -- |  Write a ToJSON datatype to a file as pretty-printed JSON.
  writeJSON :: ToJSON a => FilePath -> a -> IO ()
  writeJSON path = B.writeFile path . encodePretty
  
  -- | Write a ToJSON datatype to a file as pretty-printed JSON.
  -- Adds a newline at the end of the file
  writeJSON' :: ToJSON a => Config -> FilePath -> a -> IO ()
  writeJSON' config path obj = B.writeFile path (encodePretty' config obj <> "\n")
  
  -- | Read a UTF-8 encoded file.  Like 'readFile' but forces the
  -- encoding to UTF-8.
  readFileUTF8 :: FilePath -> IO String
  readFileUTF8 path = do
    h <- openFile path ReadMode
    hSetEncoding h utf8
    hGetContents h
  
  -- | Write a UTF-8 encoded file.  Like 'writeFile' but forces the
  -- encoding to UTF-8.
  writeFileUTF8 :: Handle -> String -> IO ()
  writeFileUTF8 h str = do
      hSetEncoding h utf8
      hPutStr h str
  
  -- | Write a UTF-8 encoded file atomically.  Writes to a temporary file
  -- first, and then atomically moves the temporary file to the destination
  -- if successful.
  --
  -- Adapted from Distribution.Simple.Utils.writeFileAtomic in Cabal.
  writeFileAtomicUTF8 :: FilePath -> String -> IO ()
  writeFileAtomicUTF8 targetPath content = do
    let (targetDir, targetFile) = splitFileName targetPath
    Exception.bracketOnError
      (do
        createDirectoryIfMissing True targetDir
        openBinaryTempFileWithDefaultPermissions targetDir $ targetFile <.> "tmp")
      (\(tmpPath, handle) -> hClose handle >> removeFile tmpPath)
      (\(tmpPath, handle) -> do
          hSetEncoding handle utf8
          hPutStr handle content
          hClose handle
          renameFile tmpPath targetPath)
  
  -- | Write a UTF-8 encoded file from text, only if it differs from the content
  -- that would be written. If the output directory will be created if it does
  -- not exist.
  writeTextIfChanged :: FilePath -> Text -> IO ()
  writeTextIfChanged path text = do
    createDirectoryIfMissing True $ takeDirectory path
    writeFileIfChanged readFileUTF8Text writeFileUTF8Text path text
  
  -- Like Util.IO.writeTextIfChanged, but announce upon success.
  loudlyWriteTextIfChanged :: FilePath -> Text -> IO ()
  loudlyWriteTextIfChanged = writeTextIfChangedWith Loud putStrLn
  
  data Verbosity = Loud | Quiet
  
  writeTextIfChangedWith
    :: Verbosity
    -> (String -> IO ())
    -> FilePath
    -> Text
    -> IO ()
  writeTextIfChangedWith verbosity f path text = do
    createDirectoryIfMissing True $ takeDirectory path
    writeFileIfChanged readFileUTF8Text writeFileUTF8Text_ path text
    where
    writeFileUTF8Text_ h txt = do
      writeFileUTF8Text h txt
      case verbosity of
        Loud -> atomicIO $ f ("Written " ++ path)
        Quiet -> return ()
  
  -- | Write a UTF-8 encoded file from text.  Like 'writeFile' but forces the
  -- encoding to UTF-8.
  writeFileUTF8Text :: Handle -> Text -> IO ()
  writeFileUTF8Text h text = do
    hSetEncoding h utf8
    TextIO.hPutStr h text
  
  -- | Read a UTF-8 encoded file as text.
  readFileUTF8Text :: FilePath -> IO Text
  readFileUTF8Text path = do
    h <- openFile path ReadMode
    hSetEncoding h utf8
    TextIO.hGetContents h
  
  -- Like Util.IO.writeTextIfChanged, but works for String
  writeUtf8StringIfChanged :: FilePath -> String -> IO ()
  writeUtf8StringIfChanged path content = do
    createDirectoryIfMissing True $ takeDirectory path
    writeFileIfChanged readFileUTF8 writeFileUTF8 path content
  
  -- Like Util.IO.loudlyWriteTextIfChanged, but works for String.
  loudlyWriteUtf8StringIfChanged :: FilePath -> String -> IO ()
  loudlyWriteUtf8StringIfChanged = writeUtf8StringIfChangedWith Loud putStrLn
  
  -- Like Util.IO.writeTextIfChangedWith, but works for String
  writeUtf8StringIfChangedWith
    :: Verbosity
    -> (String -> IO ())
    -> FilePath
    -> String
    -> IO ()
  writeUtf8StringIfChangedWith verbosity f path content = do
    createDirectoryIfMissing True $ takeDirectory path
    writeFileIfChanged readFileUTF8 writeFileUTF8_ path content
    where
    writeFileUTF8_ h txt = do
      writeFileUTF8 h txt
      case verbosity of
        Loud -> atomicIO $ f ("Written " ++ path)
        Quiet -> return ()
  
  getUserForProcess :: IO String
  getUserForProcess = fmap userName $ getUserEntryForID =<< getRealUserID
  
  getGroupForProcess :: IO String
  getGroupForProcess = fmap groupName $ getGroupEntryForID =<< getRealGroupID
  
  data HostInfo = HostInfo
    { hostname :: Text
    , username :: Text
    }
  
  getHostInfo :: IO HostInfo
  getHostInfo = do
    user <- Text.pack <$> getUsername
    host <- Text.pack <$> getHostname
    return $ HostInfo host user
  
  getDevserverHostInfo :: IO HostInfo
  getDevserverHostInfo = do
    user <- Text.pack <$> getDevserverUser
    host <- Text.pack <$> getHostname
    return $ HostInfo host user
  
  -- | Equivalent of get_devserver_username in www (see
  -- https://fburl.com/codex/05mqipn2).
  getDevserverUser :: IO String
  getDevserverUser = getUserForProcess
  
  -- | Get username using 'id' command line tool
  getUsername :: IO String
  getUsername = init <$> readProcess "id" ["-un"] ""
  
  -- | Get hostname using 'hostname' command line tool
  getHostname :: IO String
  getHostname = strip <$> readProcess "hostname" [] []
  
  -- Try to get a meaninful user unix name. This will not realize that the user
  -- was changed to a non-root account like mysql.
  getUserUnixname :: IO String
  getUserUnixname = do
    user <- getUserForProcess
    fromMaybe user <$> asum
      [ return $ helpful $ Just user
      , helpful <$> getTupperwareUser
      , helpful <$> getDevserverOwner
      ]
    where
      unhelpfulNames = ["root", "twsvcscm", "svcscm", "apache"]
  
      helpful (Just name) | name `elem` unhelpfulNames = Nothing
      helpful mname = mname
  
      getTupperwareUser = lookupEnv "TW_JOB_USER"
  
  getDevserverOwner :: IO (Maybe String)
  getDevserverOwner = do
    res <- Exception.try $ readFile "/etc/devserver.owners"
      :: IO (Either Exception.SomeException String)
    return $ case lines <$> res of
      Right (name:_) -> Just name
      _ -> Nothing
  
  redirectHandle :: Handle -> FilePath -> IO () -> IO ()
  redirectHandle handle file io =
    withFile file WriteMode $ \fh -> do
    bracket (hDuplicate handle) (`hDuplicateTo` handle) $ \_ -> do
      hDuplicateTo fh handle
      io
  
  -- | Run an IO action saving the stderr to a file
  saveStderr :: FilePath -> IO () -> IO ()
  saveStderr = redirectHandle stderr
  
  -- | Run an IO action saving the stdout to a file
  saveStdout :: FilePath -> IO () -> IO ()
  saveStdout = redirectHandle stdout
  
  -- | Copy all the contents of one directory to another directory
  -- and return all the resulting paths.
  -- This will create the target directory if it does not exist.
  
  -- | DISCLAIMER: This is not an atomic operation.
  copyDirectoryContents :: FilePath -> FilePath -> IO [FilePath]
  copyDirectoryContents dirfrom dirto = do
    paths <- listDirectoryRecursive dirfrom
    forM paths $ \path -> do
        let relfile = makeRelative dirfrom path
        let frompath = dirfrom </> relfile
        let topath = dirto </> relfile
        createDirectoryIfMissing True $ takeDirectory topath
        copyFile frompath topath
        return topath
  
  copyDirectoryContents_ :: FilePath -> FilePath -> IO ()
  copyDirectoryContents_ dirfrom dirto
    = void $ copyDirectoryContents dirfrom dirto
  
  listDirectoryRecursive
      :: FilePath
      -> IO [FilePath]
  listDirectoryRecursive dir = do
      fs <- listDirectory dir
      (concat <$>) $ forM fs $ \file -> do
          let f = dir </> file
          isFile <- doesFileExist f
          if isFile
          then return [f]
          else do
              isDir <- doesDirectoryExist f
              if isDir
              then listDirectoryRecursive f
              else return []
  
  
  -- | Perform IO in batches with delays between them. Useful to avoid
  -- filling up fixed size queues.
  slowIO
    :: ([a] -> IO ())                     -- ^ operation to apply to each batch
    -> [a]                                -- ^ all arguments
    -> Int                                -- ^ maximum batch size
    -> Int                                -- ^ delay, in milliseconds
    -> IO ()
  slowIO f args batchSize delay =
    sequence_
      $ intersperse (threadDelay (delay*1000))
      $ map f
      $ chunksOf batchSize args
  
  isModifiedAfter :: FilePath -> FilePath -> IO Bool
  isModifiedAfter fp1 fp2 = ifM (doesFileExist fp1 &&^ doesFileExist fp2)
    (do
      time <- getModificationTime fp1
      time' <- getModificationTime fp2
      return $ time > time')
    (return False)
  
  -- | System.Posix.removeLink calls @unlink()@ with an unsafe FFI call,
  -- which can cause long stalls in some cases. This affects things
  -- defined in terms of it, including 'System.Directory.removeFile',
  -- and 'System.Directory.removePathForcibly' to name a couple.
  --
  -- This function is a replacement for 'removePathForcibly' that forks
  -- an external 'rm' instead.
  safeRemovePathForcibly :: FilePath -> IO ()
  safeRemovePathForcibly path = callProcess "rm" [ "-rf", path ]
  
  -- | Turn a with-style resource scoping function into one that
  -- executes lazily. The @with@ is performed the first time the
  -- supplied getter is invoked, and is released when @withLazy@
  -- returns.
  withLazy
    :: (forall b . (a -> IO b) -> IO b)
      -- ^ with-style function to allocate a scoped resource of type @a@
    -> (IO a -> IO c)
      -- ^ provides a lazy getter for the resource @a@
    -> IO c
  withLazy wit fn = do
    barrier <- newEmptyMVar
    result <- newEmptyMVar
    done <- newEmptyMVar
    let
      get = do _ <- tryPutMVar barrier (); readMVar result
      go = do
        takeMVar barrier
        wit $ \a -> do
          putMVar result a
          takeMVar done
    snd <$> concurrently go (fn get `finally` putMVar done ())
  