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

import Control.Exception
import Data.Default
import qualified Data.Map.Strict as Map
import Test.HUnit
import Util.String.Quasi
import Data.Proxy
import qualified Data.Text as Text
import Control.Concurrent
import Control.Monad
import Control.Concurrent.Async
import qualified Data.HashMap.Strict as HashMap
import Util.STM
import System.IO.Temp (withSystemTempDirectory)
import System.FilePath

import TestRunner

import Glean.Angle.Types
import Glean.Backend.Types hiding (deriveStored)
import qualified Glean.Backend.Types as Backend
import Glean.Init
import Glean.Database.Open
import Glean.Database.Schema.Types
import Glean.Database.Test
import Glean.Database.Types
import qualified Glean.Database.Catalog as Catalog
import Glean.Internal.Types
import Glean.Typed
import Glean.Types as Thrift
import Glean.Test.HUnit
import qualified Glean.Schema.GleanTest.Types as Glean.Test
import Glean.Query.Thrift
import Glean.Write.JSON (syncWriteJsonBatch)

import TestDB

main :: IO ()
main = withUnitTest $ testRunner $ TestList
  [ TestLabel "deriveStored" $ testDerivation $ deriveSerial $ const mempty
  , TestLabel "deriveParallel" deriveParallelTest
  , TestLabel "deriveStoredLogging" testDeriveStoredLogging
  , TestLabel "deriveIncremental" deriveIncrementalTest
  ]

testDerivation :: RunDerive -> Test
testDerivation derive = TestList
  [ TestLabel "deriveTest" $ deriveTest derive
  , TestLabel "completenessTest" $ completenessTest derive
  , TestLabel "deriveDeleteDeriveTest" $ deriveDeleteDeriveTest derive
  ]

type RunDerive = forall p. Predicate p => Env -> Repo -> Proxy p -> IO Int

-- Test that deleting a DB and re-creating it works
deriveDeleteDeriveTest :: RunDerive -> Test
deriveDeleteDeriveTest runDerive = TestLabel "deriveDeleteDerive" $
  TestCase $ withTestEnv [] $ \env -> do
    let repo = Thrift.Repo "dbtest-repo" "1"
    createTestDB env repo
    deriveTestCases runDerive env repo
    void $ deleteDatabase env repo
    createTestDB env repo
    deriveTestCases runDerive env repo

-- Test deriving stored facts
deriveTest :: RunDerive -> Test
deriveTest runDerive = dbTestCaseWritable (deriveTestCases runDerive)

deriveParallelTest :: Test
deriveParallelTest = TestLabel "deriveParallelTest" $
  TestCase $ withWritableTestDB [] $ \env repo -> do
    derivedCount <-
      deriveParallel (const mempty) env repo
        (Proxy @Glean.Test.StoredDualStringPair)
        (ParallelDerivation
          { parallelDerivation_outer_predicate = "glean.test.StringPair"
          , parallelDerivation_inner_query =
              "glean.test.StoredDualStringPair { fst = X }"
          , parallelDerivation_min_batch_size = Just 1
          })
    assertEqual "deriveParallelTest" 4 derivedCount

deriveTestCases :: RunDerive -> Env -> Repo -> IO ()
deriveTestCases runDerive env repo = do
  let derive :: forall p. Predicate p => Proxy p -> IO Int
      derive = runDerive env repo
  -- initially zero facts
  results <- runQuery_ env repo $
    angle @Glean.Test.StoredRevStringPair $
    [s|
      glean.test.StoredRevStringPair _
    |]
  assertEqual "deriveTest - pre-derive" 0 (length results)

  -- compute and store. Should derive 6 facts
  derivedCount <- derive (Proxy @Glean.Test.StoredRevStringPair)
  assertEqual "deriveTest - derivation" 6 derivedCount

  -- now there should be 6 facts stored in the DB
  results <- runQuery_ env repo $
    angle @Glean.Test.StoredRevStringPair $
    [s|
      glean.test.StoredRevStringPair _
    |]
  assertEqual "deriveTest - post-derive" 6 (length results)

  -- should not confuse predicates from different repos
  let pred = Proxy @Glean.Test.StoredRevStringPairWithA
      PredicateRef{..} = getName pred
      sref = SourceRef predicateRef_name (Just predicateRef_version)
  pid <- withOpenDatabase env repo $ \opendb ->
    either (throwIO . ErrorCall . Text.unpack) return $
      lookupPredicateSourceRef sref LatestSchema (odbSchema opendb)
  addDummyDerivationForPredicate (predicateId pid) env
  derivedCount <- derive pred
  assertEqual "deriveTest - repo check" 1 derivedCount

