{-
  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 TypeApplications, QuasiQuotes #-}
module JSONWriteTest (main) where

import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BC
import Data.Default
import Data.List (isInfixOf)
import Data.Proxy
import Test.HUnit

import TestRunner
import Thrift.Protocol.JSON
import Util.String.Quasi
import Util.Control.Exception (tryAll)

import Glean (factIdRange)
import Glean.Init
import Glean.Database.Test
import Glean.Write.JSON (syncWriteJsonBatch)
import qualified Glean.Schema.GleanTest.Types as Glean.Test
import qualified Glean.Schema.Sys.Types as Sys
import Glean.Query.Angle
import Glean.Query.Thrift
import Glean.Typed hiding (end)
import Glean.Types

import TestData

mkBatch :: PredicateRef -> [ByteString] -> JsonFactBatch
mkBatch ref facts =
  JsonFactBatch
    { jsonFactBatch_predicate = ref
    , jsonFactBatch_facts = facts
    , jsonFactBatch_unit = Nothing
    }

writeJsonBatchTest :: Test
writeJsonBatchTest = TestCase $ withEmptyTestDB [] $ \env repo -> do
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob))
        [ "{ \"id\": 1025, \"key\": \"hello\"}"
        , "{ \"id\": 1026, \"key\": \"bye\"}"
        ]
      ]
    options = Just SendJsonBatchOptions
      { sendJsonBatchOptions_no_base64_binary = True
      , sendJsonBatchOptions_schema_id = Nothing
      }

    options' = Just SendJsonBatchOptions
      { sendJsonBatchOptions_no_base64_binary = False
      , sendJsonBatchOptions_schema_id = Nothing
      }

  _ <- syncWriteJsonBatch env repo batches options

  results <- runQuery_ env repo $ query $
    predicate @Sys.Blob (byteArray "hello")
  case results of
    [Sys.Blob{}] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - sys.Blob"

  -- Write another fact into the DB
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob))
        [ "{ \"id\": 1025, \"key\": \"test\"}" -- overlaps with previous
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Sys.Blob (byteArray "test")
  test_blob_id <- case results of
    [Sys.Blob{..}] -> return blob_id
    _ -> do _ <- assertFailure "syncWriteJsonBatch - sys.Blob 2"; return 0

  -- Test a predicate ref to an earlier fact
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        [ "{ \"id\": 1025, \"key\": { \"sum_\" : { \"d\": " <>
          BC.pack (show test_blob_id) <> "}, \"pred\" : " <>
          BC.pack (show test_blob_id) <> "}}"
            -- most of the fields are defaulted
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate wild
  test_predicate_id <- case results of
    [Glean.Test.Predicate{..}] -> return predicate_id
    _ -> do
      _ <- assertFailure "syncWriteJsonBatch - glean.test.Predicate";
      return 0

  -- insert the same fact again, ensure we still have just one
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate wild
  assertBool "syncWriteJsonBatch - glean.test.Predicate 2" $ case results of
    [Glean.Test.Predicate{..}] -> predicate_id == test_predicate_id
    _ -> False

  -- Test a predicate ref to a fact in the same batch
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob))
        [ "{ \"id\": 1025, \"key\": \"test2\"}"
        ]
      , mkBatch (getName (Proxy @Glean.Test.Predicate))
        [ "{ \"id\": 1026, \"key\": { \"sum_\" : { \"d\": 1025" <>
          "}, \"pred\" : 1025}}"
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "test2") end
  test_pred <- case results of
    [fact@Glean.Test.Predicate{}] -> return fact
    _ -> do
      _ <- assertFailure "syncWriteJsonBatch - glean.test.Predicate 3"
      return def

  -- test a more interesting value:
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        [ serializeJSON $ Glean.Test.Predicate 1027 $
            let blob = Sys.Blob 1028 (Just "blobby") in
            Just $ kitchenSink1 {
              Glean.Test.kitchenSink_pred = blob,
              Glean.Test.kitchenSink_sum_ = Glean.Test.KitchenSink_sum__d blob }
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options'
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "blobby") end
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - glean.test.Predicate.1"

  -- nested facts can have duplicate Ids, provided we don't refer to them
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        -- Construct a value using "def", which will set all the IDs to 0.
        -- Note that we *must* define the nested facts, otherwise they
        -- will default to a reference to fact 0, which will be an error.
        [ serializeJSON $ def
            { Glean.Test.predicate_key = Just $ def
                { Glean.Test.kitchenSink_pred =
                    def { Sys.blob_key = Just "abc" }
                , Glean.Test.kitchenSink_sum_ =
                  Glean.Test.KitchenSink_sum__d $
                    def { Sys.blob_key = Just "def" }
                , Glean.Test.kitchenSink_named_record_ =
                    def { Glean.Test.rec_beta = Glean.Test.Sum_wed True }
                , Glean.Test.kitchenSink_named_sum_ =
                    Glean.Test.Sum_wed True
                , Glean.Test.kitchenSink_maybe_ = Nothing
                }
            }
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options'
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "abc") end
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - glean.test.Predicate.1"

  -- Round-trip test: serialize the fact and insert it again
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        [ serializeJSON test_pred
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "test2") end
  void $ case results of
    [fact@Glean.Test.Predicate{}] -> return fact
    _ -> do
      _ <- assertFailure "syncWriteJsonBatch - glean.test.Predicate 4"
      return def

  -- Test for ids very far apart
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob))
        [ "{ \"id\": 1024, \"key\": \"test4\"}"
        , "{ \"id\": 4611686018427387903, \"key\": \"test5\"}"
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Sys.Blob (byteArray "test5")
  print results
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - sparse ids"

  -- Test for missing ids
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob))
        [ "{\"key\": " <> BC.pack (show s) <> "}"
        | s <- sequence ["ab","bc","cd"]
        ]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Sys.Blob (byteArray "bcd")
  print results
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - missing ids"

  -- Test for mixture of named + anon facts
  let
    batches =
      [ mkBatch (getName (Proxy @Sys.Blob)) $
        [ "{\"key\": " <> BC.pack (show s) <> "}"
        | s <- sequence ["ab","bc","cd"]
        ] ++
        [[s| {"id": 10000, "key": "abba"} |]]
      , mkBatch (getName (Proxy @Glean.Test.Predicate))
        [[s|
          { "key": { "sum_" : { "d": 10000 }, "pred" : 10000}}
        |]]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "abba") end
  print results
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - mixture of named + anon"

  -- Test for nested facts
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        [[s|
        { "key":
          { "sum_": { "d": { "id": 10000, "key": "abba" }}
          , "pred": 10000
          }
        }
        |]]  --  id 10000 is defined and used in the same fact
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "abba") end
  print results
  case results of
    [_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - nested facts"

  -- Named facts that refer to anon facts
  let
    batches =
      [ mkBatch (getName (Proxy @Glean.Test.Predicate))
        [[s|
        { "id": 3
        , "key":
          { "sum_": { "d": { "key": "abba" }}
          , "pred": { "key": "baab" }
          }
        }
        |]
        ,[s|
        { "id": 4
        , "key":
          { "sum_": { "c": 3 }
          , "pred": { "key": "baab" }
          }
        }
        |]]
      ]
  _ <- syncWriteJsonBatch env repo batches options
  results <- runQuery_ env repo $ query $
    predicate @Glean.Test.Predicate $
      rec $ field @"pred" (byteArray "baab") end
  print results
  case results of
    [_,_] -> return ()
    _ -> assertFailure "syncWriteJsonBatch - named fact that refers to anon facts"

  -- same key with different values
  withEmptyTestDB [] $ \env repo -> do
    let
      batches =
        [ mkBatch (getName (Proxy @Glean.Test.KeyValue))
          [[s| { "key":   { "kstring": "A", "knat" : 1 },
                 "value": { "vstring": "A", "vnat" : 1 } }
          |]]
        , mkBatch (getName (Proxy @Glean.Test.KeyValue))
          [[s| { "key":   { "kstring": "A", "knat" : 1 },
                 "value": { "vstring": "B", "vnat" : 2 } }
          |]]
        ]
      label = "syncWriteJsonBatch - same key with different values"
    FactIdRange _ finish <- factIdRange env repo
    r <- tryAll $ syncWriteJsonBatch env repo batches options
    print r
    case r of
      Right _ -> assertFailure label
      Left err -> assertBool label $
        ("invalid fact redefinition of $" <> show finish) `isInfixOf` show err

main :: IO ()
main = withUnitTest $ testRunner $ TestList
  [ TestLabel "writeJsonBatchTest" writeJsonBatchTest
  ]