{-
  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.
-}

module Glean.Backend.Logging
  ( LoggingBackend(..)
  ) where

import Control.Exception
import qualified Data.ByteString as ByteString
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Map as Map
import Data.Maybe
import Data.Monoid (Sum(getSum, Sum))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Text.Encoding.Error as Text
import TextShow (showt)

import Util.Logger

import Glean.Schema.Util (showRef)
import Glean.Backend.Local ()
import Glean.Logger.Server as Logger
import qualified Glean.Database.List as Database
import qualified Glean.Database.Types as Database
import Glean.Database.Writes (batchOwnedSize, batchDependenciesSize)
import Glean.Logger
import qualified Glean.Types as Thrift
import Util.Time

import Glean.Backend.Types


-- | A logging wrapper for Env. We do it this way because some backend
-- calls invoke other backend calls, and we only want to log the
-- outermost one. For example, userQuery will call queryFact a *lot*,
-- and it would be too expensive to log each and every call to
-- queryFact.
newtype LoggingBackend = LoggingBackend Database.Env

instance Backend LoggingBackend where
  queryFact (LoggingBackend env) repo id =
    loggingAction (runLogRepo "queryFact" env repo) (const mempty) $
      queryFact env repo id
  factIdRange (LoggingBackend env) repo =
    loggingAction (runLogRepo "factIdRange" env repo) (const mempty) $
      factIdRange env repo
  getSchemaInfo (LoggingBackend env) (Just repo) req =
    loggingAction (runLogRepo "getSchemaInfo" env repo) (const mempty) $
      getSchemaInfo env (Just repo) req
  getSchemaInfo (LoggingBackend env) Nothing req =
    loggingAction (runLogCmd "getSchemaInfo" env) (const mempty) $
      getSchemaInfo env Nothing req
  validateSchema (LoggingBackend env) req =
    loggingAction (runLogCmd "validateSchema" env) (const mempty) $
      validateSchema env req
  predicateStats (LoggingBackend env) repo opts =
    loggingAction (runLogRepo "predicateStats" env repo) (const mempty) $
      predicateStats env repo opts
  userQueryFacts (LoggingBackend env) repo req =
    loggingAction (runLogQueryFacts "userQueryFacts" env repo req)
      logQueryResults $
        userQueryFacts env repo req
  userQuery (LoggingBackend env) repo req =
    loggingAction (runLogQuery "userQuery" env repo req) logQueryResults $
      userQuery env repo req
  userQueryBatch (LoggingBackend env) repo reqBatch =
    loggingAction
      (runLogQueryBatch "userQueryBatch" env repo reqBatch)
      logQueryResultsOrException
      (userQueryBatch env repo reqBatch)

  deriveStored (LoggingBackend env) log repo q =
    loggingAction
      (runLogDerivePredicate "deriveStored" env repo q)
      (const mempty)
      (deriveStored env (runLogDerivationResult env log repo q) repo q)

  listDatabases (LoggingBackend env) req =
    loggingAction (runLogCmd "listDatabases" env) (const mempty) $
      Database.listDatabases env req
  getDatabase (LoggingBackend env) repo =
    loggingAction (runLogRepo "getDatabase" env repo) (const mempty) $
      getDatabase env repo

  kickOffDatabase (LoggingBackend env) rq =
    loggingAction
      (runLogKickOff "kickOff" env rq)
      (const mempty) $
      kickOffDatabase env rq
  finishDatabase (LoggingBackend env) repo =
    loggingAction
      (runLogRepo "finishDatabase" env repo)
      (const mempty) $
      finishDatabase env repo
  finalizeDatabase (LoggingBackend env) repo =
    loggingAction
      (runLogRepo "finalizeDatabase" env repo)
      (const mempty) $
      finalizeDatabase env repo
  updateProperties (LoggingBackend env) repo set unset =
    loggingAction
      (runLogRepo "updateProperties" env repo)
      (const mempty) $
      updateProperties env repo set unset

  completePredicates_ (LoggingBackend env) repo preds =
    loggingAction
       (runLogRepo "completePredicates" env repo)
       (const mempty) $
       completePredicates_ env repo preds

  restoreDatabase (LoggingBackend env) loc =
    loggingAction (runLogCmd "restoreDatabase" env) (const mempty) $
      restoreDatabase env loc
  deleteDatabase (LoggingBackend env) repo =
    loggingAction (runLogRepo "deleteDatabase" env repo) (const mempty) $
      deleteDatabase env repo
  enqueueBatch (LoggingBackend env) cbatch =
    loggingAction
      (runLogEnqueueBatch "enqueueBatch" env cbatch)
      (const mempty) $
        enqueueBatch env cbatch
  enqueueJsonBatch (LoggingBackend env) repo batch =
    loggingAction (runLogRepo "enqueueJsonBatch" env repo) (const mempty) $
      enqueueJsonBatch env repo batch
  enqueueBatchDescriptor (LoggingBackend env) repo batch waitPolicy =
    loggingAction
      (runLogRepo "enqueueBatchDescriptor" env repo)
      (const mempty) $
        enqueueBatchDescriptor env repo batch waitPolicy
  pollBatch (LoggingBackend env) handle =
    loggingAction (runLogCmd "pollBatch" env) (const mempty) $
      pollBatch env handle
  displayBackend (LoggingBackend b) = displayBackend b
  hasDatabase (LoggingBackend b) repo = hasDatabase b repo
  schemaId (LoggingBackend b) = schemaId b
  usingShards (LoggingBackend b) = usingShards b
  initGlobalState (LoggingBackend b) = initGlobalState b