addDummyDerivationForPredicate :: PredicateId -> Env -> IO ()
addDummyDerivationForPredicate pred env =
  let repo = Thrift.Repo "name" "hash"
  in
  atomically
    $ modifyTVar' (envDerivations env)
    $ HashMap.insert (repo, pred) Derivation
      { derivationStart = undefined
      , derivationFinished = undefined
      , derivationStats = error "wrong repo!"
      , derivationPendingWrites = undefined
      , derivationError = undefined
      , derivationHandle = "dummy-handle"
      }

-- Test completeness constraint enforcement
completenessTest :: RunDerive -> Test
completenessTest runDerive = dbTestCaseWritable $ \env repo -> do
  let derive :: forall p. Predicate p => Proxy p -> IO Int
      derive = runDerive env repo
  -- deriving a stored predicate depending on a non-derived one succeeds
  derivedCount <-
    derive (Proxy @Glean.Test.StoredRevStringPairWithRev)
  assertEqual "completenessTest - derived depending on non-derived"
    6 derivedCount

  -- deriving a non-derived predicate fails
  assertThrows "completenessTest - non-derived" Thrift.NotAStoredPredicate $
    void $ derive (Proxy @Glean.Test.StringPair)

  -- deriving a derived but not stored predicate fails
  assertThrows "completenessTest - non-stored" Thrift.NotAStoredPredicate $
    void $ derive (Proxy @Glean.Test.RevStringPair)

  -- deriving a stored predicate that depends on incomplete predicates fails
  assertThrows "completenessTest - incomplete dep"
    (Thrift.IncompleteDependencies
      [ getName $ Proxy @Glean.Test.StoredRevStringPair
      , getName $ Proxy @Glean.Test.StoredRevStringPairWithA
      ])
    $ void $ derive (Proxy @Glean.Test.StoredRevStringPairSum)

  -- parallel derivation works
  let run = derive (Proxy @Glean.Test.StoredRevStringPair)
  (derivedCount1, derivedCount2) <- concurrently run run
  assertEqual "deriveTest - parallel"
    (6, 6)
    (derivedCount1, derivedCount2)

  -- should derive 6 facts now that the dependency is complete
  derivedCount <- derive (Proxy @Glean.Test.StoredRevStringPairWithA)
  assertEqual "deriveTest - subsequent derivation" 1 derivedCount

  -- deriving a complete predicate is a no-op
  derivedCount <- derive (Proxy @Glean.Test.StoredRevStringPair)
  assertEqual "deriveTest -  complete" 6 derivedCount

  -- deriving a predicate with no facts records it as complete
  let prox = Proxy @Glean.Test.EmptyStoredStringPair
      pred = getName prox
  derivedCount <- derive prox
  complete <- atomically $
    metaCompletePredicates <$> Catalog.readMeta (envCatalog env) repo
  assertEqual "deriveTest -  empty"
    (0, True)
    (derivedCount, pred `elem` complete)

testDeriveStoredLogging :: Test
testDeriveStoredLogging = dbTestCaseWritable $ \env repo -> do
  tvar <- newTVarIO (0 :: Int)
  let log _ = atomically $ modifyTVar' tvar (+ 1)

  -- The result is logged once and only once.
  replicateM_ 10 $
    deriveSerial log env repo $ Proxy @Glean.Test.StoredRevStringPairWithRev
  logCalls <- readTVarIO tvar
  assertEqual "deriveStoredLogging - logs once" 1 logCalls


