{-
  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 AllowAmbiguousTypes, TypeApplications, CPP, InstanceSigs #-}
module Glean.Typed.BuildFact
  ( NewFact(newFact,withUnit, derivedFrom)
  , makeFact
  , makeFact_
  , makeFactV
  , makeFactV_
  , Facts
  , newFacts
  , serializeFacts
  , factsMemory
  , FactBuilder
  , buildFacts
  , extendFacts
  , buildBatch
  ) where

import Data.Coerce
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Int
import Data.IORef
import Data.Maybe
import qualified Data.Vector.Storable as Vector
import Control.Monad
import Control.Monad.IO.Class
#if !MIN_VERSION_base(4,13,0)
import Control.Monad.Fail (MonadFail)
#endif
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader

import Glean.RTS.Builder (sizeOfBuilder, withBuilder)
import Glean.RTS.Foreign.Define (defineFact)
import Glean.RTS.Foreign.FactSet (FactSet)
import qualified Glean.RTS.Foreign.FactSet as FactSet
import Glean.RTS.Foreign.Lookup
import Glean.RTS.Types (lowestFid)
import Glean.Typed.Binary
import Glean.Typed.Id
import Glean.Typed.Predicate
import qualified Glean.Types as Thrift

-- -----------------------------------------------------------------------------

-- | Class of monads which support creation of new typed facts
--
-- Include 'MonadFail' constraint to avoid needing 'error'
class (MonadFail m, Monad m) => NewFact m where
  -- | Create a new fact with a the given key and value
  newFact :: Predicate p => KeyType p -> ValueType p -> m (IdOf p)

  -- | Create some facts owned by the given UnitName
  withUnit ::Thrift.UnitName -> m a -> m a

  -- | Set dependencies of an externally derived fact
  derivedFrom :: Predicate p => [Fid] -> [p] -> m ()

-- | Create a new fact in a 'NewFact' monad and return the corresponding Thrift
-- structure which will have 'Just' the passed key and value.
makeFactV
  :: forall p m. (Predicate p, NewFact m)
  => KeyType p -> ValueType p -> m p
makeFactV key value =
  (\i -> mkFact i (Just key) (Just value)) <$> newFact key value

-- | Create a new fact in a 'NewFact' monad.
makeFactV_
  :: forall p m. (Predicate p, NewFact m)
  => KeyType p -> ValueType p -> m ()
makeFactV_ key value = void $ newFact @m @p key value

-- | Create a new key-only fact in a 'NewFact' monad and return the
-- corresponding Thrift structure which will have 'Just' the passed key.
makeFact
  :: forall p m. (Predicate p, ValueType p ~ (), NewFact m)
  => KeyType p -> m p
makeFact key = makeFactV key ()

-- | Create a new key-only fact in a 'NewFact' monad.
makeFact_
  :: forall p m. (Predicate p, ValueType p ~ (), NewFact m)
  => KeyType p -> m ()
makeFact_ key = makeFactV_ @p key ()

-- | A collection of facts which can be written to and then serialized.
data Facts = Facts
  { factsPredicates :: Predicates
  , factsData :: FactSet
  , factsOwnership :: IORef (HashMap Thrift.UnitName [Int64])
  , factsDerivations :: IORef (HashMap Pid [([Fid],[Fid])])
  }

-- | Create a new empty collection of facts. New facts will be assigned
-- ids >= start id if a start id is supplied; otherwise, they will be
-- assigned ids which are guaranteed not to clash with any pids in the
-- 'Predicates' map.
newFacts
  :: Predicates -- ^ pid map
  -> Maybe Fid -- ^ start id
  -> IO Facts
newFacts ps start =
  Facts ps
    <$> FactSet.new (fromMaybe lowestFid start)
    <*> newIORef HashMap.empty
    <*> newIORef HashMap.empty

-- | Serialize the facts into a batch which can be sent via Thrift.
serializeFacts :: Facts -> IO Thrift.Batch
serializeFacts Facts{..} = do
  batch <- FactSet.serialize factsData
  ownership <- readIORef factsOwnership
  derivations <- readIORef factsDerivations
  return batch
    { Thrift.batch_owned = fmap Vector.fromList ownership
    , Thrift.batch_dependencies = HashMap.fromList
        [ (fromPid pid, map toFactDependencies deps)
        | (pid, deps) <- HashMap.toList derivations ]
    }
  where
    toFactDependencies :: ([Fid],[Fid]) -> Thrift.FactDependencies
    toFactDependencies (deps, facts) =
      Thrift.FactDependencies
        (Vector.fromList (coerce facts))
        (Vector.fromList (coerce deps))

-- | Return a rough estimate of how much memory is used by the facts.
factsMemory :: Facts -> IO Int
factsMemory = FactSet.factMemory . factsData

-- | A monad for creating fact batches.
--
-- We expose 'MonadFail' to avoid needing 'error'.
--
-- NOTE: This is very specifically not an instance of MonadIO or
-- MonadTrans as we want to restrict side effects just to fact creation.
newtype FactsM a = FactsM { runFactsM :: ReaderT Facts IO a }
  deriving(Functor,Applicative,Monad,MonadFail)

instance NewFact FactsM where
  newFact key value = FactsM $ do
    facts <- ask
    lift $ withBuilder $ \builder -> do
      buildRtsValue builder key
      key_size <- sizeOfBuilder builder
      buildRtsValue builder value
      mk facts $ \pid -> IdOf <$> defineFact
        (factsData facts)
        (pidOf pid)
        builder
        key_size
      where
        mk :: Predicate p => Facts -> (PidOf p -> f (IdOf p)) -> f (IdOf p)
        mk facts f = f $ getPid $ factsPredicates facts

  withUnit unit build = FactsM $ do
    Facts{..} <- ask
    firstId <- liftIO $ firstFreeId factsData
    a <- runFactsM build
    lastId <- liftIO $ firstFreeId factsData
    when (lastId > firstId) $ liftIO $
      modifyIORef' factsOwnership $
        HashMap.insertWith (++) unit
          [fromFid firstId, fromFid lastId - 1]
    return a

  derivedFrom :: forall p. Predicate p => [Fid] -> [p] -> FactsM ()
  derivedFrom deps facts = FactsM $ do
    Facts{..} <- ask
    let pid = pidOf (getPid factsPredicates :: PidOf p)
    liftIO $ modifyIORef' factsDerivations $
        HashMap.insertWith (<>) pid [(deps, map (idOf . getId) facts)]

-- | A fact builder
type FactBuilder = forall m. NewFact m => m ()

-- | Run a fact builder to produce facts based on the supplied
-- 'Predicates' map. The facts will be assigned ids >= start id if a
-- start id is supplied; otherwise, they will be assigned ids which are
-- guaranteed not to clash with any pids in the 'Predicates' map.
buildFacts
  :: Predicates -- ^ pid map
  -> Maybe Fid -- ^ start id
  -> FactBuilder -- ^ builder
  -> IO Facts
buildFacts ps start builder = do
  facts <- newFacts ps start
  runReaderT (runFactsM builder) facts
  return facts

extendFacts :: Facts -> FactBuilder -> IO ()
extendFacts facts builder = runReaderT (runFactsM builder) facts

-- | Run a fact builder and produce a batch of facts based on
-- the supplied 'Predicates' map. The facts will be assigned
-- ids >= start id if a start id is supplied; otherwise, they will be
-- assigned ids which are guaranteed not to clash with any pids in the
-- 'Predicates' map.
buildBatch
  :: Predicates -- ^ pid map
  -> Maybe Fid -- ^ start id
  -> FactBuilder -- ^ builder
  -> IO Thrift.Batch
buildBatch ps start builder =
  serializeFacts =<< buildFacts ps start builder