runLogKickOff
  :: Text
  -> Database.Env
  -> Thrift.KickOff
  -> GleanServerLog
  -> IO ()
runLogKickOff cmd env Thrift.KickOff{..} log =
  runLogRepo cmd env kickOff_repo $ log <> schemaId
  where
  schemaId = maybe mempty Logger.SetSchemaId $
    HashMap.lookup "glean.schema_id" kickOff_properties

runLogQueryFacts
  :: Text
  -> Database.Env
  -> Thrift.Repo
  -> Thrift.UserQueryFacts
  -> GleanServerLog -> IO ()
runLogQueryFacts cmd env repo Thrift.UserQueryFacts{..} log =
  runLogRepo cmd env repo $ log
    <> maybe mempty logQueryOptions userQueryFacts_options
    <> maybe mempty logQueryClientInfo userQueryFacts_client_info
    <> maybe mempty (Logger.SetSchemaId . Thrift.unSchemaId)
        userQueryFacts_schema_id

runLogQuery
  :: Text
  -> Database.Env
  -> Thrift.Repo
  -> Thrift.UserQuery
  -> GleanServerLog
  -> IO ()
runLogQuery cmd env repo Thrift.UserQuery{..} log = do
  runLogRepo cmd env repo $ mconcat
    [ log
    , Logger.SetQuery
        (Text.decodeUtf8With Text.lenientDecode $
          if ByteString.length userQuery_query > 1024
            then "[truncated] " <> ByteString.take 1024 userQuery_query
            else userQuery_query)
    , Logger.SetPredicate userQuery_predicate
    , maybe mempty (Logger.SetPredicateVersion . fromIntegral)
        userQuery_predicate_version
    , maybe mempty (Logger.SetSchemaId . Thrift.unSchemaId)
        userQuery_schema_id
    , maybe mempty logQueryOptions userQuery_options
    , maybe mempty logQueryClientInfo userQuery_client_info
    ]

runLogQueryBatch
  :: Text
  -> Database.Env
  -> Thrift.Repo
  -> Thrift.UserQueryBatch
  -> GleanServerLog
  -> IO ()
runLogQueryBatch cmd env repo Thrift.UserQueryBatch{..} log =
  runLogRepo cmd env repo $ mconcat
    [ log
    , Logger.SetQuery $ case userQueryBatch_queries of
        [] -> "0 batched queries"
        q:rest -> Text.unlines $
          Text.decodeUtf8With Text.lenientDecode q :
          [" + " <> showt n <> " batched queries"
          | let n = length rest
          , n > 1
          ]
    , Logger.SetPredicate userQueryBatch_predicate
    , maybe mempty (Logger.SetPredicateVersion . fromIntegral)
        userQueryBatch_predicate_version
    , maybe mempty (Logger.SetSchemaId . Thrift.unSchemaId)
        userQueryBatch_schema_id
    , maybe mempty logQueryOptions userQueryBatch_options
    , maybe mempty logQueryClientInfo userQueryBatch_client_info
    ]

runLogEnqueueBatch
  :: Text
  -> Database.Env
  -> Thrift.ComputedBatch
  -> GleanServerLog
  -> IO ()
runLogEnqueueBatch cmd env Thrift.ComputedBatch{..} log =
  let !Thrift.Batch{..} = computedBatch_batch in
  runLogRepo cmd env computedBatch_repo $ mconcat
    [ log
    , Logger.SetBatchFactsSize $ ByteString.length batch_facts
    , Logger.SetBatchFactsCount $ fromIntegral $
        Thrift.batch_count computedBatch_batch
    , Logger.SetBatchOwnedSize $ batchOwnedSize batch_owned
    , Logger.SetBatchDependenciesSize $ batchDependenciesSize batch_dependencies
    ]