deriveIncrementalTest :: Test
deriveIncrementalTest = TestLabel "incremental" $ TestList
  [ TestLabel "statement" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node : nat
        predicate IsNode : nat stored X where Node X
      }|]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 4}|]
          ]
      ]
      (PredicateRef "all.IsNode" 1)
      $ \stats -> do
        assertEqual "results" 1 (userQueryStats_result_count stats)
        assertEqual "facts searched" 1 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)
  , TestLabel "sequence" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node : nat
        predicate NodeXNode : { a: Node, b: Node } stored
          { A, B } where
            A = Node _;
            B = Node _;
      }|]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 4}|]
          ]
      ]
      (PredicateRef "all.NodeXNode" 1)
      $ \stats -> do
        -- (new x old) + (old x new) + (new x new)
        assertEqual "results" 7 (userQueryStats_result_count stats)
        assertEqual "facts searched" 11 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)
  , TestLabel "disjunction" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node1 : nat
        predicate Node2 : nat
        predicate NodeCombinations : { a: nat, b: nat } stored
          { A, B } where
            (Node1 A; Node2 B) |
            (Node1 B; Node2 A)
      }|]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]

      , mkBatch (PredicateRef "all.Node2" 1)
          [ [s|{"key": 4}|]
          , [s|{"key": 5}|]
          , [s|{"key": 6}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 7}|]
          ]
      ]
      (PredicateRef "all.NodeCombinations" 1)
      $ \stats -> do
        -- (new x old) + (old x new)
        assertEqual "results" 6 (userQueryStats_result_count stats)
        assertEqual "facts searched" 8 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)
  , TestLabel "disjunction followed by sequence" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node1 : nat
        predicate Node2 : nat
        predicate Node1and2XNode2 : { a: nat, b: nat } stored
          { A, B } where
            (Node1 A | Node2 A);
            Node2 B
      }|]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]

      , mkBatch (PredicateRef "all.Node2" 1)
          [ [s|{"key": 4}|]
          , [s|{"key": 5}|]
          , [s|{"key": 6}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 7}|]
          ]
      ]
      (PredicateRef "all.Node1and2XNode2" 1)
      $ \stats -> do
        -- (new x old)
        assertEqual "results" 3 (userQueryStats_result_count stats)
        assertEqual "facts searched" 4 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)

  , TestLabel "disjunction followed by sequence(2)"  $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node1 : nat
        predicate Node2 : nat
        predicate Node1and2XNode2 : { a: nat, b: nat } stored
          { A, B } where
            (Node1 A | Node2 A);
            Node2 B
      }|]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node1" 1)
          [ [s|{"key": 7}|]
          ]
      , mkBatch (PredicateRef "all.Node2" 1)
          [ [s|{"key": 4}|]
          , [s|{"key": 5}|]
          , [s|{"key": 6}|]
          ]
      ]
      (PredicateRef "all.Node1and2XNode2" 1)
      $ \stats -> do
        -- (new x old)
        assertEqual "results" 21 (userQueryStats_result_count stats)
        assertEqual "facts searched" 28 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)

  , TestLabel "disjunction of expression and FactGenerator" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node : nat
        predicate Node2 : { a: nat, b: nat } stored
          { A, B } where
            Node A;
            B = (X where Node X) | [11,12,13,14,15][..];
      }|]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 1}|]
          , [s|{"key": 2}|]
          , [s|{"key": 3}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": 4}|]
          ]
      ]
      (PredicateRef "all.Node2" 1)
      $ \stats -> do
        -- crucially we don't want to have the 3*5 case of matching
        -- Node<base> + the array element generator.
        assertEqual "results" 12 -- 1*5 + 1*3 + 3*1 + 1*1
          (userQueryStats_result_count stats)
        assertEqual "facts searched" 11 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)

  , TestLabel "continuations" $ TestCase $
    derivationStats
      def { derivePredicateOptions_max_results_per_query = Just 1 }
      [s|schema all.1 {
        predicate Node : string
        predicate Node2 : string stored A where Node A
      }|]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": "D"}|]
          , [s|{"key": "E"}|]
          , [s|{"key": "F"}|]
          , [s|{"key": "G"}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Node" 1)
          [ [s|{"key": "A"}|]
          , [s|{"key": "B"}|]
          , [s|{"key": "C"}|]
          ]
      ]
      (PredicateRef "all.Node2" 1)
      $ \stats -> do
        -- the fact iterators should be serialised correctly such that
        -- after the query is resumed it will still not pick-up the
        -- lexicographically greater facts of the base db.
        assertEqual "incremental" 3 (userQueryStats_result_count stats)
  , TestLabel "empty derivation" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate Node : string
        predicate Node2 : string stored A where Node A
      }|]
      [ mkBatch (PredicateRef "all.Node" 1) [] ]
      [ mkBatch (PredicateRef "all.Node" 1) [] ]
      (PredicateRef "all.Node2" 1)
      $ \stats -> do
        -- should not throw an error
        assertEqual "incremental" 0 (userQueryStats_result_count stats)

  , TestLabel "optimise - sequence of prefix seeks" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate P : { a: nat, b: nat }
        predicate Q : { a: nat, b: nat }
        predicate R : { a: nat, b: nat } stored
          { A, B } where
            P { K, A };
            Q { K, B };
      }|]
      [ mkBatch (PredicateRef "all.P" 1)
          [ [s|{"key": { "a": 1, "b": 1}}|]
          , [s|{"key": { "a": 2, "b": 2}}|]
          , [s|{"key": { "a": 3, "b": 3}}|]
          , [s|{"key": { "a": 4, "b": 4}}|]
          ]
      , mkBatch (PredicateRef "all.Q" 1)
          [ [s|{"key": { "a": 1, "b": 1}}|]
          , [s|{"key": { "a": 2, "b": 2}}|]
          , [s|{"key": { "a": 3, "b": 3}}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Q" 1)
          [ [s|{"key": { "a": 4, "b": 4}}|]
          ]
      ]
      (PredicateRef "all.R" 1)
      $ \stats -> do
        assertEqual "results" 1 (userQueryStats_result_count stats)
        assertEqual "facts searched" 5 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)

  , TestLabel "optimise - disjunction with prefix seeks" $ TestCase $
    derivationStats def
      [s|schema all.1 {
        predicate P : { a: nat, b: nat }
        predicate Q : { a: nat, b: nat }
        predicate R : { a: nat, b: nat }
        predicate S : { a: nat, b: nat } stored
          { A, B } where
            P { K, A };
            (Q { K, B } | R { K, B });
      }|]
      [ mkBatch (PredicateRef "all.P" 1)
          [ [s|{"key": { "a": 1, "b": 1}}|]
          , [s|{"key": { "a": 2, "b": 2}}|]
          , [s|{"key": { "a": 3, "b": 3}}|]
          , [s|{"key": { "a": 4, "b": 4}}|]
          ]
      , mkBatch (PredicateRef "all.Q" 1)
          [ [s|{"key": { "a": 1, "b": 1}}|]
          ]
      , mkBatch (PredicateRef "all.R" 1)
          [ [s|{"key": { "a": 2, "b": 2}}|]
          ]
      ]
      [ mkBatch (PredicateRef "all.Q" 1)
          [ [s|{"key": { "a": 3, "b": 3}}|]
          ]
      , mkBatch (PredicateRef "all.R" 1)
          [ [s|{"key": { "a": 4, "b": 4}}|]
          , [s|{"key": { "a": 5, "b": 5}}|]
          , [s|{"key": { "a": 6, "b": 6}}|]
          , [s|{"key": { "a": 7, "b": 7}}|]
          , [s|{"key": { "a": 8, "b": 8}}|]
          , [s|{"key": { "a": 9, "b": 9}}|]
          ]
      ]
      (PredicateRef "all.S" 1)
      $ \stats -> do
        assertEqual "results" 2 (userQueryStats_result_count stats)
        assertEqual "facts searched" 6 $ maybe 0 (sum . Map.elems)
          (userQueryStats_facts_searched stats)
    ]
  where
    mkBatch ref facts =
      JsonFactBatch
        { jsonFactBatch_predicate = ref
        , jsonFactBatch_facts = facts
        , jsonFactBatch_unit = Nothing
        }

    derivationStats
      :: Thrift.DerivePredicateOptions
      -> String
      -> [JsonFactBatch]
      -> [JsonFactBatch]
      -> PredicateRef
      -> (UserQueryStats -> IO a)
      -> IO a
    derivationStats opts schema baseFacts topFacts pref action =
      withSchemaFile latestAngleVersion schema $ \root file -> do
      let settings =
            [ setRoot root
            , setSchemaPath file
            ]

          opts' = opts { derivePredicateOptions_collect_facts_searched = True }

      withTestEnv settings $ \env -> do
        let base = Repo "base" "0"
        kickOffTestDB env base id
        void $ syncWriteJsonBatch env base baseFacts Nothing
        void $ completePredicates env base (CompletePredicates_axiom def)
        _ <- deriveStored' env base opts' pref
        completeTestDB env base

        stacked <- stackedDB env base (Repo "base-stacked" "0")
        void $ syncWriteJsonBatch env stacked topFacts Nothing

        stats <- deriveStored' env stacked opts' pref
        action stats

    stacked (Repo name hash) = Dependencies_stacked $ Stacked name hash Nothing
    stackedDB :: Env -> Repo -> Repo -> IO Repo
    stackedDB env base top = do
      kickOffTestDB env top $ \kickOff ->
        kickOff { kickOff_dependencies = Just $ stacked base }
      return top

    -- get derivation stats
    deriveStored'
      :: Env
      -> Repo
      -> Thrift.DerivePredicateOptions
      -> PredicateRef
      -> IO UserQueryStats
    deriveStored' env repo opts pref = loop
      where
        loop = do
          let query = derivePredicateQuery opts pref Nothing
          res <- Backend.deriveStored env (const mempty) repo query
          case res of
            DerivationStatus_complete complete ->
              return $ derivationComplete_stats complete
            DerivationStatus_ongoing _ -> do
              threadDelay (ceiling @Double 1e6)
              loop

    withSchemaFile :: AngleVersion -> String -> (FilePath -> FilePath -> IO a) -> IO a
    withSchemaFile (AngleVersion version) str action = do
      withSystemTempDirectory "glean-dbtest" $ \root -> do
        let newSchemaFile = root </> "schema"
        appendFile newSchemaFile $ "version: " <> show version
        appendFile newSchemaFile str
        action root newSchemaFile

