{-
  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 ApplicativeDo #-}
{-# OPTIONS_GHC -Wno-orphans #-}
--
-- | An abstraction over the Glean API that can be served by
-- either a local DB or a remote server.
--
module Glean.Backend.Local
  (
    Env(..)
  , Config(..)
  , options
  , withDatabases

    -- * Schemas
  , loadDbSchema
  , serializeInventory
  ) where

import Control.Applicative
import Control.Concurrent (getNumCapabilities)
import Control.Concurrent.Stream (stream)
import Control.Exception (catches, Handler (Handler), throwIO)
import Control.Monad (forM_)
import Data.ByteString (ByteString)
import Data.Default
import Data.IORef.Extra
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Typeable
import qualified Haxl.Core as Haxl

import Util.Control.Exception
import qualified Util.Control.Exception.CallStack as CallStack
import Util.STM

import Glean.Backend.Types
import qualified Glean.Database.Catalog as Catalog
import qualified Glean.Database.CompletePredicates as Database
import qualified Glean.Database.Create as Database
import Glean.Database.Config
import qualified Glean.Database.Delete as Database
import Glean.Database.Env
import Glean.Database.Open as Database
import qualified Glean.Database.List as Database
import qualified Glean.Database.PredicateStats as Database (predicateStats)
import qualified Glean.Database.Restore as Database
import qualified Glean.Database.Schema as Database
import Glean.Database.Schema hiding (getSchemaInfo, schemaId)
import qualified Glean.Database.Types as Database
import Glean.Database.Types (Env(..))
import qualified Glean.Database.Finish as Database
import qualified Glean.Database.Writes as Database
import Glean.Internal.Types
import qualified Glean.Query.UserQuery as UserQuery
import qualified Glean.Query.Derive as Derive
import Glean.Query.Thrift
import Glean.Query.Thrift.Internal
import Glean.RTS (Fid(..))
import qualified Glean.RTS.Foreign.Inventory as Inventory
import qualified Glean.RTS.Foreign.Lookup as Lookup
import qualified Glean.Types as Thrift
import Glean.Util.Observed as Observed
import Haxl.DataSource.Glean.Common


instance Backend Database.Env where
  queryFact env repo id = readDatabase env repo $ \_ db ->
    Lookup.lookupFact db (Fid id)

  factIdRange env repo = do
    (starting, next) <- readDatabase env repo $ \_ db ->
      (,) <$> Lookup.startingId db <*> Lookup.firstFreeId db
    return $ Thrift.FactIdRange (fromFid starting) (fromFid next)

  getSchemaInfo env (Just repo) req =
    withOpenDatabase env repo $ \odb -> do
      index <- Observed.get (Database.envSchemaSource env)
      Database.getSchemaInfo (Database.odbSchema odb) index req

  getSchemaInfo env Nothing Thrift.GetSchemaInfo{..} = do
      index <- Observed.get (Database.envSchemaSource env)
      sid <- case getSchemaInfo_select of
          Thrift.SelectSchema_schema_id sid -> return sid
          other -> throwIO $ userError $ "unsupported: " <> show other
      Database.getSchemaInfoForSchema index sid (envDebug env)

  validateSchema env (Thrift.ValidateSchema str) = do
    schema <- Observed.get (Database.envSchemaSource env)
    conf <- Observed.get (Database.envServerConfig env)
    validateNewSchema conf str schema (envDebug env)

  predicateStats env repo opts = Database.predicateStats env repo opts

  userQueryFacts env repo q
    | isNothing (Thrift.userQueryFacts_schema_id q) =
      UserQuery.userQueryFacts env repo q {
        Thrift.userQueryFacts_schema_id = schemaId env
      }
    | otherwise = UserQuery.userQueryFacts env repo q

  userQuery env repo q
    | isNothing (Thrift.userQuery_schema_id q) =
      UserQuery.userQuery env repo q {
        Thrift.userQuery_schema_id = schemaId env
      }
    | otherwise = UserQuery.userQuery env repo q

  userQueryBatch env repo Thrift.UserQueryBatch{..} = do
      resultsRef <- newIORef mempty
      numCaps <- getNumCapabilities
      stream numCaps (forM_ queries) $ \q -> do
        res <- runOne q
        atomicModifyIORef'_ resultsRef $ Map.insert q res
      results <- readIORef resultsRef
      return $ map (results Map.!) queries
      where
        queries =
          [ Thrift.UserQuery
            { userQuery_predicate = userQueryBatch_predicate
            , userQuery_predicate_version = userQueryBatch_predicate_version
            , userQuery_encodings  = userQueryBatch_encodings
            , userQuery_client_info = userQueryBatch_client_info
            , userQuery_schema_id = userQueryBatch_schema_id <|> schemaId env
            , userQuery_options = userQueryBatch_options
            , userQuery_query = q
            }
          | q <- userQueryBatch_queries
          ]
        runOne query =
          (Thrift.UserQueryResultsOrException_results
            <$> userQuery env repo query)
          `catches`
          [ Handler $ \(Thrift.Retry r) -> return $
              Thrift.UserQueryResultsOrException_retry $
                Thrift.RetryException r
          , Handler $ \(Thrift.BadQuery b) -> return $
              Thrift.UserQueryResultsOrException_badQuery $
                Thrift.BadQueryException b
          ]


  deriveStored = Derive.deriveStored

  listDatabases = Database.listDatabases
  getDatabase env repo =
    maybe (CallStack.throwIO $ Thrift.UnknownDatabase repo) return
    =<< atomically (Catalog.getLocalDatabase (Database.envCatalog env) repo)

  kickOffDatabase = Database.kickOffDatabase
  finishDatabase = Database.finishDatabase
  finalizeDatabase = Database.finalizeDatabase
  updateProperties env repo set unset = do
    Database.updateProperties env repo set unset
    return def

  completePredicates_ = Database.completePredicates

  restoreDatabase = Database.restoreDatabase

  deleteDatabase env repo = do
    Database.deleteDatabase env repo
    return def

  enqueueBatch env cbatch = Database.enqueueBatch env cbatch Nothing
  enqueueJsonBatch env cbatch = Database.enqueueJsonBatch env cbatch
  enqueueBatchDescriptor env repo batch waitPolicy =
    Database.enqueueBatchDescriptor env repo batch waitPolicy
  pollBatch env handle = Database.pollBatch env handle

  displayBackend _ = "(local backend)"

  hasDatabase env repo = do
    Thrift.GetDatabaseResult { getDatabaseResult_database = Thrift.Database{..}}
      <- getDatabase env repo
    return $ case database_status of
      Thrift.DatabaseStatus_Restorable -> False
      _ -> True

  schemaId = Database.envSchemaId

  usingShards _ = False

  initGlobalState backend = do
    capabilities <- getNumCapabilities
    let streamMapM_ :: (a -> IO ()) -> [a] -> IO ()
        streamMapM_ f xx = stream capabilities (forM_ xx) f
    if capabilities == 1
      then return
        ( GleanGetState $ Haxl.SyncFetch $
            mapM_ (syncGetOne backend) . HashMap.toList . requestByRepo
        , GleanQueryState $ Haxl.SyncFetch $
            mapM_ (syncQueryOne backend)
        )
      else return
        ( GleanGetState $ Haxl.BackgroundFetch $
            streamMapM_ (syncGetOne backend) . HashMap.toList . requestByRepo
        , GleanQueryState $ Haxl.BackgroundFetch $
            streamMapM_ (syncQueryOne backend)
        )


-- -----------------------------------------------------------------------------
-- Haxl

syncGetOne :: Env -> (Thrift.Repo, [Haxl.BlockedFetch GleanGet]) -> IO ()
syncGetOne env (repo, requests) = do
  let schema = schemaId env
  results <- userQueryFacts env repo (mkUserQueryFacts Nothing schema requests)
  putResults results requests

syncQueryOne :: Env -> Haxl.BlockedFetch GleanQuery -> IO ()
syncQueryOne env
    (Haxl.BlockedFetch (QueryReq (Query q :: Query q) repo stream) rvar) =
  runSyncQuery repo env
    (Query (mkUserQuery Nothing (schemaId env) q) :: Query q)
    (if stream then Just mempty else Nothing) rvar

runSyncQuery
  :: forall q r. (Show q, Typeable q, QueryResult q r)
  => Thrift.Repo
  -> Env
  -> Query q
  -> Maybe r -- results so far
  -> Haxl.ResultVar (r, Bool)
  -> IO ()
runSyncQuery repo env q@(Query req) acc rvar = do
  r <- tryAll $ userQuery env repo req
  case r of
    Left e -> Haxl.putFailure rvar e
    Right results ->
      putQueryResults q results acc rvar $
        \(q :: Query q) acc -> runSyncQuery repo env q acc rvar

-- -----------------------------------------------------------------------------
-- DbSchema

loadDbSchema :: Backend a => a -> Thrift.Repo -> IO DbSchema
loadDbSchema backend repo = do
  Thrift.SchemaInfo schema pids _ dbSchemaIds _ _ <-
    getSchemaInfo backend (Just repo) def
      { Thrift.getSchemaInfo_select = Thrift.SelectSchema_stored def }
  fromStoredSchema Nothing (StoredSchema schema pids dbSchemaIds)
    readWriteContent def

-- | Serialize the inventory for the schema used by this repo.
-- The inventory is some information about the predicates in the schema
serializeInventory
  :: Backend backend
  => backend
  -> Thrift.Repo
  -> IO ByteString
serializeInventory backend repo = do
  dbSchema <- loadDbSchema backend repo
  return $ Inventory.serialize $ schemaInventory dbSchema