{-
  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.RTS.Foreign.FactSet
  ( FactSet
  , new
  , factCount
  , factMemory
  , allocatedMemory
  , predicateStats
  , firstFreeId
  , serialize
  , serializeReorder
  , append
  , rebase
  , renameFacts
  , DefineFlags(..)
  ) where

import Control.Exception
import Data.Int
import Data.Vector.Storable as Vector hiding (create)
import Data.Word
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Ptr

import Util.FFI

import Glean.FFI
import Glean.RTS.Foreign.Define
import Glean.RTS.Foreign.Inventory (Inventory)
import Glean.RTS.Foreign.Lookup (Lookup(..), CanLookup(..))
import Glean.RTS.Foreign.LookupCache (LookupCache)
import Glean.RTS.Foreign.Stacked (stacked)
import Glean.RTS.Foreign.Stats (marshalPredicateStats)
import Glean.RTS.Foreign.Subst (Subst)
import Glean.RTS.Types (Fid(..), Pid(..))
import qualified Glean.Types as Thrift

-- An environment for defining facts
newtype FactSet = FactSet (ForeignPtr FactSet)

instance Object FactSet where
  wrap = FactSet
  unwrap (FactSet p) = p
  destroy = glean_factset_free

instance CanLookup FactSet where
  lookupName _ = "factset"
  withLookup x f = with x $ f . glean_factset_lookup

instance CanDefine FactSet where
  withDefine x f = with x $ f . glean_factset_define

new :: Fid -> IO FactSet
new next_id = construct $ invoke $ glean_factset_new next_id

factCount :: FactSet -> IO Int
factCount facts = fromIntegral <$> with facts glean_factset_fact_count

factMemory :: FactSet -> IO Int
factMemory facts = fromIntegral <$> with facts glean_factset_fact_memory

allocatedMemory :: FactSet -> IO Int
allocatedMemory facts =
  fromIntegral <$> with facts glean_factset_allocated_memory

predicateStats :: FactSet -> IO [(Pid, Thrift.PredicateStats)]
predicateStats facts = with facts
    $ marshalPredicateStats . glean_factset_predicateStats

firstFreeId :: FactSet -> IO Fid
firstFreeId facts = with facts glean_factset_first_free_id

mkBatch :: IO (Fid, CSize, Ptr (), CSize) -> IO Thrift.Batch
mkBatch fn = mask_ $ do
  (Fid first_id, count, facts_data, facts_size) <- fn
  Thrift.Batch first_id (fromIntegral count)
    <$> unsafeMallocedByteString facts_data facts_size
    <*> pure Nothing
    <*> pure mempty
    <*> pure mempty
    <*> pure Nothing

serialize :: FactSet -> IO Thrift.Batch
serialize facts =
  with facts $ \facts_ptr -> do
  mkBatch $ invoke $ glean_factset_serialize facts_ptr

serializeReorder :: FactSet -> Vector Int64 -> IO Thrift.Batch
serializeReorder facts order =
  with facts $ \facts_ptr -> do
  unsafeWith order $ \order_ptr -> do
    mkBatch $ invoke $ glean_factset_serializeReorder
      facts_ptr
      order_ptr
      (fromIntegral (Vector.length order))

rebase :: Inventory -> Subst -> LookupCache -> FactSet -> IO (FactSet, Subst)
rebase inventory subst cache facts =
  with inventory $ \inventory_ptr ->
  with subst $ \subst_ptr ->
  with cache $ \cache_ptr ->
  with facts $ \facts_ptr ->
  mask_ $ do
    (p_factset, p_subst) <- invoke $
      glean_factset_rebase
        facts_ptr
        inventory_ptr
        subst_ptr
        cache_ptr
    (,) <$> create p_factset <*> create p_subst

append :: FactSet -> FactSet -> IO ()
append target source =
  with target $ \target_ptr ->
  with source $ \source_ptr ->
  invoke $ glean_factset_append target_ptr source_ptr

-- Prepare a Thrift batch for writing into the database by renaming and
-- deduplicating facts.
renameFacts
  :: CanLookup l
  => Inventory          -- ^ where to lookup predicates
  -> l                  -- ^ where to lookup facts
  -> Fid                -- ^ first free fact id in the database
  -> Thrift.Batch       -- ^ batch to rename
  -> DefineFlags        -- ^ flags
  -> IO (FactSet, Subst)
                        -- ^ resulting facts and substitution
renameFacts inventory base next batch flags = do
  added <- new next
  subst <- defineBatch (stacked base added) inventory batch flags
  return (added, subst)

foreign import ccall unsafe glean_factset_new
  :: Fid -> Ptr (Ptr FactSet) -> IO CString
foreign import ccall unsafe "&glean_factset_free" glean_factset_free
  :: FunPtr (Ptr FactSet -> IO ())

foreign import ccall unsafe glean_factset_fact_count
  :: Ptr FactSet -> IO CSize

foreign import ccall unsafe glean_factset_fact_memory
  :: Ptr FactSet -> IO CSize

foreign import ccall unsafe glean_factset_allocated_memory
  :: Ptr FactSet -> IO CSize

foreign import ccall safe glean_factset_predicateStats
  :: Ptr FactSet
  -> Ptr CSize
  -> Ptr (Ptr Int64)
  -> Ptr (Ptr Word64)
  -> Ptr (Ptr Word64)
  -> IO CString

foreign import ccall unsafe glean_factset_first_free_id
  :: Ptr FactSet -> IO Fid

foreign import ccall unsafe glean_factset_lookup
  :: Ptr FactSet -> Ptr Lookup

foreign import ccall unsafe glean_factset_define
  :: Ptr FactSet -> Define

foreign import ccall unsafe glean_factset_serialize
  :: Ptr FactSet
  -> Ptr Fid
  -> Ptr CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> IO CString

foreign import ccall unsafe glean_factset_serializeReorder
  :: Ptr FactSet
  -> Ptr Int64
  -> CSize
  -> Ptr Fid
  -> Ptr CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> IO CString

foreign import ccall unsafe glean_factset_rebase
  :: Ptr FactSet
  -> Ptr Inventory
  -> Ptr Subst
  -> Ptr LookupCache
  -> Ptr (Ptr FactSet)
  -> Ptr (Ptr Subst)
  -> IO CString

foreign import ccall unsafe glean_factset_append
  :: Ptr FactSet -> Ptr FactSet -> IO CString