deriveSerial
  :: forall p. Predicate p
  => LogDerivationResult
  -> Env
  -> Repo
  -> Proxy p
  -> IO Int
deriveSerial log env repo proxy = deriveStored log env repo proxy Nothing

deriveParallel
  :: forall p. Predicate p
  => LogDerivationResult
  -> Env
  -> Repo
  -> Proxy p
  -> ParallelDerivation
  -> IO Int
deriveParallel log env repo proxy par =
  deriveStored log env repo proxy (Just par)

deriveStored
  :: forall p. Predicate p
  => LogDerivationResult
  -> Env
  -> Repo
  -> Proxy p
  -> Maybe ParallelDerivation
  -> IO Int
deriveStored log env repo proxy par = do
  () <- loop
  length <$> runQuery_ env repo (allFacts @p)
  where
    pred = getName proxy
    loop = do
      res <- Backend.deriveStored env log repo
        $ derivePredicateQuery def pred par
      case res of
        DerivationStatus_ongoing _ -> threadDelay (ceiling @Double 1e6) >> loop
        DerivationStatus_complete _ -> return ()

derivePredicateQuery
  :: Thrift.DerivePredicateOptions
  -> PredicateRef
  -> Maybe Thrift.ParallelDerivation
  -> Thrift.DerivePredicateQuery
derivePredicateQuery opts (PredicateRef name version) par = def
  { derivePredicateQuery_predicate = name
  , derivePredicateQuery_predicate_version = Just version
  , derivePredicateQuery_parallel = par
  , derivePredicateQuery_options = Just opts
  }