Main.hs

{-# LANGUAGE CPP #-}

module Main where

import Control.Concurrent.Async
import Control.Exception (evaluate)
import NixTree.App
import NixTree.PathStats
import qualified Options.Applicative as Opts
import qualified Options.Applicative.Help.Pretty as Opts
import System.Directory (XdgDirectory (XdgState), doesDirectoryExist, getHomeDirectory, getXdgDirectory, pathIsSymbolicLink)
import System.Exit (ExitCode (..))
import System.FilePath ((</>))
import System.IO (hPutStrLn)
import System.ProgressBar hiding (msg)

version :: Text
version = VERSION_nix_tree

data Opts = Opts
  { oInstallables :: [Installable],
    oStore :: Maybe String,
    oFile :: Maybe FilePath,
    oVersion :: Bool,
    oDerivation :: Bool,
    oImpure :: Bool,
    oDot :: Bool
  }

optsParser :: Opts.ParserInfo Opts
optsParser =
  Opts.info (parser <**> Opts.helper) $
    mconcat
      [ Opts.progDesc "Interactively browse dependency graphs of Nix derivations.",
        Opts.fullDesc,
        Opts.footerDoc (Just keybindingsHelp)
      ]
  where
    parser :: Opts.Parser Opts
    parser =
      Opts
        <$> many
          ( Installable
              <$> Opts.strArgument @Text
                ( Opts.metavar "INSTALLABLE"
                    <> Opts.helpDoc
                      ( Just $
                          Opts.vsep
                            [ "A store path or a flake reference.",
                              "Paths default to \"~/.nix-profile\" and \"/var/run/current-system\""
                            ]
                      )
                )
          )
        <*> optional
          ( Opts.strOption
              ( Opts.long "store"
                  <> Opts.metavar "STORE"
                  <> Opts.helpDoc
                    ( Just $
                        Opts.vsep
                          [ "The URL of the Nix store, e.g. \"daemon\" or \"https://cache.nixos.org\"",
                            "See \"nix help-stores\" for supported store types and settings."
                          ]
                    )
              )
          )
        <*> optional
          ( Opts.strOption
              ( Opts.long "file"
                  <> Opts.metavar "FILE"
                  <> Opts.helpDoc
                    ( Just $ Opts.vsep ["Interpret installables as attribute paths relative to the Nix expression stored in file."]
                    )
              )
          )
        <*> Opts.switch (Opts.long "version" <> Opts.help "Show the nix-tree version")
        <*> Opts.switch (Opts.long "derivation" <> Opts.help "Operate on the store derivation rather than its outputs")
        <*> Opts.switch (Opts.long "impure" <> Opts.help "Allow access to mutable paths and repositories")
        <*> Opts.switch (Opts.long "dot" <> Opts.help "Print the dependency graph in dot format")

    keybindingsHelp :: Opts.Doc
    keybindingsHelp =
      Opts.vsep
        [ "Keybindings:",
          Opts.indent 2 . Opts.vsep $ map Opts.pretty (lines helpText)
        ]

showAndFail :: Text -> IO a
showAndFail msg = do
  hPutStrLn stderr . toString $ "Error: " <> msg
  exitWith (ExitFailure 1)

isValidRoot :: FilePath -> IO Bool
isValidRoot path = do
  isExistingDirectory <- doesDirectoryExist path
  if isExistingDirectory
    then -- We need to check that it's a symlink (presumably to nix store),
    -- because if it is just a directory, nix will try to interpret it as a flake.
    -- We do doesDirectoryExist before pathIsSymbolicLink,
    -- because the latter will fail if the path does not exist.
      pathIsSymbolicLink path
    else return False

main :: IO ()
main = do
  opts <-
    Opts.customExecParser
      (Opts.prefs $ Opts.columns 120)
      optsParser

  when (opts & oVersion) $ do
    putTextLn $ "nix-tree " <> version
    exitSuccess

  installables <- case opts & oInstallables of
    p : ps ->
      return $ p :| ps
    [] -> do
      home <- getHomeDirectory
      nixXdgDirectory <- getXdgDirectory XdgState "nix/profile"
      homeManagerDirectory <- getXdgDirectory XdgState "nix/profiles/home-manager"
      roots <-
        filterM
          isValidRoot
          [ home </> ".nix-profile",
            nixXdgDirectory,
            homeManagerDirectory,
            "/var/run/current-system"
          ]
      case roots of
        [] -> showAndFail "No store path given."
        p : ps -> return . fmap (Installable . toText) $ p :| ps

  let seo =
        StoreEnvOptions
          { seoIsDerivation = opts & oDerivation,
            seoIsImpure = opts & oImpure,
            seoStoreURL = opts & oStore,
            seoFile = opts & oFile
          }

  withStoreEnv seo installables $ \env' -> do
    let env = calculatePathStats env'
        allPaths = seAll env

    bar <- newProgressBar defStyle {stylePostfix = exact} 4 (Progress 0 (length allPaths) ())
    allPaths
      & toList
      & chunks 50
      & mapConcurrently_ (mapM_ (\p -> evaluate (rnf p) >> incProgress bar 1))

    if opts & oDot
      then putTextLn $ storeEnvToDot env
      else run env

chunks :: Int -> [a] -> [[a]]
chunks _ [] = []
chunks n xs =
  let (ys, zs) = splitAt n xs
   in ys : chunks n zs