{-
  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 MultiWayIf #-}

module Glean.Write.JSON
  ( buildJsonBatch
  , syncWriteJsonBatch
  , writeJsonBatch
  ) where

import Control.Exception hiding (catch, throw)
import Control.Monad
import Control.Monad.Reader
import Control.Monad.Catch
import Data.ByteString (ByteString)
import qualified Data.ByteString.Unsafe as BS
import Data.Coerce (coerce)
import Data.Default
import Data.IORef
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Maybe
import Compat.Prettyprinter hiding ((<>))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Vector.Storable as Vector
import Foreign hiding (void)
import Foreign.C.Types (CSize)
import TextShow hiding (Builder)

import Thrift.Protocol.JSON.Base64
import Util.FFI (invoke)

import Glean.Database.Open
import Glean.Database.Write.Batch
import Glean.Database.Schema.Types
import Glean.Database.Types as Database
import Glean.Display
import qualified Glean.FFI as FFI
import Glean.RTS as RTS
import Glean.RTS.Builder
import Glean.RTS.Set
import Glean.RTS.Constants
import qualified Glean.RTS.Foreign.JSON as J
import Glean.RTS.Types
import Glean.Angle.Types hiding (Type)
import Glean.Schema.Util
import Glean.Types as Thrift hiding (Value, Nat, Byte)


syncWriteJsonBatch
  :: Env
  -> Repo
  -> [Thrift.JsonFactBatch]
  -> Maybe Thrift.SendJsonBatchOptions
  -> IO ()
syncWriteJsonBatch env repo batches opts = do
  let batch =
        Thrift.SendJsonBatch
          { Thrift.sendJsonBatch_batches = batches
          , Thrift.sendJsonBatch_options = opts
          , Thrift.sendJsonBatch_remember = False }
  content <- writeJsonBatch env repo batch
  void $ syncWriteContentDatabase env repo content

writeJsonBatch
  :: Env
  -> Repo
  -> SendJsonBatch
  -> IO WriteContent
writeJsonBatch env repo SendJsonBatch{..} = do
  dbSchema <- withOpenDatabase env repo (return . Database.odbSchema)
  writeContentFromBatch <$>
    buildJsonBatch dbSchema sendJsonBatch_options sendJsonBatch_batches

buildJsonBatch
  :: DbSchema
  -> Maybe SendJsonBatchOptions
  -> [JsonFactBatch]
  -> IO Batch
buildJsonBatch dbSchema opts batches =
  withFactBuilder $ \builders ->
    forM_ batches $ \JsonFactBatch{..} ->
      writeFacts dbSchema (fromMaybe def opts) builders
        jsonFactBatch_predicate jsonFactBatch_facts jsonFactBatch_unit

writeFacts
  :: DbSchema
  -> SendJsonBatchOptions
  -> FactBuilder
  -> PredicateRef
  -> [ByteString]  -- ^ The facts to write, in JSON
  -> Maybe ByteString -- ^ The unit that owns the facts, if any
  -> IO ()
writeFacts dbSchema opts builder@FactBuilder{..} pred factList maybeUnit = do
  details <- predDetailsForWriting dbSchema pred
  before <- readIORef nextId
  mapM_ (writeFact dbSchema opts builder details) factList
  after <- readIORef nextId
  forM_ maybeUnit $ \unit ->
    when (after > before) $ do
      -- merge adjacent ranges
      let merge _ (Fid last : rest) | last+1 == before = Fid (after-1) : rest
          merge new old = new ++ old
      modifyIORef' owned $
        HashMap.insertWith merge unit [Fid (after-1), Fid before]

predDetailsForWriting :: DbSchema -> PredicateRef -> IO PredicateDetails
predDetailsForWriting dbSchema (PredicateRef name ver) = do
  let sourceRef = SourceRef name (Just ver)
  case lookupPredicateSourceRef sourceRef LatestSchema dbSchema of
    Right info ->
      assert (predicateInStoredSchema info) $ return info
      -- it should be impossible for predicateInStoredSchema to be
      -- False because we don't update the schema of a writable DB.
    Left err ->
      throwIO $ Thrift.Exception err

writeFact
  :: DbSchema
  -> SendJsonBatchOptions
  -> FactBuilder
  -> PredicateDetails
  -> ByteString
  -> IO ()
writeFact dbSchema opts builders details str =
  J.withParsed str $ \json ->
    runReaderT (writeJsonFact dbSchema opts details json) builders


data FactBuilder = FactBuilder
  { facts :: Builder
  , nextId :: {-# UNPACK #-} !(IORef Int64)
  , idsRef :: {-# UNPACK #-} !(IORef [Fid])  -- TODO: use a Storable Vector
  , owned :: {-# UNPACK #-} !(IORef (HashMap ByteString [Fid]))
      -- maps each unit to a *reversed* list of inclusive fact ID ranges
      -- e.g. [end2,start2,end1,start1,end0,start0]
      -- reversed so that we can prepend the next one easily and merge it
      -- with the previous one.
  }

withFactBuilder :: (FactBuilder -> IO ()) -> IO Thrift.Batch
withFactBuilder action =
  withBuilder $ \facts -> do
  nextId <- newIORef firstAnonId
  idsRef <- newIORef []
  owned <- newIORef HashMap.empty
  action FactBuilder{..}
  mem <- finishBuilder facts
  ids <- readIORef idsRef
  ownerMap <- readIORef owned
  return $ Thrift.Batch
    firstAnonId
    (fromIntegral (length ids))
    mem
    (Just $ Vector.fromList $ coerce $ reverse ids)
    (fmap (Vector.fromList . coerce . reverse) ownerMap)
    mempty
    Nothing -- TODO: we should have a schema ID for JSON batches


type WriteFacts a = ReaderT FactBuilder IO a


writeBatchFact :: Builder -> Pid -> Builder -> CSize -> WriteFacts ()
writeBatchFact builder pid clause key_size = liftIO $
  -- TODO: avoid copy
  invoke $ glean_push_fact builder pid clause key_size

namedFact :: Fid -> Pid -> Builder -> CSize -> WriteFacts Fid
namedFact fid pid clause key_size = do
  when (fromFid fid >= firstAnonId) $
    liftIO $ throwIO $ Thrift.Exception $ "id too high: " <> showt (fromFid fid)
  FactBuilder{..} <- ask
  writeBatchFact facts pid clause key_size
  id <- liftIO $ readIORef nextId
  liftIO $ writeIORef nextId $! id+1
  liftIO $ modifyIORef' idsRef (fid:)
  return (Fid id)

anonFact :: Pid -> Builder -> CSize -> WriteFacts Fid
anonFact pid clause key_size = do
  FactBuilder{..} <- ask
  writeBatchFact facts pid clause key_size
  id <- liftIO $ readIORef nextId
  liftIO $ writeIORef nextId $! id+1
  liftIO $ modifyIORef' idsRef (Fid Thrift.iNVALID_ID :)
     -- Thrift.iNVALID_ID: this tells the fact renamer that this Id
     -- maps to itself. This is so that we avoid needing to construct
     -- a mapping with all the Ids for anonymous facts.
  return (Fid id)

--
-- | Convert a JSON value to a Term, according to the type (Type) given
-- by the schema for this predicate.
--
writeJsonFact
  :: DbSchema                           -- ^ needed for looking up typerefs
  -> Thrift.SendJsonBatchOptions        -- ^ needed for no_base64_binary
  -> PredicateDetails
  -> J.Value
  -> WriteFacts ()
writeJsonFact
    dbSchema
    Thrift.SendJsonBatchOptions{..}
    details json =
  wrapJsonContextM json $ void $ factToTerm details json
  where

  factToTerm PredicateDetails{..} json@(J.Object obj) = do
    r <- lift $ J.field obj "id"
    case r of
      Just (J.Int id)
        | J.arity obj == 1 ->
          if id == fromFid invalidFid
            then invalidFactIdError
            else return (Fid id)
        | otherwise -> fact 1 (namedFact (Fid id)) obj
      _ -> fact 0 anonFact obj
    where
      -- id_arity is 0 if we don't have an id field and 1 otherwise
      fact id_arity create obj = do
        key <- do
          r <- lift $ J.field obj "key"
          case r of
            Just key -> return key
            Nothing -> badFact json
        val <- lift $ J.field obj "value"
        when (J.arity obj /= id_arity + if isJust val then 2 else 1) $
          badFact json
        withBuilder $ \clause -> do
          jsonToTerm clause predicateKeyType key
          key_size <- liftIO $ sizeOfBuilder clause
          forM_ val $ \v ->
            jsonToTerm clause predicateValueType v
          create predicatePid clause key_size
  factToTerm _ json = badFact json

  wrapJsonContextM :: J.Value -> WriteFacts a -> WriteFacts a
  wrapJsonContextM json act =
    act `catch` \(Thrift.Exception msg) -> do
      msg' <- wrapJsonContext json (pretty msg)
      throwM $ Thrift.Exception $ Text.pack $ show msg'

  wrapJsonContext :: MonadIO m => J.Value -> Doc ann -> m (Doc ann)
  wrapJsonContext json doc = do
    enc <- liftIO $ J.encode json
    return $ vcat
      [ doc
      , "in JSON term:"
      , indent 2 (pretty (Text.decodeUtf8 enc))
      ]

  badFact :: J.Value -> WriteFacts a
  badFact json = do
    msg <- wrapJsonContext json $ vcat
      [ "Expecting a fact, which should be of the form:"
      , indent 2 "{[\"id\" : N, ] \"key\": ... [, \"value\": ...]}"
      ]
    throwM $ Thrift.Exception $ Text.pack $ show msg

  -- Defining a fact with ID 0 means "don't care"; we want this
  -- behaviour because Thrift serialization will use ID 0 when the
  -- "id" field is unspecified. However, we don't support referring to
  -- fact ID 0 (see anonFact above), so catch and report that error
  -- here rather than waiting for the typechecker.
  invalidFactIdError :: WriteFacts a
  invalidFactIdError = liftIO $ do
    enc <- J.encode json
    throwIO $ Thrift.Exception $ Text.pack $ show $ vcat
      [ "Cannot use " <> pretty (fromFid invalidFid) <> " as a fact ID:"
      , indent 2 (pretty (Text.decodeUtf8 enc))
      ]

  jsonToTerm
    :: Builder
    -> Type                               -- ^ the schema type
    -> J.Value                            -- ^ the JSON value
    -> WriteFacts ()
  jsonToTerm b typ v = case (typ, v) of
    (NatTy, J.Int n) ->
      lift $ invoke $ glean_push_value_nat b $ fromIntegral n
    (ByteTy, J.Int n) ->
      lift $ invoke $ glean_push_value_byte b $ fromIntegral n
    (StringTy, J.String (J.ByteStringRef p n)) ->
      lift $ invoke $ glean_push_value_string b (castPtr p) n
    (ArrayTy ByteTy, J.String (J.ByteStringRef p n))
      | sendJsonBatchOptions_no_base64_binary -> lift $ do
          invoke $ glean_push_value_array b n
          invoke $ glean_push_value_bytes b (castPtr p) n
      | otherwise -> lift $ do
          bytes <-
            decodeBase64 <$> BS.unsafePackCStringLen (castPtr p, fromIntegral n)
          FFI.unsafeWithBytes bytes $ \ptr len -> do
            invoke $ glean_push_value_array b len
            invoke $ glean_push_value_bytes b (castPtr ptr) len
    (ArrayTy ty, J.Array arr) -> do
      let !n = J.size arr
      lift $ invoke $ glean_push_value_array b $ fromIntegral n
      when (n > 0) $ forM_ [0 .. n-1] $ \i -> do
        x <- lift $ J.index arr i
        jsonToTerm b ty x
    (SetTy byteTy, J.String (J.ByteStringRef p n))
      | ByteTy == repType byteTy -> if
        | sendJsonBatchOptions_no_base64_binary ->
          withWordRtsSet $ \rtsset ->
            lift $ do
            insertBytesRtsSet rtsset p n
            buildWordSetBytes rtsset b
        | otherwise -> lift $ do
            bytes <- decodeBase64 <$>
              BS.unsafePackCStringLen (castPtr p, fromIntegral n)
            FFI.unsafeWithBytes bytes $ \ptr len ->
              withWordRtsSet $ \rtsset -> do
                insertBytesRtsSet rtsset (castPtr ptr) len
                buildWordSetBytes rtsset b
    (SetTy natTy, J.Array set)
      | NatTy == repType natTy ->
      withWordRtsSet $ \rtsset -> do
        let !n = J.size set
        forM_ [0 .. n-1] $ \i -> do
          J.Int n <- lift $ J.index set i
          lift $ insertWordRtsSet rtsset n
        lift $ buildWordSet rtsset b
    (SetTy ty, J.Array set) ->
      withBuilder $ \tb ->
      withRtsSet $ \rtsset -> do
        let !n = J.size set
        forM_ [0 .. n-1] $ \i -> do
          x <- lift $ J.index set i
          jsonToTerm tb ty x
          lift $ insertBuilder rtsset tb
          lift $ resetBuilder tb
        lift $ buildSet rtsset b
    (RecordTy fields, J.Object obj) -> do
      let
        doField !n (FieldDef name ty) = do
          r <- lift $ J.field obj $ Text.encodeUtf8 name
          case r of
            Just val -> do
              jsonToTerm b ty val
              return $! n+1
            Nothing -> do
              lift $ defaultValue b v (name, ty)
              return n
      n <- foldM doField 0 fields
      -- ensure that all the fields mentioned in the fact are valid
      when (J.arity obj /= n) $ termError typ v
    (SumTy fields, J.Object obj)
      | J.arity obj == 1 -> do
          let -- this is O(number of alternatives) but I don't expect this is
              -- a problem
              get !n (FieldDef name ty : rest) = do
                -- TODO: avoid the encodeUtf8 (name should really be a
                -- bytestring)
                r <- lift $ J.field obj $ Text.encodeUtf8 name
                case r of
                  Just val -> do
                    lift $ invoke $ glean_push_value_selector b n
                    jsonToTerm b ty val
                  Nothing -> get (n+1) rest
              get _ _ = termError typ v
          get 0 fields
    (NamedTy _ (ExpandedType _ ty), val) -> jsonToTerm b ty val
    (PredicateTy{}, J.Int n)
      | n == fromFid invalidFid -> invalidFactIdError
      | otherwise -> lift $ invoke $ glean_push_value_fact b $ Fid n
    -- allow { "id": N } for predicate refs, this allows us to accept
    -- JSON-serialized Thrift facts.
    (PredicateTy _ (PidRef pid ref), val) ->
      -- Facts can be nested. We know from the schema the predicate of
      -- the fact at this position.
      case lookupPid pid dbSchema of
        Nothing ->
          throwError $ "unknown predicate " ++ show (displayDefault ref)
        Just deets -> do
          fid <- factToTerm deets val
          lift $ invoke $ glean_push_value_fact b fid
    (EnumeratedTy vals, J.Int n)
      | fromIntegral n < length vals ->
        lift $ invoke $ glean_push_value_selector b $ fromIntegral n
    (MaybeTy ty, val) -> do
      lift $ invoke $ glean_push_value_selector b 1
      jsonToTerm b ty val
    (BooleanTy, J.Bool False) ->
      lift $ invoke $ glean_push_value_selector b 0
    (BooleanTy, J.Bool True) ->
      lift $ invoke $ glean_push_value_selector b 1
    _otherwise -> termError typ v

  -- Remove named types so that we can match on the underlying
  -- representation of the type
  repType :: Type -> Type
  repType (NamedTy _ (ExpandedType _ ty)) = ty
  repType ty = ty

  -- Thrift might omit fields from the output if they have the
  -- default value, so we have to reconstruct the default value
  -- here.
  defaultValue :: Builder -> J.Value -> (Text, Type) -> IO ()
  defaultValue b json (fieldName, typ) = case typ of
    ByteTy -> invoke $ glean_push_value_byte b 0
    NatTy -> invoke $ glean_push_value_nat b 0
    StringTy -> invoke $ glean_push_value_string b nullPtr 0
    ArrayTy{} -> invoke $ glean_push_value_array b 0
    SetTy{} -> invoke $ glean_push_value_set b 0
    RecordTy fields ->
      mapM_ (defaultValue b json) [(nm, ty) | FieldDef nm ty <- fields ]
    SumTy (FieldDef _ ty : _) -> do
      invoke $ glean_push_value_selector b 0
      defaultValue b json (fieldName, ty)
    NamedTy _ (ExpandedType _ ty) -> defaultValue b json (fieldName, ty)
    PredicateTy{} -> do
      msg <- wrapJsonContext json $
        "Field '" <> pretty fieldName <>
        "' is missing, but it has a predicate type: " <> displayDefault typ
      throwError $ show msg
    EnumeratedTy{} -> invoke $ glean_push_value_selector b 0
    MaybeTy ty -> defaultValue b json (fieldName, lowerMaybe ty)
    BooleanTy -> invoke $ glean_push_value_selector b 0
    _otherwise -> throwError $ "internal: defaultValue: " <> show typ

  throwError :: MonadIO m => String -> m a
  throwError str = liftIO $ throwIO $ Thrift.Exception $ Text.pack str

  termError typ v = do
    msg <- wrapJsonContext v $ vcat
      [ "Error in fact. Expecting an expression of type:"
      , indent 2 (displayDefault typ)
      , "which should be of the form:"
      , indent 2 (expecting typ)
      ]
    throwError $ show msg

  allowed fields = hcat $
    punctuate ", " [ dquotes (pretty f) | FieldDef f _ <- fields ]

  expecting :: Type -> Doc ann
  expecting NatTy = "number"
  expecting ByteTy = "number"
  expecting StringTy = "string"
  expecting (ArrayTy ByteTy) = "string"
  expecting ArrayTy{} = "[..]"
  expecting SetTy{} = "set"
  expecting (RecordTy fields) =
    "{..} (allowed fields: " <> allowed fields <> ")"
  expecting (SumTy fields) =
    "{\"field\" : value} (allowed fields: " <> allowed fields <> ")"
  expecting (EnumeratedTy vals) =
    "number (< " <> pretty (length vals) <> ")"
  expecting (MaybeTy ty) =
    expecting (lowerMaybe ty)
  expecting BooleanTy = "true or false"
  expecting PredicateTy{} = "N, or {\"id\": N }, where N is a fact ID"
  expecting _ = error "expecting"