{-
  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 DeriveTraversable #-}
{- | A Catalog stores metadata about available databases.
A thread running Glean.Database.Backup.backuper will continuously read the
state of the catalog to decide which dbs to restore/backup/finalize
-}

module Glean.Database.Catalog
  ( Catalog
  , Entry
  , CatalogClosedError(..)
  , EntryAlreadyExists(..)
  , open
  , close
  , create
  , delete
  , list
  , listMostRecent
  , exists
  , readMeta
  , writeMeta
  , modifyMeta
  , readExpiring
  , setExpiring
  , unsetExpiring
  , startRestoring
  , finishRestoring
  , abortRestoring
  , resetElsewhere
  , getLocalDatabases
  , getLocalDatabase
  , dbFailed
  , resetFailed
  -- for testing
  , getEntries
  , list'
  , Entries
  , EntriesF(..)
  , EntryF(..)
  , mkEntry
  ) where

import Control.Concurrent.Async
import Control.Exception
    ( Exception,
      SomeException,
      bracket,
      bracket_,
      try,
      mask_,
      uninterruptibleMask_ )
import Control.Monad
import Control.Monad.Catch (throwM)
import Data.Functor.Identity
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import Data.Maybe
import qualified Data.Text as Text
import Data.Time
import Data.Typeable
import GHC.Stack (HasCallStack)
import System.IO (fixIO)

import Util.Control.Exception (tryBracket)
import qualified Util.Control.Exception.CallStack as CallStack
import Util.Log
import Util.STM

import Glean.Database.Catalog.Filter
import Glean.Database.Catalog.Store (Store)
import qualified Glean.Database.Catalog.Store as Store
import Glean.Database.Exception
import Glean.Database.Meta
import Glean.Database.Repo
import Glean.Types (Repo(..))
import qualified Glean.Internal.Types as Thrift
import qualified Glean.Types as Thrift

-- | Catalog entry
type Entry = EntryF TVar

-- | Catalog entry, parameterized by the state type;
--   This extra generality comes handy in tests.
data EntryF ref = Entry
  { entryRepo :: Repo
  , entryStatus :: ref ItemStatus
      -- ^ This is the combined status of this db and the entryStatus of its
      -- dependencies. Sent with restorable status as 'Thrift.DatabaseStatus'.
  , entryMeta :: ref Meta
  , entryDirty :: ref Bool
      -- ^ does the metadata need to be written back (Entry must be in
      -- 'catDirtyQueue' in that case)
  , entryComitting :: ref Bool
      -- ^ is the metadata being committed (Entry can't be deleted while this
      -- is happening)
  , entryExpiring :: ref (Maybe UTCTime)
  }

deriving instance Show (EntryF Identity)
-- deriving instance Eq (EntryF Identity)

-- | All databases known to the 'Catalog'.
type Entries = EntriesF Entry

-- | The type of all databases in the 'Catalog', parameterized by the entry type
--   This extra generality comes handy in tests.
data EntriesF entry = Entries
  { entriesLiveHere :: HashMap Repo entry
      -- ^ available databases
  , entriesLiveElsewhere :: HashMap Repo entry
      -- ^ available databases in another shard
      --   ignored for all purposes except by getLocalDatabases
  , entriesRestoring :: HashMap Repo Meta
      -- ^ databases that are being restored
  , entriesEphemeral :: HashSet Repo
      -- ^ databases that are being created or deleted
  , entriesFailed :: HashMap Repo SomeException
      -- ^ databases that failed to open, and their exceptions
  }
  deriving (Show, Functor, Foldable, Traversable)


-- | Catalog of database entries
data Catalog = forall local. Store local => Catalog
  { catLocal :: local
      -- ^ local metadata storage
  , catEntries :: TVar (Maybe Entries)
      -- ^ all known entries ('Nothing' means catalog is closed)
  , catDirtyQueue :: TQueue Entry
      -- ^ queue of entries which need to have their metadata committed
  , catCommitter :: Async ()
      -- ^ metadata committer thread
  , catRepoDependents :: TVar (HashMap Repo [Repo])
      -- ^ keep track of repos that depend on a repo
  }

data CatalogClosedError = CatalogClosedError
  deriving(Show,Typeable)

instance Exception CatalogClosedError

newtype EntryAlreadyExists = EntryAlreadyExists Repo
  deriving(Eq,Show,Typeable)

instance Exception EntryAlreadyExists

mkEntry :: Repo -> Meta -> STM Entry
mkEntry repo meta = Entry repo
  <$> newTVar ItemMissing
  <*> newTVar meta
  <*> newTVar False
  <*> newTVar False
  <*> newTVar Nothing

-- | Updates entryStatus. Needs be called whenever the status of this db or one
-- of its dependencies changes
recalculateStatus :: Catalog -> Entry -> STM ()
recalculateStatus Catalog{..} entry = do
  meta <- readTVar $ entryMeta entry
  entries <- readTVar catEntries
  case entries of
    Nothing -> return ()
    Just Entries{..} -> do
      let
        repo = entryRepo entry

        missingStatus repo | repo `HashMap.member` entriesRestoring =
          ItemRestoring
        missingStatus _ = ItemMissing

        dependencies = case metaDependencies meta of
          Just (Thrift.Dependencies_stacked Thrift.Stacked{..}) ->
            [(Repo stacked_name stacked_hash, stacked_guid)]
          Just (Thrift.Dependencies_pruned up) ->
            [(Thrift.pruned_base up, Thrift.pruned_guid up)]
          Nothing -> []
        live = repo `HashMap.member` entriesLiveHere
      forM_ dependencies $ \(dep, _) -> if live then
          modifyTVar' catRepoDependents $
            HashMap.insertWith (<>) dep [repo]
        else
          modifyTVar' catRepoDependents $
            HashMap.adjust (filter (repo /=)) dep

      let
        itemStatusFor :: Thrift.Completeness -> ItemStatus
        itemStatusFor Thrift.Complete{} = ItemComplete
        itemStatusFor Thrift.Incomplete{} = ItemIncomplete
        itemStatusFor Thrift.Broken{} = ItemBroken
        itemStatusFor Thrift.Finalizing{} = ItemFinalizing

      dependencyStatuses <- forM dependencies $ \(dep, guid) -> do
        let missing = missingStatus dep
        case HashMap.lookup dep entriesLiveHere of
          Nothing -> return missing
          Just entry -> do
            meta <- readTVar $ entryMeta entry
            status <- readTVar $ entryStatus entry
            let properties = Thrift.metaProperties meta
            if HashMap.lookup "glean.guid" properties == guid then
              return status
            else
              return missing

      meta <- readTVar $ entryMeta entry
      oldStatus <- readTVar $ entryStatus entry
      let
        status
          | Just{} <- HashMap.lookup repo entriesFailed = ItemBroken
          | live = itemStatusFor $ metaCompleteness meta
          | otherwise = missingStatus repo
        newStatus = mconcat $ status:dependencyStatuses

      when (newStatus /= oldStatus) $ do
        writeTVar (entryStatus entry) newStatus
        recalculateDepsStatus Catalog{..} repo

recalculateDepsStatus :: Catalog -> Repo -> STM ()
recalculateDepsStatus Catalog{..} repo = do
  repoDependents <- readTVar catRepoDependents
  entries <- readTVar catEntries
  forM_ (HashMap.lookupDefault [] repo repoDependents) $
    \dep -> forM_ (entries >>= HashMap.lookup dep . entriesLiveHere) $
      \entry -> recalculateStatus Catalog{..} entry

itemDatabaseStatus :: ItemStatus -> Thrift.DatabaseStatus
itemDatabaseStatus ItemComplete = Thrift.DatabaseStatus_Complete
itemDatabaseStatus ItemIncomplete = Thrift.DatabaseStatus_Incomplete
itemDatabaseStatus ItemRestoring = Thrift.DatabaseStatus_Restoring
itemDatabaseStatus ItemBroken = Thrift.DatabaseStatus_Broken
itemDatabaseStatus ItemMissing = Thrift.DatabaseStatus_Missing
itemDatabaseStatus ItemFinalizing = Thrift.DatabaseStatus_Finalizing

dirtyEntry :: Catalog -> Entry -> STM ()
dirtyEntry cat entry = do
  dirty <- swapTVar (entryDirty entry) True
  when (not dirty) $ writeTQueue (catDirtyQueue cat) entry

-- | Metadata committer thread. This thread should *never* be cancelled since it
-- needs to process outstanding commit requests before shutting down.
-- Accordingly, it runs in uninterruptibleMask_. Note that crashing in case of
-- heap/stack overflow is preferable to this thread dying.
commit :: Catalog -> IO ()
commit cat@Catalog{..} = uninterruptibleMask_ loop
  where
    loop = do
      continue <- bracket
        -- get Just the next entry or Nothing if the Catalog has been closed
        (atomically $
          do
            entry <- readTQueue catDirtyQueue
            writeTVar (entryComitting entry) True
            return $ Just entry
          `orElse`
          do
            open <- isJust <$> readTVar catEntries
            when open retry
            return Nothing)
        (\r -> forM_ r $ \entry ->
            atomically $ writeTVar (entryComitting entry) False)
        $ \r -> case r of
            Just entry@Entry{..} -> do
              meta <- atomically $ do
                writeTVar entryDirty False
                readTVar entryMeta
              -- catch *all* exceptions
              r <- try $ Store.put catLocal entryRepo meta
              case r of
                Right ok -> when (not ok) $
                  -- TODO: What do we do in this case?
                  logError $ inRepo entryRepo
                    "couldn't commit meta because entry no longer exists"
                Left (exc :: SomeException) -> do
                  -- TODO: this is lame - it's not clear what to do if we can't
                  -- write the metadata
                  logError $ inRepo entryRepo $
                    "couldn't commit meta: " ++ show exc
                  atomically $ dirtyEntry cat entry
              return True
            Nothing -> return False
      when continue loop

getEntries :: Catalog -> STM Entries
getEntries cat = do
  r <- readTVar $ catEntries cat
  case r of
    Just entries -> return entries
    Nothing -> throwSTM CatalogClosedError

lookupEntry :: Catalog -> Repo -> STM (Maybe Entry)
lookupEntry cat repo = HashMap.lookup repo . entriesLiveHere <$> getEntries cat

getEntry :: HasCallStack => Catalog -> Repo -> STM Entry
getEntry cat repo = do
  r <- lookupEntry cat repo
  case r of
    Just entry -> return entry
    Nothing -> CallStack.throwSTM $ Thrift.UnknownDatabase repo

-- | Open a 'Catalog'
open :: Store local => local -> IO Catalog
open local = do
  metas <- Store.list local
  live <- atomically $ HashMap.traverseWithKey mkEntry metas
  entries <- newTVarIO $ Just Entries
    { entriesLiveHere = live
    , entriesLiveElsewhere = mempty
    , entriesRestoring = mempty
    , entriesEphemeral = mempty
    , entriesFailed = mempty
    }
  dependents <- newTVarIO mempty
  dirty_queue <- newTQueueIO
  cat <- fixIO $ \cat -> do
    committer <- mask_ $ async $ commit cat
    return Catalog
      { catLocal = local
      , catEntries = entries
      , catDirtyQueue = dirty_queue
      , catCommitter = committer
      , catRepoDependents = dependents
      }
  forM_ (HashMap.elems live) $ \entry ->
    atomically $ recalculateStatus cat entry
  return cat

-- | Close a 'Catalog'
close :: Catalog -> IO ()
close cat = do
  atomically $ do
    Entries{..} <- getEntries cat
    when (not $ HashSet.null entriesEphemeral) retry
    writeTVar (catEntries cat) Nothing
  void $ waitCatch $ catCommitter cat

-- | Create a new entry in a 'Catalog' and execute the supplied action if
-- creation was successful.
create :: Catalog -> Repo -> Meta -> STM () -> IO ()
create cat@Catalog{..} repo meta on_success = tryBracket
  (atomically $ do
    Entries{..} <- getEntries cat
    when
      (repo `HashMap.member` entriesLiveHere
        || repo `HashSet.member` entriesEphemeral
        || repo `HashMap.member` entriesRestoring) $
      throwSTM $ EntryAlreadyExists repo
    writeTVar catEntries $ Just Entries
      { entriesEphemeral = HashSet.insert repo entriesEphemeral, .. })
  (\_ r -> atomically $ do
    entry <- mkEntry repo meta
    modifyTVar' catEntries $ fmap $ \Entries{..} -> Entries
      { entriesEphemeral = HashSet.delete repo entriesEphemeral
      , entriesLiveHere = case r of
          Right _ -> HashMap.insert repo entry entriesLiveHere
          Left _ -> entriesLiveHere
      , ..
      }
    recalculateStatus cat entry
    forM_ r $ const on_success)
  $ const $ do
    ok <- Store.create catLocal repo meta
    -- This is *not* an EntryAlreadyExists error
    when (not ok) $ dbError repo
      "database already exists in local metadata store"

-- | Permanently delete an entry from a 'Catalog'
delete :: Catalog -> Repo -> IO ()
delete cat@Catalog{..} repo = bracket_
  (atomically $ do
    entry <- getEntry cat repo
    committing <- readTVar (entryComitting entry)
    when committing retry
    dirty <- readTVar (entryDirty entry)
    when dirty $ do
      es <- flushTQueue catDirtyQueue
      mapM_ (writeTQueue catDirtyQueue)
        $ filter (\e -> entryRepo e /= repo) es
    modifyTVar' catEntries $ fmap $ \Entries{..} -> Entries
      { entriesLiveHere = HashMap.delete repo entriesLiveHere
      , entriesEphemeral = HashSet.insert repo entriesEphemeral
      , ..
      }
    recalculateStatus cat entry)
  (atomically $ modifyTVar' catEntries $ fmap $ \Entries{..} -> Entries
    { entriesEphemeral = HashSet.delete repo entriesEphemeral, .. })
  -- Ignore if it doesn't exist for now
  $ void $ Store.delete catLocal repo

-- | List all databases in the given localities that match the given filters
list :: Catalog -> [Locality] -> Filter () -> STM [Item]
list cat locs f = getEntries cat >>= list' readTVar locs f

list'
  :: Monad m
  => (forall a. f a -> m a)
  -> [Locality]
  -> Filter ()
  -> EntriesF (EntryF f)
  -> m [Item]
list' read locs f Entries{..} = do
  fmap (runFilter f . concat) $ forM locs $ \loc -> do
    xs <- case loc of
      Local -> mapM statusAndMeta entriesLiveHere
      Restoring -> return $ fmap (ItemRestoring,) entriesRestoring
      Cloud -> mapM statusAndMeta entriesLiveElsewhere
    return $
      [ Item repo loc meta status
      | (repo, (status, meta)) <- HashMap.toList xs]
  where
    -- statusAndMeta :: Entry -> f (ItemStatus, Meta)
    statusAndMeta Entry{..} = do
      status <- read entryStatus
      meta <- read entryMeta
      return (status, meta)

-- | List the most recent instances of all DBs in the Catalog
listMostRecent :: Catalog -> IO [Item]
listMostRecent catalog =  atomically $
  list catalog [Local, Restoring, Cloud] $ groupF repoNameV $ do
    sortF createdV Descending
    limitF 1

-- | Check if a database exists in the catalog
exists :: Catalog -> [Locality] -> Repo -> STM Bool
exists cat locs repo = do
  Entries{..} <- getEntries cat
  let exists_in Local = HashMap.member repo entriesLiveHere
      exists_in Restoring = HashMap.member repo entriesRestoring
      exists_in Cloud = False
  return $ any exists_in locs

-- | Read the metadata of a database
readMeta :: HasCallStack => Catalog -> Repo -> STM Meta
readMeta cat repo = do
  entry <- getEntry cat repo
  readTVar (entryMeta entry)

-- | Update the metadata of a database
writeMeta :: Catalog -> Repo -> Meta -> STM ()
writeMeta cat repo meta = do
  entry <- getEntry cat repo
  writeTVar (entryMeta entry) meta
  recalculateStatus cat entry
  dirtyEntry cat entry

-- | Modify the metadata of a database and return the new metadata
modifyMeta :: Catalog -> Repo -> (Meta -> STM Meta) -> STM Meta
modifyMeta cat repo f = do
  entry <- getEntry cat repo
  old_meta <- readTVar (entryMeta entry)
  new_meta <- f old_meta
  writeTVar (entryMeta entry) new_meta
  recalculateStatus cat entry
  dirtyEntry cat entry
  return new_meta

dbFailed :: Catalog -> Repo -> SomeException -> STM ()
dbFailed cat repo exception = do
  Entries{..} <- getEntries cat
  writeTVar (catEntries cat) $ Just Entries
    { entriesFailed = HashMap.insert repo exception entriesFailed
    , ..
    }
  entry <- getEntry cat repo
  recalculateStatus cat entry

resetFailed :: Catalog -> STM ()
resetFailed cat = do
  Entries{..} <- getEntries cat
  writeTVar (catEntries cat) $ Just Entries
    { entriesFailed = HashMap.empty
    , ..
    }

readExpiring :: Catalog -> Repo -> STM (Maybe UTCTime)
readExpiring cat repo = do
  entry <- getEntry cat repo
  readTVar $ entryExpiring entry

setExpiring :: Catalog -> Repo -> UTCTime -> STM ()
setExpiring cat repo time = do
  entry <- getEntry cat repo
  writeTVar (entryExpiring entry) $ Just time

-- | Unset the expiring bit, to no longer mark the 'Repo' as expiring
unsetExpiring :: Catalog -> Repo -> STM ()
unsetExpiring cat repo = do
  entry <- getEntry cat repo
  writeTVar (entryExpiring entry) Nothing

-- | Schedule a database for download/restore
startRestoring :: Catalog -> Repo -> Meta -> STM ()
startRestoring cat repo meta = do
  Entries{..} <- getEntries cat
  when
    (repo `HashMap.member` entriesLiveHere
      || repo `HashMap.member` entriesRestoring
      || repo `HashSet.member` entriesEphemeral)
    $ throwM DBAlreadyExists
  writeTVar (catEntries cat) $ Just Entries
    { entriesRestoring = HashMap.insert repo meta entriesRestoring
    , ..
    }
  recalculateDepsStatus cat repo

-- | Notify the catalog that the database has been restored and is available
-- locally
finishRestoring :: Catalog -> Repo -> IO ()
finishRestoring cat@Catalog{..} repo = tryBracket
  (atomically $ do
    Entries{..} <- getEntries cat
    case HashMap.lookup repo entriesRestoring of
      Just meta -> do
        writeTVar catEntries $ Just Entries
          { entriesRestoring = HashMap.delete repo entriesRestoring
          , entriesEphemeral = HashSet.insert repo entriesEphemeral
          , ..
          }
        return meta
      Nothing -> dbError repo "finishRestoring: unknown database")
  (\meta r -> atomically $ do
    entry <- mkEntry repo meta
    modifyTVar' catEntries $ fmap $ \Entries{..} -> Entries
      { entriesEphemeral = HashSet.delete repo entriesEphemeral
      , entriesLiveHere = case r of
          Right _ -> HashMap.insert repo entry entriesLiveHere
          Left _ -> entriesLiveHere
      , .. }
    recalculateStatus cat entry)
  $ \meta -> do
    ok <- Store.create catLocal repo meta
    when (not ok) $ dbError repo
      "finishRestoring: entry already exists in local metadata store"

-- | Notify the catalog that the database is no longer being restored
abortRestoring :: Catalog -> Repo -> STM ()
abortRestoring cat repo = do
  Entries{..} <- getEntries cat
  when (not $ repo `HashMap.member` entriesRestoring) $ dbError repo
    "abortRestoring: unknown database"
  writeTVar (catEntries cat) $ Just Entries
    { entriesRestoring = HashMap.delete repo entriesRestoring
    , .. }
  recalculateDepsStatus cat repo

-- | Reset the catalog of items available elsewhere to the ones given
resetElsewhere :: Catalog -> [Item] -> STM ()
resetElsewhere cat items = do
  Entries{..} <- getEntries cat
  entriesElsewhere <- forM items $ \Item{..} -> do
    e <- mkEntry itemRepo itemMeta
    return (itemRepo, e)
  writeTVar (catEntries cat) $ Just Entries
    { entriesLiveElsewhere =
        HashMap.fromList entriesElsewhere
    , ..
    }

getLocalDatabases :: Catalog -> STM (HashMap Repo Thrift.GetDatabaseResult)
getLocalDatabases cat = do
  Entries{..} <- getEntries cat
  local <- HashMap.traverseWithKey (local_db Nothing) entriesLiveHere
  elsewhere <- HashMap.traverseWithKey
    (local_db (Just Thrift.DatabaseStatus_Available))
    entriesLiveElsewhere
  return $ mconcat
    [ elsewhere
    , updateFailed entriesFailed local
    , HashMap.mapWithKey restoring_db entriesRestoring
    ]
  where
    updateFailed entriesFailed = HashMap.mapWithKey $
      \repo r@Thrift.GetDatabaseResult{..} ->
        if
          | Just exception <- HashMap.lookup repo entriesFailed ->
            Thrift.GetDatabaseResult {
              getDatabaseResult_database = getDatabaseResult_database {
                Thrift.database_broken = Just $ Thrift.DatabaseBroken
                  { databaseBroken_task = "open"
                  , databaseBroken_reason = Text.pack (show exception)
                  }
                },
              ..
            }
          | otherwise -> r

    local_db overrideStatus repo entry = do
      meta <- readTVar $ entryMeta entry
      exp <- readTVar $ entryExpiring entry
      status <- readTVar $ entryStatus entry
      return Thrift.GetDatabaseResult
        { getDatabaseResult_database = metaToThriftDatabase
            (fromMaybe (itemDatabaseStatus status) overrideStatus)
            exp
            repo
            meta
        }

    restoring_db repo meta = Thrift.GetDatabaseResult
      { getDatabaseResult_database = metaToThriftDatabase
          Thrift.DatabaseStatus_Restoring
          Nothing
          repo
          meta
      }

getLocalDatabase :: Catalog -> Repo -> STM (Maybe Thrift.GetDatabaseResult)
-- TODO: This isn't very efficient
getLocalDatabase cat repo = HashMap.lookup repo <$> getLocalDatabases cat