content: {-
    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 TemplateHaskell, QuasiQuotes, DeriveGeneric #-}
  -- | Utilities for writing data to Glean
  module Glean.Write
    ( fileToBatches
    , parseRef
    , parseJsonFactBatches
    ,schemaIdToOpts) where
  
  import Control.Exception (ErrorCall(ErrorCall), throwIO)
  import Control.Monad.Extra
  import Data.Aeson
  import qualified Data.Aeson.Types as Aeson
  import Data.ByteString (ByteString)
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy.Char8 as LB
  import Data.Maybe
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import qualified Data.Vector as Vector
  import Foreign(Ptr)
  import Foreign.C(CString, CChar, CLong(..))
  
  import Mangle.TH
  import Foreign.CPP.Dynamic (Dynamic, callJSONParserFFI)
  import Util.String.Quasi
  
  import Glean.Types hiding (Value)
  import Glean.Schema.Util
  import System.FilePath (takeExtension)
  import qualified System.Process.ByteString as BS
  import System.Exit (ExitCode(ExitSuccess))
  import Control.Applicative ((<|>))
  import Data.Either (partitionEithers)
  
  $(mangle
    [s|
      folly::dynamic* facebook::glean::cpp::parseJSONFacts(
          const char*, int64_t, char **)
    |] [d|
      foreign import ccall safe c_parseJsonFacts
        :: CString -> CLong -> Ptr (Ptr CChar) -> IO (Ptr Dynamic)
    |])
  
  newtype ParseJsonFactBatchForWriteServer = ParseJsonFactBatchForWriteServer
    {getBatch :: JsonItem}
  instance FromJSON ParseJsonFactBatchForWriteServer where
    parseJSON =
      fmap ParseJsonFactBatchForWriteServer <$>
        parseJsonFactBatchGen (withText "fact" $ return . Text.encodeUtf8)
  
  fileToBatches :: FilePath -> IO ([JsonFactBatch], Maybe SchemaId)
  fileToBatches file = do
    bs <- if takeExtension file == ".zst" then do
      (exit, bs, err) <- BS.readProcessWithExitCode "zstd" [file,"-d","-c"] ""
      when (exit /= ExitSuccess) $
        throwIO $ ErrorCall $ file ++ ": " ++ show err
      return bs
      else B.readFile file
    r <- Foreign.CPP.Dynamic.callJSONParserFFI c_parseJsonFacts bs
    case r of
      Right val -> case Aeson.parse parseJSON val of
        Aeson.Error str -> throwIO $ ErrorCall $ file ++ ": " ++ str
        Aeson.Success x -> return $ jsonItemsToPair $ map getBatch x
      Left err -> throwIO $ ErrorCall $ file ++ ": " ++ Text.unpack err
  
  parsePredicate :: Value -> Aeson.Parser PredicateRef
  parsePredicate = withText "predicate" $ \txt -> do
    let SourceRef pred mbVersion = parseRef txt
    return (PredicateRef pred (fromMaybe 1 mbVersion))
  
  parsePredicateRef :: Value -> Aeson.Parser PredicateRef
  parsePredicateRef = withObject "predicate ref" $ \obj ->
    PredicateRef
      <$> obj .: "name"
      <*> obj .: "version"
  
  parseFact :: Value -> Aeson.Parser ByteString
  parseFact = withObject "fact" $ \obj -> return (LB.toStrict (encode obj))
  
  -- | Given a fact parser, returns a 'JsonFactBatch' parser
  parseJsonFactBatchGen
    :: (Value -> Aeson.Parser Json) -> Value -> Aeson.Parser JsonItem
  parseJsonFactBatchGen parseFact = withObject "JsonFactBatch" $ \v ->
    fmap JsonItemBatch (JsonFactBatch
      <$> Aeson.explicitParseField parsePred v "predicate"
      <*> Aeson.explicitParseField parseFacts v "facts"
      <*> fmap (fmap Text.encodeUtf8) (v .:? "unit"))
    <|>
    (JsonItemSchemaId
      <$> fmap SchemaId (v .: "schema_id"))
    where
      parsePred v = parsePredicate v `mplus` parsePredicateRef v
      parseFacts = withArray "facts" (mapM parseFact . Vector.toList)
  
  -- | Parser that expects facts to be JSON objects
  parseJsonFactBatches :: Value -> Aeson.Parser ([JsonFactBatch], Maybe SchemaId)
  parseJsonFactBatches = withArray "JsonFactBatch" $ \vec ->
    jsonItemsToPair <$> mapM (parseJsonFactBatchGen parseFact) (Vector.toList vec)
  
  data JsonItem = JsonItemBatch JsonFactBatch | JsonItemSchemaId SchemaId
  
  jsonItemsToPair :: [JsonItem] -> ([JsonFactBatch], Maybe SchemaId)
  jsonItemsToPair items = (batches, listToMaybe schemas)
   where
    (batches, schemas) = partitionEithers (map toEither items)
    toEither (JsonItemBatch item) = Left item
    toEither (JsonItemSchemaId schema_id) = Right schema_id
  
  schemaIdToOpts :: Maybe SchemaId -> Maybe SendJsonBatchOptions
  schemaIdToOpts schema_id = case schema_id of
    Nothing -> Nothing
    Just s -> Just SendJsonBatchOptions {
        sendJsonBatchOptions_schema_id = Just s,
        sendJsonBatchOptions_no_base64_binary = False
      }
  