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.
  -}
  
  {-# LANGUAGE DeriveTraversable #-}
  module Glean.Shell.Types (
    Parse(..), Statement(..), JSONQuery(..), AngleQuery(..),
    SchemaQuery(..),
    Stats(..),
    ShellState(..),
    Eval(..),
    ExpandResults(..),
    withBackend,
    getState,
    getRepo,
    setRepo,
  ) where
  
  import Control.Concurrent
  import Control.Exception
  import Data.Default
  import Data.Bifunctor
  import Data.Functor (($>))
  import qualified Control.Monad.Catch as C
  import Control.Monad.IO.Class (MonadIO, liftIO)
  import qualified Control.Monad.Trans.State.Strict as State
  import Data.Int
  import Data.List
  import qualified Data.Map as Map
  import Data.Maybe
  import Compat.Prettyprinter as Pretty
  #if !MIN_VERSION_haskeline(0,8,0)
  import qualified System.Console.Haskeline as Haskeline
  #endif
  import System.IO
  import qualified Text.JSON as JSON
  import qualified Text.Parsec as P
  import qualified Text.Parsec.Language as P
  import qualified Text.Parsec.Token as P
  
  import Glean
  import qualified Glean.Types as Thrift
  import Glean.LocalOrRemote (LocalOrRemote)
  import Glean.Database.Config (ProcessedSchema(..), processSchema)
  import Glean.Util.Some
  
  data Statement pat
    = Command String String
    | Pattern pat
    | FactRef Fid
  
  data ExpandResults
    = ExpandRecursive
    | ExpandPredicates [SourceRef]
  
  type Parser = P.Parsec String ()
  
  class Parse a where
    parse :: Parser a
  
  instance Parse Fid where
    parse = Fid . fromInteger <$> P.braces lexer (P.natural lexer)
  
  instance Parse pat => Parse (Statement pat) where
    parse = P.choice [command, P.try factref, ptrn]
      where
        command = split <$> (P.char ':' *> P.getInput <* P.setInput "")
        ptrn = Pattern <$>  parse
        factref = FactRef <$> parse
        split s
          | (cmd,' ':arg) <- break (==' ') s = Command cmd arg
          | otherwise = Command s ""
  
  data AngleQuery = AngleQuery
    { angleQueryDeprecatedRec :: Bool
    , angleQueryStored :: Bool
    , angleQuery :: String
    }
  
  instance Parse AngleQuery where
    parse = AngleQuery
      <$> P.option False (P.char '!' $> True)
      <*> P.option False (P.char '*' $> True)
      <*> P.many P.anyChar
  
  data JSONQuery = JSONQuery
    { jsonQueryPred :: String
    , jsonQueryDeprecatedRec :: Bool
    , jsonQueryStored :: Bool
    , jsonQuery :: String
    }
  
  instance Parse JSONQuery where
    parse = JSONQuery
      <$> P.identifier lexer
      <*> P.option False (P.char '!' $> True)
      <*> P.option False (P.char '*' $> True)
      <*> P.many P.anyChar
  
  lexer :: P.TokenParser st
  lexer = P.makeTokenParser P.emptyDef
    { P.identLetter = P.alphaNum P.<|> P.oneOf "_." }
  
  
  data SchemaQuery = SchemaQuery
    { sqPredicate :: String
    , sqRecursive :: ExpandResults
    , sqStored :: Bool
    , sqQuery :: String
    , sqCont :: Maybe Thrift.UserQueryCont
    , sqTransform :: Maybe (JSON.JSValue -> JSON.Result JSON.JSValue)
    , sqSyntax :: Thrift.QuerySyntax
    , sqOmitResults :: Bool
    , sqContinue :: Bool
    }
  
  data Stats = NoStats | SummaryStats | FullStats
    deriving Eq
  
  data ShellState = ShellState
    { backend :: Some LocalOrRemote
    , repo :: Maybe Repo
    , schemas :: Maybe ProcessedSchema
    , schemaInfo :: Maybe Thrift.SchemaInfo
    , useSchemaId :: Thrift.SelectSchema
    , limit :: Int64
    , timeout :: Maybe Int64
    , stats :: Stats
    , lastSchemaQuery :: Maybe SchemaQuery
    , updateSchema :: Maybe (Eval ())
    , isTTY :: Bool
    , pageWidth :: Maybe PageWidth
    , expandResults :: ExpandResults
    , pager :: Bool
    , outputHandle :: MVar System.IO.Handle
    , debug :: Thrift.QueryDebugOptions
    , client_info :: Thrift.UserQueryClientInfo
    , query_file :: FilePath
    }
  
  newtype Eval a = Eval
    { unEval :: State.StateT ShellState IO a
    }
    deriving
      ( Functor
      , Applicative
      , Monad
  #if !MIN_VERSION_haskeline(0,8,0)
      , Haskeline.MonadException
  #endif
      , MonadIO
      , C.MonadCatch
      , C.MonadMask
      , C.MonadThrow )
  
  withBackend :: (forall b . LocalOrRemote b => b -> Eval a) -> Eval a
  withBackend f = do
    state <- getState
    case backend state of
      Some b -> f b
  
  getState :: Eval ShellState
  getState = Eval State.get
  
  getRepo :: Eval (Maybe Repo)
  getRepo = repo <$> getState
  
  setRepo :: Repo -> Eval ()
  setRepo r = do
    sel <- useSchemaId <$> getState
    withBackend $ \backend -> do
      info@SchemaInfo{..} <- liftIO $
        getSchemaInfo backend (Just r) def { getSchemaInfo_select = sel }
      let sids = map (first Thrift.SchemaId) $ Map.toList $ case sel of
            Thrift.SelectSchema_current{} -> schemaInfo_schemaIds
            Thrift.SelectSchema_stored{} -> schemaInfo_dbSchemaIds
            Thrift.SelectSchema_schema_id id
              | Map.member (unSchemaId id) schemaInfo_schemaIds ->
                schemaInfo_schemaIds
              | Map.member (unSchemaId id) schemaInfo_dbSchemaIds ->
                schemaInfo_dbSchemaIds
              | Just ids <- find (unSchemaId id `Map.member`)
                schemaInfo_otherSchemaIds -> ids
            _otherwise -> Map.empty
      Eval $ State.modify $ \s -> s
        { repo = Just r
        , schemaInfo = Just info }
      proc <- either (liftIO . throwIO . ErrorCall) return $
        processSchema (listToMaybe sids) schemaInfo_schema
      Eval $ State.modify $ \s -> s { schemas = Just proc }
  