Completion.hs

{-# LANGUAGE TupleSections #-}

module Hackport.Completion
    ( categoryCompleter
    , cabalPackageCompleter
    , trieFile
    ) where

import Control.Applicative
import Control.Monad.Trans.Maybe
import Data.Binary
import qualified Data.List as L
import Data.Maybe (catMaybes)
import qualified Data.Set as Set
import Data.Set (Set)
import GHC.Compact
import qualified Options.Applicative as Opt
import System.Exit (ExitCode(..))
import System.FilePath
import System.IO.Error (tryIOError)
import System.Process hiding (env)
import Text.PrettyPrint (render)

import qualified Data.Trie as Trie
import Data.Trie (Trie)
import Hackport.Env
import Hackport.Util (withHackportContext, hackportDir)
import Util

import Distribution.Client.GlobalFlags (RepoContext(..))
import Distribution.Client.IndexUtils (getSourcePackages)
import Distribution.Client.Types (packageIndex)
import Distribution.Compat.Parsing (choice)
import Distribution.Pretty (Pretty, pretty)
import Distribution.Solver.Types.PackageIndex (allPackages)
import Distribution.Solver.Types.SourcePackage (SourcePackage(..))
import qualified Distribution.Verbosity as V

-- | Read the @profiles/category@ file from:
--
--   * @gentoo@  (determined by portageq)
--
--   The following might be read from in the future but it probably isn't
--   worth it for the extra time taken -- there most likely won't be any
--   extra categories that don't exist in @::gentoo@:
--
--   * @haskell@ (determined by portageq)
--   * The repository reported by 'getOverlayPath'
--
--   This is used exclusively for shell completion for the @make-ebuild@
--   subcommand.
categoryCompleter :: Opt.Completer
categoryCompleter = Opt.mkCompleter $ \s -> do
    let env = GlobalEnv V.silent Nothing Nothing
    gentoo <- runEnv gentooOverlay () env

    let repos = catMaybes [gentoo]

    sets <- traverse repoCategories repos
    pure $ Set.toAscList $ Set.filter (s `L.isPrefixOf`) $ Set.unions sets

  where
    gentooOverlay :: Env () (Maybe FilePath)
    gentooOverlay = getRepoPath "gentoo"

    getRepoPath :: String -> Env () (Maybe FilePath)
    getRepoPath repo =
        runOrDie
            "/usr/bin/portageq"
            ["get_repo_path", "/", repo]
            $ \out -> case out of
                [p] -> Just p
                _   -> Nothing

    repoCategories :: FilePath -> IO (Set String)
    repoCategories repoPath = do
        let catsFile = repoPath </> "profiles" </> "categories"
        foldMap (Set.fromList . filter (not . null) . lines)
            -- Returns a Nothing (wrapped in IO) if `readFile` fails
            <$> optional (readFile catsFile)

-- | Read a list of package names/verisons on Hackage. This requires the user
--   to have run @hackport update@ at least once.
--
--   Internally, this creates a 'Trie' which is cached in
--   @~/.hackport/packages.cache@. The cache is removed when @hackport update@
--   is called, which causes the cache file to be regenerated.
--
--   This is used exclusively for shell completion for the @merge@
--   subcommand.
cabalPackageCompleter :: Opt.Completer
cabalPackageCompleter = Opt.mkCompleter $ \s ->
    runEnv (go s) () $ GlobalEnv V.silent Nothing Nothing
  where
    go s = withHackportContext $ \_ repoContext -> do
        Just trie <- runMaybeT $ choice
            [ -- Check to see if there is an existing cache.
              readTrie
              -- Otherwise, create a new trie from the Hackage database
            , lift $ createTrie repoContext
            ]

        let endings = map fst $ maybe [] Trie.toList
                $ Trie.lookup (getCompact trie) s

        pure $ map (s ++) endings

    -- A (Trie Char ()) is useful for storing/searching a list of strings
    createTrie repoContext = do
        pkgs <- allPV repoContext
        trie <- Trie.buildCompact (,()) pkgs
        writeTrie (getCompact trie)
        pure trie

    readTrie :: MonadIO m
        => MaybeT m (Compact (Trie Char ()))
    readTrie = MaybeT $ liftIO $ optional
        $ trieFile >>= decodeFile >>= compact

    writeTrie :: MonadIO m => Trie Char () -> m ()
    writeTrie t = do
        f <- trieFile
        liftIO $ encodeFile f t

    allPV :: MonadIO m => RepoContext -> m [String]
    allPV repoContext = do
        index <- liftIO $ packageIndex <$> getSourcePackages V.silent repoContext
        let allP = allPackages index
        pure $ (\(SourcePackage i _ _ _) -> showPretty i) <$> allP

trieFile :: MonadIO m => m FilePath
trieFile = do
    d <- hackportDir
    pure $ d </> "packages" <.> "cache"

runOrDie :: FilePath -> [String] -> ([String] -> Maybe a) -> Env env (Maybe a)
runOrDie p ss f = do
    r <- liftIO
        $ tryIOError
        $ readProcessWithExitCode p ss ""
    case r of
        Left  e            -> die (show e)
        Right (ec, out, err) -> do
            case ec of
                ExitSuccess ->
                    pure $ f $ filter (not . null) $ lines $ out
                _ -> die $ unwords [out, err]

showPretty :: Pretty a => a -> String
showPretty = render . pretty