logQueryOptions :: Thrift.UserQueryOptions -> GleanServerLog
logQueryOptions Thrift.UserQueryOptions{..} = mconcat
  [ Logger.SetNoBase64Binary userQueryOptions_no_base64_binary
  , Logger.SetExpandResults userQueryOptions_expand_results
  , Logger.SetRecursive userQueryOptions_recursive
  , maybe mempty (Logger.SetMaxResults . fromIntegral)
      userQueryOptions_max_results
  , Logger.SetSyntax $ case userQueryOptions_syntax of
      Thrift.QuerySyntax_JSON -> "JSON"
      Thrift.QuerySyntax_ANGLE -> "Angle"
  , maybe mempty
      ( Logger.SetRequestContinuationSize
      . ByteString.length
      . Thrift.userQueryCont_continuation
      )
      userQueryOptions_continuation
  ]

logQueryClientInfo :: Thrift.UserQueryClientInfo -> GleanServerLog
logQueryClientInfo Thrift.UserQueryClientInfo{..} = mconcat
  [ maybe mempty Logger.SetClientUnixname userQueryClientInfo_unixname
  , Logger.SetClientApplication userQueryClientInfo_application
  , Logger.SetClientName userQueryClientInfo_name
  ]

logQueryResultsOrException
  :: [Thrift.UserQueryResultsOrException] -> GleanServerLog
logQueryResultsOrException results = mconcat
  [
    Logger.SetResults $ getSum $ foldMap (Sum . countQueryResults)
      [ r | Thrift.UserQueryResultsOrException_results r <- results]
  ]

logQueryResults :: Thrift.UserQueryResults -> GleanServerLog
logQueryResults it@Thrift.UserQueryResults{..} = mconcat
  [ Logger.SetResults $ countQueryResults it
  , Logger.SetTruncated (isJust userQueryResults_continuation)
  , maybe mempty logQueryStats userQueryResults_stats
  , maybe mempty Logger.SetType userQueryResults_type
  , maybe mempty
      ( Logger.SetResponseContinuationSize
      . ByteString.length
      . Thrift.userQueryCont_continuation
      )
      userQueryResults_continuation
  ]

countQueryResults :: Thrift.UserQueryResults -> Int
countQueryResults Thrift.UserQueryResults{..} =
  case userQueryResults_results of
    Thrift.UserQueryEncodedResults_bin bin ->
      Map.size (Thrift.userQueryResultsBin_facts bin)
    Thrift.UserQueryEncodedResults_json json ->
      length (Thrift.userQueryResultsJSON_facts json)
    Thrift.UserQueryEncodedResults_compact compact ->
      length (Thrift.userQueryResultsCompact_facts compact)
    _ ->
      length userQueryResults_facts

logQueryStats :: Thrift.UserQueryStats -> GleanServerLog
logQueryStats Thrift.UserQueryStats{..} = mconcat
  [ Logger.SetResults (fromIntegral userQueryStats_result_count)
  , Logger.SetFacts (fromIntegral userQueryStats_num_facts)
  , Logger.SetFullScans (showRef <$> userQueryStats_full_scans)
  , maybe mempty (Logger.SetBytecodeSize . fromIntegral)
      userQueryStats_bytecode_size
  , maybe mempty (Logger.SetCompileTimeUs . fromIntegral . (`quot` 1000))
      userQueryStats_compile_time_ns
  , maybe mempty (Logger.SetExecuteTimeUs . fromIntegral . (`quot` 1000))
      userQueryStats_execute_time_ns
  , maybe mempty (Logger.SetQueryResultBytes . fromIntegral)
      userQueryStats_result_bytes
  ]

runLogDerivePredicate
  :: Text
  -> Database.Env
  -> Thrift.Repo
  -> Thrift.DerivePredicateQuery
  -> GleanServerLog
  -> IO ()
runLogDerivePredicate cmd env repo Thrift.DerivePredicateQuery {..} log =
  runLogRepo cmd env repo $ mconcat
    [ log
    , Logger.SetPredicate derivePredicateQuery_predicate
    , maybe mempty (Logger.SetPredicateVersion . fromIntegral)
        derivePredicateQuery_predicate_version
    , maybe mempty logQueryClientInfo derivePredicateQuery_client_info
    ]

runLogDerivationResult
  :: Database.Env
  -> LogDerivationResult
  -> Thrift.Repo
  -> Thrift.DerivePredicateQuery
  -> Either (DiffTimePoints, SomeException) Thrift.UserQueryStats
  -> IO ()
runLogDerivationResult env log repo Thrift.DerivePredicateQuery{..} res = do
  log res
  runLogRepo "deriveStored(completed)" env repo $ mconcat
    [ Logger.SetPredicate derivePredicateQuery_predicate
    , maybe mempty (Logger.SetPredicateVersion . fromIntegral)
        derivePredicateQuery_predicate_version
    , maybe mempty logQueryClientInfo derivePredicateQuery_client_info
    , case res of
        Left (_,err) -> failureLog err
        Right stats -> successLog <> logQueryStats stats
    , timeLog $ toDiffSeconds $ case res of
        Left (duration, _) -> duration
        Right Thrift.UserQueryStats{..} ->
          nanoseconds (fromIntegral userQueryStats_elapsed_ns)
    ]