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.
  -}
  
  {-# OPTIONS_GHC -Wno-orphans #-}
  
  module HieIndexer.Main (main) where
  
  import Control.Monad
  import Control.Monad.State (StateT, evalStateT, get, liftIO, put)
  import qualified Data.ByteString as BS
  import Data.Default
  import Data.IORef
  import Data.List.NonEmpty (NonEmpty)
  import qualified Data.List.NonEmpty as NonEmpty
  import qualified Data.Map as AMap
  import qualified Data.Set as Set
  import qualified Data.HashSet as HashSet
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Options.Applicative
  import System.Directory
  import System.FilePath
  import qualified Thrift.Protocol.Compact
  import Text.Printf
  
  import HieDb (withHieFile)
  import HieDb.Utils (makeNc)
  import HieDb.Types (runDbM)
  
  import Util.EventBase
  import Util.Log
  
  import qualified Glean
  import qualified Glean.LocalOrRemote as Glean
  import Glean.Impl.ConfigProvider (ConfigAPI)
  import Glean.Init (withOptions)
  import Glean.Schema.Builtin.Types (schema_id)
  import qualified Glean.Schema.Hs as Hs
  import qualified Glean.Schema.Src as Src
  import qualified Glean.Util.ConfigProvider as Glean
  
  import qualified HieIndexer.Options as Hie
  import HieIndexer.Index
  import HieIndexer.Options
  
  main :: IO ()
  main = do
    let parser =
          (,) <$> infoParser Hie.options <*> Glean.options
        opts = info (helper <*> parser) fullDesc
    withOptions opts $ \((cfg, mode), svc) ->
      withEventBaseDataplane $ \evb ->
        Glean.withConfigProvider Glean.defaultConfigOptions $
          \(cfgAPI :: ConfigAPI) ->
            Glean.withBackendWithDefaultOptions
              evb
              cfgAPI
              svc
              (case mode of
                BinaryMode _ given_schema_id -> Just given_schema_id
                _ -> Just schema_id)
              $ case mode of
                  Hie.BinaryMode{..} ->
                    outputMain cfg outputPath schemaId
                  Hie.WriteMode{..} ->
                    defaultMain cfg repo
  
  defaultMain
    :: Glean.Backend b
    => HieIndexerOptions
    -> Glean.Repo
    -> b
    -> IO ()
  defaultMain cfg repo backend = do
    paths <- getHieFilesIn (NonEmpty.toList (hiePaths cfg))
  
    logInfo $ printf "Indexing %d files" (HashSet.size paths)
  
    let Glean.Repo{..} = repo
  
    logInfo $ printf "Writing to Glean DB %s/%s" repo_name repo_hash
  
    Glean.withSender backend repo allPredicates def $ \sender -> do
      Glean.withWriter sender def $
        indexHieFiles paths (srcPaths cfg)
  
    predicates <-
      Glean.schemaInfo_predicateIds
        <$> Glean.getSchemaInfo backend (Just repo)
              def { Glean.getSchemaInfo_omit_source = True }
    repoStats <- Glean.predicateStats backend repo Glean.ExcludeBase
    let readableStats =
          [ printf " - %s: count = %d, size = %d"
            (Text.unpack predicateRef_name)
            predicateStats_count
            predicateStats_size
          | (p, Glean.PredicateStats {..}) <- AMap.toList repoStats
          , Just Glean.PredicateRef {..} <- [AMap.lookup p predicates]
          ]
  
    logInfo $ unlines $ "Repo stats: " : readableStats
  
  allPredicates :: [Glean.SchemaPredicates]
  allPredicates = [Src.allPredicates, Hs.allPredicates]
  
  outputMain
    :: Glean.Backend b
    => HieIndexerOptions
    -> FilePath
    -> Glean.SchemaId
    -> b
    -> IO ()
  outputMain cfg out schema_id backend = do
    paths <- getHieFilesIn (NonEmpty.toList (hiePaths cfg))
    ((), batch) <-
      Glean.withBatchWriter backend schema_id Nothing def $
        indexHieFiles paths (srcPaths cfg)
    BS.writeFile out (Thrift.Protocol.Compact.serializeCompact batch)
  
  indexHieFiles
    :: HashSet.HashSet FilePath
    -> NonEmpty Text
    -> Glean.Writer
    -> IO ()
  indexHieFiles paths srcs writer =
    forM_ (HashSet.toList paths) $ \f -> do
      nc <- newIORef =<< makeNc
      runDbM nc $ do
        withHieFile f $ \h ->
          liftIO $ indexHieFile writer srcs f h
  
  {- | Recursively search for @.hie@ and @.hie-boot@  files in given directory
     avoiding loops due to symlinks
  -}
  getHieFilesIn :: [FilePath] -> IO (HashSet.HashSet FilePath)
  getHieFilesIn = flip evalStateT mempty . foldMap go
    where
      go path = do
        seen <- get
        path' <- liftIO $ canonicalizePath path
        if Set.member path' seen
          then return HashSet.empty
          else do
            put (Set.insert path' seen)
            isFile <- liftIO $ doesFileExist path'
            if isFile
              && ( "hie" `isExtensionOf` path'
                    || "hie-boot" `isExtensionOf` path'
                 )
              then do
                return (HashSet.singleton path')
              else do
                isDir <- liftIO $ doesDirectoryExist path'
                if isDir
                  then do
                    cnts <- liftIO $ listDirectory path'
                    foldMap (go . (path' </>)) cnts
                  else return HashSet.empty
  
  instance (Monoid a, Monad m) => Semigroup (StateT s m a) where
    (<>) = liftA2 (<>)
  
  instance (Monoid a, Monad m) => Monoid (StateT s m a) where
    mempty = pure mempty
  