{-
  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.Storage
  ( Mode(..)
  , CreateSchema(..)
  , Storage(..)
  , DBVersion(..)
  , AxiomOwnership
  , WriteLock(..)
  , canOpenVersion
  , currentVersion
  , writableVersions
  ) where

import Data.ByteString (ByteString)
import Data.HashMap.Strict (HashMap)
import qualified Data.Vector.Storable as VS

import Glean.Database.Backup.Backend (Data)
import Glean.Internal.Types (StoredSchema)
import Glean.RTS.Foreign.FactSet (FactSet)
import Glean.RTS.Foreign.Inventory (Inventory)
import Glean.RTS.Foreign.Lookup (CanLookup, Lookup)
import Glean.RTS.Foreign.Ownership
import Glean.RTS.Types (Fid, Pid)
import Glean.ServerConfig.Types (DBVersion(..))
import Glean.Types (PredicateStats, Repo, SchemaId)

-- | List of binary representation versions we can read
readableVersions :: [DBVersion]
readableVersions = [DBVersion 3]

-- | List of binary representation versions we can write
writableVersions :: [DBVersion]
writableVersions = [DBVersion 3]

-- | Check whether we can open a particular database version
canOpenVersion :: Mode -> DBVersion -> Bool
canOpenVersion mode version = version `elem` versions
  where
    versions = case mode of
      ReadOnly -> readableVersions
      ReadWrite -> writableVersions
      Create{} -> writableVersions

-- | Default current binary representation version
currentVersion :: DBVersion
currentVersion = maximum writableVersions

-- Choose which schema goes into a newly created DB
data CreateSchema
  = UseDefaultSchema
  | UseSpecificSchema SchemaId
  | UseThisSchema StoredSchema
  deriving (Show)

-- | Database opening mode
data Mode
  = ReadOnly
  | ReadWrite
  | Create
      Fid  -- starting fact id
      (Maybe Ownership)  -- base DB ownership
      CreateSchema

-- | Raw ownership data for axiomatic (non-derived) facts: a mapping
-- from unit name to ranges of fact IDs.
type AxiomOwnership = HashMap ByteString (VS.Vector Fid)

-- | Token representing the write lock
data WriteLock w = WriteLock w

-- | An abstract storage for fact database
class CanLookup (Database s) => Storage s where
  -- | A fact database
  data Database s

  -- | A short, user-readable description of the storage
  describe :: s -> String

  -- | Open a database
  open :: s -> Repo -> Mode -> DBVersion -> IO (Database s)

  -- | Close a database
  close :: Database s -> IO ()

  -- | Delete a database if it exists
  delete :: s -> Repo -> IO ()

  -- | Unconditionally remove a database or anything that might be stored where
  -- the database would exist. For disk-based storage, this would remove the
  -- directory where the database would be stored.
  safeRemoveForcibly :: s -> Repo -> IO ()

  -- | Obtain the 'PredicateStats' for each predicate
  predicateStats :: Database s -> IO [(Pid, PredicateStats)]

  -- | Store an arbitrary binary key/value pair in the database. This data is
  -- completely separate from the facts.
  --
  -- NOTE: It is expected that 'store' and 'retrieve' are used sparingly and
  -- there are no performance guarantees. A typical use case for this is
  -- storing the serialised schema in the database.
  store :: Database s -> ByteString -> ByteString -> IO ()

  -- | Retrieve the value of a key previously stored with 'store'.
  retrieve :: Database s -> ByteString -> IO (Maybe ByteString)

  -- | Commit a set of facts to the database. The facts must have the right ids,
  -- they are NOT renamed.
  commit :: Database s -> FactSet -> IO ()

  -- | Add ownership data about a set of (committed) facts.
  addOwnership :: Database s -> WriteLock w -> AxiomOwnership -> IO ()

  -- | Optimise a database for reading. This is typically done before backup.
  optimize :: Database s -> Bool {- compact -} -> IO ()

  computeOwnership
    :: Database s
    -> Maybe Lookup
       -- ^ Base DB lookup if this is a stacked DB, because ownership may
       -- need to propagate ownership through facts in the base DB.
    -> Inventory
    -> IO ComputedOwnership

  storeOwnership :: Database s -> WriteLock w -> ComputedOwnership -> IO ()

  -- | Fetch the 'Ownership' interface for this DB. This is used to
  -- make a 'Slice' (a view of a subset of the facts in the DB).
  --
  -- Can return 'Nothing' if this database backend doesn't support
  -- ownership. (TODO: support ownership in the memory backend and
  -- remove this 'Maybe').
  getOwnership :: Database s -> IO (Maybe Ownership)

  getUnitId :: Database s -> ByteString -> IO (Maybe UnitId)
  getUnit :: Database s -> UnitId -> IO (Maybe ByteString)

  -- | Called once per batch.
  addDefineOwnership :: Database s -> WriteLock w -> DefineOwnership -> IO ()

  -- | Called once per derived predicate at the end of its derivation.
  computeDerivedOwnership
    :: Database s
    -> WriteLock w
    -> Ownership
    -> Maybe Lookup
       -- ^ Base DB lookup if this is a stacked DB, because we may
       -- derive facts that already exist in the base DB and the
       -- ownership of those facts will need to be extended.
    -> Pid
    -> IO ComputedOwnership

  -- | After writing has finished, cache ownership data to support
  -- faster getOwner() operations. Takes time to cache the data and
  -- memory to retain the cache. Only useful if this DB will be used
  -- in an incremental stack.
  cacheOwnership :: Database s -> IO ()

  prepareFactOwnerCache :: Database s -> IO ()

  -- | Determine the total capacity of the storage medium (e.g., disk size).
  getTotalCapacity :: s -> IO (Maybe Int)

  -- | Determine the used capacity of the storage medium (e.g., how much of the
  -- disk is in use).
  getUsedCapacity :: s -> IO (Maybe Int)

  -- | Determine the free capacity of the storage medium (e.g., how much of the
  -- disk is free).
  getFreeCapacity :: s -> IO Int

  -- | Execute the action, passing to it a path to a scratch directory which can
  -- be used, e.g., for downloading databases. This directory is not guaranteed
  -- to persist beyond the call and is not guaranteed to be empty.
  withScratchRoot :: s -> (FilePath -> IO a) -> IO a

  -- | Backup a database. The scratch directory which can be used for storing
  -- intermediate files is guaranteed to be empty and will be deleted after
  -- the operation completes.
  backup
    :: Database s  -- ^ database
    -> FilePath  -- ^ scratch directory
    -> (FilePath -> Data -> IO a)
          -- ^ function which expects the serialised database
    -> IO a

  -- | Restore a database. The scratch directory which can be used for
  -- storing intermediate files is guaranteed to be empty and will be
  -- deleted after the operation completes. The implementation may
  -- delete the serialized database file after it has been consumed,
  -- to reduce the number of copies of the DB on disk during a restore.
  restore
    :: s   -- ^ storage
    -> Repo  -- ^ repo
    -> FilePath  -- ^ scratch directory
    -> FilePath  -- ^ file containing the serialiased database (produced by 'backup')
    -> IO ()