{-
  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-incomplete-uni-patterns #-}
module BackupTest (main) where

import Util.STM
import Control.Monad
import Data.Default
import qualified Data.HashSet as HashSet
import Data.Int (Int64)
import Data.List
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Time.Clock
import System.IO.Temp
import Test.HUnit

import TestRunner
import Util.EventBase

import Glean.Backend.Local ()
import Glean.Backend.Types
import Glean.Database.Backup (Event(..))
import qualified Glean.Database.Backup.Backend as Backup
import qualified Glean.Database.Backup.Mock as Backup.Mock
import Glean.Database.Catalog as Catalog
import Glean.Database.Config
import Glean.Database.Env
import Glean.Database.Janitor
import Glean.Database.Meta
import Glean.Database.Types
import Glean.Database.Finish (finalizeWait)
import Glean.Init
import Glean.ServerConfig.Types as ServerTypes
import Glean.Types as Thrift
import Glean.Util.ConfigProvider
import Glean.Util.ThriftSource as ThriftSource
import Glean.Util.Trace
import Glean.Util.IO (withTempFileContents)
import Glean.Internal.Types (Completeness(Complete))
import qualified Glean.Database.Storage as Storage

withTest
  :: (EventBaseDataplane -> NullConfigProvider -> FilePath -> IO ())
  -> IO ()
withTest action =
  withEventBaseDataplane $ \evb ->
  withConfigProvider defaultConfigOptions $ \cfgAPI ->
  withSystemTempDirectory "glean-dbtest-backup" $ \backupdir ->
    action evb cfgAPI backupdir

data TestEnv = forall site. Backup.Site site => TestEnv
  { testEnv :: Env
  , testBackup :: site
  , testUpdConfig :: (ServerTypes.Config -> ServerTypes.Config) -> IO ()
  , testEvents :: TQueue Event
  }

data TestDbSpec
  = TestDbSpec
    { testDbRepo :: Repo
    , testDbAge :: Int
    , testDbDeps :: Maybe Dependencies
    }
  | CloudTestDbSpec
    { testDbRepo :: Repo
    , testDbAge :: Int
    , testDbSize :: Int64
    , testDbDeps :: Maybe Dependencies
    }

goodDb :: Text -> Text -> TestDbSpec
goodDb name hash = TestDbSpec (Repo name hash) 0 Nothing

goodDbAge :: Text -> Text -> Int -> TestDbSpec
goodDbAge name hash age = TestDbSpec (Repo name hash) age Nothing

goodDbAgeDeps :: Text -> Text -> Int -> Dependencies -> TestDbSpec
goodDbAgeDeps name hash age deps =
  TestDbSpec (Repo name hash) age (Just deps)

withTestEnv
  :: [TestDbSpec]
  -> (ServerTypes.Config -> ServerTypes.Config)
  -> (TestEnv -> IO ())
  -> EventBaseDataplane
  -> NullConfigProvider
  -> FilePath
  -> IO ()
withTestEnv dbs init_server_cfg action evb cfgAPI backupdir = do
  (server_cfg, update_server_cfg) <- ThriftSource.mutable $ init_server_cfg def
    { config_backup = def
      { databaseBackupPolicy_allowed = mempty
      , databaseBackupPolicy_location = "mock:" <> Text.pack backupdir
      }
    , config_restore = def
      { databaseRestorePolicy_enabled = False
      }
    , config_janitor_period = Nothing -- no auto janitor, we'll run it manually
    , config_db_rocksdb_cache_mb = 0
    }
  (l, events) <- recorder
  let config = def
        { cfgDataStore = tmpDataStore
        , cfgSchemaLocation = Just schemaLocationFiles
        , cfgServerConfig = server_cfg
        , cfgReadOnly = False
        , cfgMockWrites = False
        , cfgListener = l
        }

  withDatabases evb config cfgAPI $ \env -> do
    now <- getCurrentTime
    let testEnv = TestEnv
          { testEnv = env
          , testBackup = Backup.Mock.mockSite backupdir
          , testUpdConfig = update_server_cfg
          , testEvents = events
          }
    mapM_ (makeDB testEnv now) dbs
    action testEnv

expectBackups :: [Repo] -> Expect Event
expectBackups repos = parallel
  [ opt (want Waiting) <> wants [ BackupStarted repo, BackupFinished repo ]
  | repo <- repos ]

expectRestore :: [Repo] -> Expect Event
expectRestore repos = parallel
  [ opt (want Waiting) <> wants [ RestoreStarted repo, RestoreFinished repo ]
  | repo <- repos ]

expectFinalize :: [Repo] -> Expect Event
expectFinalize repos = parallel
  [ opt (want Waiting) <> wants [ FinalizeStarted repo, FinalizeFinished repo ]
  | repo <- repos ]

created :: Integral a => a -> UTCTime -> UTCTime
created t = addUTCTime (negate (fromIntegral t))

makeDB :: TestEnv -> UTCTime -> TestDbSpec -> IO ()
makeDB TestEnv{testEnv = env} now (TestDbSpec repo age deps) = do
  KickOffResponse False <- kickOffDatabase env def
    { kickOff_repo = repo
    , kickOff_dependencies = deps
    }
  void $ atomically $ Catalog.modifyMeta (envCatalog env) repo $ \meta ->
    return meta { metaCreated = utcTimeToPosixEpochTime (created age now) }
  void $ finishDatabase env repo
  finalizeWait env repo

makeDB TestEnv{testBackup} now CloudTestDbSpec{..} =
  withTempFileContents ("" :: String) $ \path ->
  void $ Backup.backup testBackup testDbRepo props Nothing path
  where
    dbtime = utcTimeToPosixEpochTime (created testDbAge now)
    props = Meta
        { metaVersion = Storage.currentVersion
        , metaCreated = dbtime
        , metaRepoHashTime = Nothing
        , metaCompleteness =
          Complete $ DatabaseComplete dbtime (Just testDbSize)
        , metaBackup = Nothing
        , metaProperties = mempty
        , metaDependencies = testDbDeps
        , metaCompletePredicates = mempty
        , metaAxiomComplete = False
        }

basicBackupTest :: Test
basicBackupTest = TestCase $ withTest $ withTestEnv
  repos
  id
  $ \TestEnv{..} -> do
    expect testEvents $ mconcat
      [ expectFinalize [repo | TestDbSpec repo _ _ <- repos] ]
    testUpdConfig $ \scfg -> scfg
      { config_backup = (config_backup scfg)
          { databaseBackupPolicy_allowed = HashSet.fromList ["test"] } }
    expect testEvents $ mconcat
      [ expectBackups [repo | TestDbSpec repo _ _ <- repos]
      , want Waiting ]

    backups <- Backup.enumerate testBackup
    assertEqual "repos"
      (sort [repo | TestDbSpec repo _ _ <- repos])
      (sort [repo | (repo, _) <- backups])
  where
    repos :: [TestDbSpec]
    repos =
      [ goodDb "test" "1"
      , goodDb "test" "2"
      , goodDb "test" "3" ]

allowedTest :: Test
allowedTest = TestCase $ withTest $ withTestEnv
  repos
  id
  $ \TestEnv{..} -> do
    expect testEvents $ mconcat
      [ expectFinalize [repo | TestDbSpec repo _ _ <- repos] ]
    testUpdConfig $ \scfg -> scfg
      { config_backup = (config_backup scfg)
          { databaseBackupPolicy_allowed = HashSet.fromList ["bar"] }}
    expect testEvents $ mconcat
      [ expectBackups [repo | TestDbSpec repo _ _ <- repos,
          repo_name repo == "bar"]
      , want Waiting ]
  where
    repos :: [TestDbSpec]
    repos =
      [ goodDb "foo" "1"
      , goodDb "bar" "2"
      , goodDb "foo" "3"
      , goodDb "baz" "4"
      , goodDb "bar" "5" ]

restoreOrderTest :: Test
restoreOrderTest = TestCase $ withTest $ withTestEnv
  repos
  id
  $ \TestEnv{..} -> do
    expect testEvents $ mconcat
      [ expectFinalize [repo | TestDbSpec repo _ _ <- repos] ]
    testUpdConfig $ \scfg -> scfg
      { config_backup = (config_backup scfg)
          { databaseBackupPolicy_allowed =
              HashSet.fromList ["bar","foo","baz"] }}
    expect testEvents $ mconcat
      [ expectBackups [repo | TestDbSpec repo _ _ <- repos ]
      , want Waiting ]

    -- delete all the repos and run the janitor to start restoring
    forM_ repos $ \(TestDbSpec repo _ _) -> deleteDatabase testEnv repo
    testUpdConfig $ \scfg -> scfg
      { config_restore = def { databaseRestorePolicy_enabled = True } }
    runDatabaseJanitor testEnv

    -- restoring picks the newest DB of each repo to restore, and then
    -- picks the first of those ordered by staleness, and then most
    -- recent.

    -- local: []
    -- newest restorable for each repo:
    --    Repo "foo" "3", True, 1
    --    Repo "bar" "5", True, 0
    -- all are equally stale, so we pick the newest, bar/5
    expect testEvents $ expectRestore [ Repo "bar" "5" ]

    -- local: [bar/5]
    -- newest restorable for each repo:
    --    Repo "foo" "3", True, 1
    --    Repo "baz" "4", True, 3
    --    Repo "bar" "2", True, 1
    -- foo, baz are stale, foo/3 is the newest
    expect testEvents $ expectRestore [ Repo "foo" "3" ]

    -- local: [bar/5, foo/3]
    -- newest restorable for each repo:
    --    Repo "foo" "1", True, 2
    --    Repo "baz" "4", True, 3
    --    Repo "bar" "2", True, 1
    -- baz is stale
    expect testEvents $ expectRestore [ Repo "baz" "4" ]

    -- local: [bar/5, foo/3, baz/4]
    -- newest restorable for each repo:
    --    Repo "foo" "1", True, 2
    --    Repo "bar" "2", True, 1
    -- none stale, bar/2 is the newest, but it depends on foo/1, so
    -- we should restore that first.
    expect testEvents $ expectRestore [ Repo "foo" "1" ]
    expect testEvents $ expectRestore [ Repo "bar" "2" ]
  where
    repos :: [TestDbSpec]
    repos =
      [ goodDbAge "foo" "1" 2
      , goodDbAge "foo" "3" 1
      , goodDbAgeDeps "bar" "2" 1 $
          Dependencies_stacked (Stacked "foo" "1" Nothing)
      , goodDbAge "bar" "5" 0
      , goodDbAge "baz" "4" 3
      ]

restoreNoDiskSpaceTest :: Test
restoreNoDiskSpaceTest =
  TestCase $ withTest $ withTestEnv repos id $ \TestEnv{..} -> do
    testUpdConfig $ \scfg -> scfg
      { config_restore = def { databaseRestorePolicy_enabled = True }
      }
    runDatabaseJanitor testEnv
    expect testEvents $
       opt (want Waiting) <> wants [ RestoreStarted repo, RestoreFailed repo ]

  where
    repo = Repo "foo" "1"
    repos = [ CloudTestDbSpec repo 1 (10*tb) Nothing]
    tb :: Int64
    tb = 1000000000000

main :: IO ()
main = withUnitTest $ testRunner $ TestList
  [ TestLabel "basicBackup" basicBackupTest
  , TestLabel "allowedTest" allowedTest
  , TestLabel "restoreOrderTest" restoreOrderTest
  , TestLabel "restoreNoDiskSpace" restoreNoDiskSpaceTest
  ]