{-
  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.Database.Test
  ( Setting
  , setRoot
  , setSchemaLocation
  , setSchemaIndex
  , setSchemaPath
  , setSchemaId
  , disableStrictSchemaId
  , setMemoryStorage
  , setDBVersion
  , setCompactOnCompletion
  , setMaxSetSize
  , enableTcDebug
  , enableQueryDebug
  , withTestEnv
  , kickOffTestDB
  , waitUntilComplete
  , completeTestDB
  , withEmptyTestDB
  , writeFactsIntoDB
  ) where

import Util.STM
import Data.Default
import Data.Functor
import Data.Int
import Data.List (foldl')
import qualified Data.Text as Text

import Util.EventBase

import Glean.Backend.Local ( {- instance Backend Env -} )
import Glean.Backend.Types
import qualified Glean.Database.Catalog as Catalog
import Glean.Database.Config
import Glean.Database.Env
import Glean.Database.Write.Batch
import Glean.Database.Types
import qualified Glean.Internal.Types as Thrift
import qualified Glean.ServerConfig.Types as ServerConfig
import Glean.Typed
import qualified Glean.Types as Thrift
import Glean.Util.ConfigProvider
import qualified Glean.Util.ThriftSource as ThriftSource

type Setting = Config -> Config

setRoot :: FilePath -> Setting
setRoot path cfg = cfg{ cfgDataStore = fileDataStore path }

setSchemaLocation :: SchemaLocation -> Setting
setSchemaLocation loc cfg = cfg{ cfgSchemaLocation = Just loc }

setSchemaPath :: FilePath -> Setting
setSchemaPath = setSchemaLocation . SchemaLocation_file . Text.pack

setSchemaIndex :: FilePath -> Setting
setSchemaIndex = setSchemaLocation . SchemaLocation_index . Text.pack

-- | Set the schema that will be used for queries, otherwise defaults
-- to the highest all.N in the latest schema.
setSchemaId :: Thrift.SchemaId -> Setting
setSchemaId id cfg = cfg { cfgSchemaId = Just id }

disableStrictSchemaId :: Setting
disableStrictSchemaId cfg = cfg {
  cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
    { ServerConfig.config_strict_query_schema_id = False }
  }

setMemoryStorage :: Setting
setMemoryStorage cfg = cfg{ cfgDataStore = memoryDataStore }

setDBVersion :: ServerConfig.DBVersion -> Setting
setDBVersion ver cfg = cfg
  { cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
      { ServerConfig.config_db_create_version = Just ver } }

setCompactOnCompletion :: Setting
setCompactOnCompletion cfg = cfg
  { cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
      { ServerConfig.config_compact_on_completion = True } }

setMaxSetSize :: Int64 -> Setting
setMaxSetSize i cfg = cfg
  { cfgServerConfig = cfgServerConfig cfg <&> \scfg -> scfg
      { ServerConfig.config_max_set_size_bytes = Just i } }

enableTcDebug :: Setting
enableTcDebug cfg = cfg
  { cfgDebug = (cfgDebug cfg) { tcDebug = True } }

enableQueryDebug :: Setting
enableQueryDebug cfg = cfg
  { cfgDebug = (cfgDebug cfg) { queryDebug = True } }

withTestEnv
  :: [Setting]
  -> (Env -> IO a)
  -> IO a
withTestEnv settings action =
  withEventBaseDataplane $ \evb ->
  withConfigProvider defaultConfigOptions $
      \(cfgAPI :: NullConfigProvider) -> do
    let
      dbConfig = foldl' (\acc f -> f acc)
        def
          { cfgDataStore = tmpDataStore
          , cfgSchemaLocation = Just schemaLocationFiles
          , cfgServerConfig = ThriftSource.value def
              { ServerConfig.config_db_rocksdb_cache_mb = 0 }
          }
        settings

    withDatabases evb dbConfig cfgAPI action

kickOffTestDB
  :: Env -> Thrift.Repo -> (Thrift.KickOff -> Thrift.KickOff) -> IO ()
kickOffTestDB env repo update = do
  void $ kickOffDatabase env $ update def
    { Thrift.kickOff_repo = repo
    }

writeFactsIntoDB
  :: Env
  -> Thrift.Repo
  -> [SchemaPredicates]
  -> (forall m. NewFact m => m ())
  -> IO ()
writeFactsIntoDB env repo allPredicates facts = do
  !predicates <- loadPredicates env repo allPredicates
  batch <- buildBatch predicates Nothing facts
  void $ syncWriteDatabase env repo batch

waitUntilComplete :: Env -> Thrift.Repo -> IO ()
waitUntilComplete Env{..} repo = atomically $ do
  meta <- Catalog.readMeta envCatalog repo
  case Thrift.metaCompleteness meta of
    Thrift.Complete{} -> return ()
    Thrift.Broken{} -> errorWithoutStackTrace "database failed unexpectedly"
    _ -> retry

completeTestDB :: Env -> Thrift.Repo -> IO ()
completeTestDB env repo = do
  void $ finishDatabase env repo
  waitUntilComplete env repo

withEmptyTestDB
  :: [Setting]
  -> (Env -> Thrift.Repo -> IO a) -- ^ action
  -> IO a
withEmptyTestDB settings action = withTestEnv settings $ \env -> do
  kickOffTestDB env repo id
  action env repo
  where
    repo = Thrift.Repo "dbtest-repo" "f00baa"