    {-# LANGUAGE LambdaCase #-}

    module Test.DocTest.Internal.Nix where

    import Control.Monad (msum)
    import Control.Monad.Extra (ifM)
    import Control.Monad.Trans.Maybe
    import Data.Bool (bool)
    import Data.List (intercalate, isSuffixOf)
    import Data.Maybe (isJust)
    import Data.Version
    import GHC.Base (mzero)
    import System.Directory
    import System.Environment (lookupEnv)
    import System.FilePath ((</>), isDrive, takeDirectory)
    import System.Process (readProcess)

    import GHC.Data.Maybe (liftMaybeT)
    import System.Info (fullCompilerVersion)

    -- | E.g. @9.0.2@
    compilerVersionStr :: String
    compilerVersionStr :: String
    compilerVersionStr = String -> [String] -> String
    forall a. [a] -> [[a]] -> [a]
    intercalate String
    "." ((Int -> String) -> [Int] -> [String]
    forall a b. (a -> b) -> [a] -> [b]
    map Int -> String
    forall a. Show a => a -> String
    show (Version -> [Int]
    versionBranch Version
    fullCompilerVersion))

    -- | Traverse upwards until one of the following conditions is met:
    --
    --   * Current working directory is either root or a home directory
    --   * The predicate function returns 'Just'
    --
    findDirectoryUp :: (FilePath -> IO (Maybe a)) -> MaybeT IO a
    findDirectoryUp :: forall a. (String -> IO (Maybe a)) -> MaybeT IO a
    findDirectoryUp String -> IO (Maybe a)
    f = do
      home <- IO String -> MaybeT IO String
    forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
    liftMaybeT IO String
    getHomeDirectory
      MaybeT (go home =<< getCurrentDirectory)
     where
      go :: String -> String -> IO (Maybe a)
    go String
    home String
    cwd
        | String -> Bool
    isDrive String
    cwd = Maybe a -> IO (Maybe a)
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe a
    forall a. Maybe a
    Nothing
        | String
    cwd String -> String -> Bool
    forall a. Eq a => a -> a -> Bool
    == String
    home = Maybe a -> IO (Maybe a)
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe a
    forall a. Maybe a
    Nothing
        | Bool
    otherwise =
          String -> IO (Maybe a)
    f String
    cwd IO (Maybe a) -> (Maybe a -> IO (Maybe a)) -> IO (Maybe a)
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
            Just a
    a -> Maybe a -> IO (Maybe a)
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (a -> Maybe a
    forall a. a -> Maybe a
    Just a
    a)
            Maybe a
    Nothing -> String -> String -> IO (Maybe a)
    go String
    home (String -> String
    takeDirectory String
    cwd)

    -- | Like 'findDirectoryUp', but takes a predicate function instead. If the predicate
    -- yields 'True', the filepath is returned.
    findDirectoryUpPredicate :: (FilePath -> IO Bool) -> MaybeT IO FilePath
    findDirectoryUpPredicate :: (String -> IO Bool) -> MaybeT IO String
    findDirectoryUpPredicate String -> IO Bool
    f = (String -> IO (Maybe String)) -> MaybeT IO String
    forall a. (String -> IO (Maybe a)) -> MaybeT IO a
    findDirectoryUp (\String
    fp -> Maybe String -> Maybe String -> Bool -> Maybe String
    forall a. a -> a -> Bool -> a
    bool Maybe String
    forall a. Maybe a
    Nothing (String -> Maybe String
    forall a. a -> Maybe a
    Just String
    fp) (Bool -> Maybe String) -> IO Bool -> IO (Maybe String)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> IO Bool
    f String
    fp)

    -- | Find the root of the Cabal project relative to the current directory.
    findCabalProjectRoot :: MaybeT IO FilePath
    findCabalProjectRoot :: MaybeT IO String
    findCabalProjectRoot =
      [MaybeT IO String] -> MaybeT IO String
    forall (t :: * -> *) (m :: * -> *) a.
    (Foldable t, MonadPlus m) =>
    t (m a) -> m a
    msum
        [ (String -> IO Bool) -> MaybeT IO String
    findDirectoryUpPredicate String -> IO Bool
    containsCabalProject
        , (String -> IO Bool) -> MaybeT IO String
    findDirectoryUpPredicate String -> IO Bool
    containsCabalPackage
        ]
     where
      containsCabalPackage :: FilePath -> IO Bool
      containsCabalPackage :: String -> IO Bool
    containsCabalPackage String
    fp = String -> [String] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem String
    "cabal.project" ([String] -> Bool) -> IO [String] -> IO Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> IO [String]
    getDirectoryContents String
    fp

      containsCabalProject :: FilePath -> IO Bool
      containsCabalProject :: String -> IO Bool
    containsCabalProject String
    fp = (String -> Bool) -> [String] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (String
    ".cabal" String -> String -> Bool
    forall a. Eq a => [a] -> [a] -> Bool
    `isSuffixOf`) ([String] -> Bool) -> IO [String] -> IO Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> IO [String]
    getDirectoryContents String
    fp

    -- | Find the local package database in @dist-newstyle@.
    findLocalPackageDb :: MaybeT IO FilePath
    findLocalPackageDb :: MaybeT IO String
    findLocalPackageDb = do
      projectRoot <- MaybeT IO String
    findCabalProjectRoot
      let
        relDir = String
    "dist-newstyle" String -> String -> String
    </> String
    "packagedb" String -> String -> String
    </> (String
    "ghc-" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    compilerVersionStr)
        absDir = String
    projectRoot String -> String -> String
    </> String
    relDir
      ifM
        (liftMaybeT (doesDirectoryExist absDir))
        (return absDir)
        mzero

    -- | Are we running in a Nix shell?
    inNixShell :: IO Bool
    inNixShell :: IO Bool
    inNixShell = Maybe String -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> IO (Maybe String)
    lookupEnv String
    "IN_NIX_SHELL"

    -- | Are we running in a Nix build environment?
    inNixBuild :: IO Bool
    inNixBuild :: IO Bool
    inNixBuild = Maybe String -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> IO (Maybe String)
    lookupEnv String
    "NIX_BUILD_TOP"

    getLocalCabalPackageDbArgs :: IO [String]
    getLocalCabalPackageDbArgs :: IO [String]
    getLocalCabalPackageDbArgs = do
      MaybeT IO String -> IO (Maybe String)
    forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
    runMaybeT MaybeT IO String
    findLocalPackageDb IO (Maybe String) -> (Maybe String -> IO [String]) -> IO [String]
    forall a b. IO a -> (a -> IO b) -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
         Maybe String
    Nothing -> [String] -> IO [String]
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []
         Just String
    s -> [String] -> IO [String]
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [String
    "-package-db", String
    s]

    getLocalNixPackageDbArgs :: IO [String]
    getLocalNixPackageDbArgs :: IO [String]
    getLocalNixPackageDbArgs = do
      pkgDb <- String -> IO String
    makeAbsolute (String
    "dist" String -> String -> String
    </> String
    "package.conf.inplace")
      ifM
        (doesDirectoryExist pkgDb)
        (pure ["-package-db", pkgDb])
        (pure [])

    -- | Get global package db; used in a NIX_SHELL context
    getGlobalPackageDb :: IO String
    getGlobalPackageDb :: IO String
    getGlobalPackageDb = String -> String
    forall a. HasCallStack => [a] -> [a]
    init (String -> String) -> IO String -> IO String
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> String -> [String] -> String -> IO String
    readProcess String
    "ghc" [String
    "--print-global-package-db"] String
    ""

    -- | Get flags to be used when running in a Nix context (either in a build, or a
    -- shell).
    getNixGhciArgs :: IO [String]
    getNixGhciArgs :: IO [String]
    getNixGhciArgs =
      IO Bool -> IO [String] -> IO [String] -> IO [String]
    forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
    ifM IO Bool
    inNixShell IO [String]
    goShell (IO Bool -> IO [String] -> IO [String] -> IO [String]
    forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
    ifM IO Bool
    inNixBuild IO [String]
    goBuild ([String] -> IO [String]
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure []))
     where
      goShell :: IO [String]
    goShell = do
        globalPkgDb <- IO String
    getGlobalPackageDb
        localPkgDbFlag <- getLocalCabalPackageDbArgs
        let globalDbFlag = [String
    "-package-db", String
    globalPkgDb]
        pure (defaultArgs ++ globalDbFlag ++ localPkgDbFlag)

      goBuild :: IO [String]
    goBuild = do
        localDbFlag <- IO [String]
    getLocalNixPackageDbArgs
        pure (defaultArgs ++ localDbFlag)

      defaultArgs :: [String]
    defaultArgs =
        [ String
    "-package-env", String
    "-"

        -- Nix doesn't always expose the GHC library (_specifically_ the GHC lib) even
        -- if a package lists it as a dependency. This simply always exposes it as a
        -- workaround.
        , String
    "-package", String
    "ghc"
        ]