{-
  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 #-}
module Schema.Multi (main) where

import Control.Exception
import Control.Monad
import Data.Either
import Data.Default (def)
import Data.List (isInfixOf)
import qualified Data.Map as Map
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text
import System.Directory
import System.FilePath
import System.IO.Temp
import Test.HUnit

import TestRunner
import Thrift.Util
import Util.String.Quasi

import Glean.Angle.Types (Type_(..), FieldDef_(..))
import Glean.Backend.Types (userQueryFacts, userQuery)
import Glean.Database.Config
import Glean.Database.Schema
import Glean.Database.Test
import Glean.Init
import qualified Glean.RTS.Term as RTS
import Glean.Types as Thrift
import qualified Glean.Internal.Types as Internal
import Glean.Write.SendBatch
import Glean.Write.JSON


import Schema.Lib
import Glean.ServerConfig.Types (config_strict_query_schema_id)

multiSchemaTest :: Test
multiSchemaTest = TestCase $
  withSystemTempDirectory "glean-dbtest" $ \root -> do
    let
      schema_v0_file = root </> "schema0"
      schema_v0 =
        [s|
          schema x.1 {
            predicate P : { a: string }
          }

          schema y.1 {
            import x.1
            predicate Q : { p : x.P }
          }

          schema z.1 {
            predicate R : string
          }

          schema derived.1 {
            import x.1
            predicate D : { a : string }
              { a = A } where x.P { a = A }

            type T = string
            predicate S : T
              stored X where x.P { a = X }
          }

          schema evolved.1 {
            predicate P : enum { A | }
            predicate Q : { ref : P }
          }

          schema evolved.2 {
            predicate P : enum { A | B }
            predicate Q : { ref : P }
          }

          schema evolved.2 evolves evolved.1

          schema evolved.3 {
            predicate P : enum { A | B | C }
            predicate Q : { ref : P }
          }

          schema all.1 : x.1, y.1, z.1, derived.1, evolved.1, evolved.2, evolved.3 {}
        |]

      schema_v1_file = root </> "schema1"
      schema_v1 =
        [s|
          schema x.1 {
            # P has a new field
            predicate P : { a: string, b: nat }
          }

          schema y.1 {
            import x.1
            predicate Q : { p : x.P }
          }

          # z.1 has been deleted

          schema derived.1 {
            import x.1
            predicate D : { a : string }
              { a = A } where x.P { a = A, b = 3 }
              # now matches on the new field

            # check that we can change the definition of a stored predicate,
            # deleting a type synonym and using a different definition.
            predicate S : string
              stored X where x.P { a = X }
          }

          schema evolved.1 {
            # evolved.P.2 should be able to evolve both the new and old
            # versions of evolved.P.1
            predicate P : enum { A | B }
            predicate Q : { ref : P }
          }

          schema evolved.2 {
            predicate P : enum { A | B }
            predicate Q : { ref : P }
          }

          schema evolved.2 evolves evolved.1

          schema evolved.3 {
            predicate P : enum { A | B | C }
            predicate Q : { ref : P }
          }

          # added evolution
          schema evolved.3 evolves evolved.2

          schema all.1 : x.1, y.1, derived.1, evolved.1, evolved.2, evolved.3 {}
        |]

      schema_index_file_0 = root </> "schema_index_0"
      schema_index_0 = Internal.SchemaIndex
        { schemaIndex_current = Internal.SchemaInstance
          { schemaInstance_versions = Map.fromList [ ("v0", 0) ]
          , schemaInstance_file = "schema0"
          }
        , schemaIndex_older = []
        }

      schema_index_file_1 = root </> "schema_index_1"
      schema_index_1 = Internal.SchemaIndex
        { schemaIndex_current = Internal.SchemaInstance
          { schemaInstance_versions = Map.fromList [ ("v1", 1) ]
          , schemaInstance_file = "schema1"
          }
        , schemaIndex_older = [
            Internal.SchemaInstance
              { schemaInstance_versions = Map.fromList [ ("v0", 1) ]
              , schemaInstance_file = "schema0"
              }
        ]
        }

    writeFile schema_v0_file schema_v0
    writeFile schema_v1_file schema_v1
    saveJSON schema_index_file_0 schema_index_0
    saveJSON schema_index_file_1 schema_index_1

    let dbRoot = root </> "db"
    createDirectory dbRoot

    let
      mkRepo schema hash upd facts =
        withTestEnv [
            setRoot dbRoot,
            setSchemaIndex schema ] $ \env -> do
          let repo = Repo "test" hash
          kickOffTestDB env repo upd
          facts env repo
          completeTestDB env repo
          return repo

      testEnv schema schema_id set act =
        withTestEnv (
          [ setRoot dbRoot,
            setSchemaIndex schema,
            maybe id (setSchemaId . SchemaId) schema_id ] ++ set)
          act

      testQuery name repo schema schema_id query result =
        testQuerySet name [] repo schema schema_id query result

      testQuerySet name set repo schema schema_id query result =
        testEnv schema schema_id set $ \env -> do
          r <- try $ angleQuery env repo query
          case result of
            Just n -> case r :: Either BadQuery UserQueryResults of
              Right UserQueryResults{..} ->
                assertEqual name n (length userQueryResults_facts)
              _ -> assertFailure (name <> ": " <> show r)
            Nothing -> assertBool name $ case r of
              Left{} -> True
              _ -> False

      testQueryFacts name repo schema schema_id query ty results =
        testQueryFactsSet name [] repo schema schema_id query ty results

      testQueryFactsSet name set repo schema schema_id query ty results =
        testEnv schema schema_id set $ \env -> do
          -- get fids
          fids <- do
            UserQueryResults{..} <- userQuery env repo def
              { userQuery_query = query
              , userQuery_options = Just def
                { userQueryOptions_syntax = QuerySyntax_ANGLE
                }
              , userQuery_encodings = [ UserQueryEncoding_bin def ]
              }
            UserQueryEncodedResults_bin b <- return userQueryResults_results
            return $ Map.keys (userQueryResultsBin_facts b)

          -- ask for facts
          r <- try $ userQueryFacts env repo def
            { userQueryFacts_facts =
                [ def { factQuery_id = fromIntegral fid } | fid <- fids ]
            , userQueryFacts_encodings = [ UserQueryEncoding_bin def ]
            }

          -- check that the result is as expected
          decoded <- decodeResults ty userQueryResultsBin_facts r
          case decoded of
            Left err -> assertFailure err
            Right values -> assertEqual name values results

      v0_facts =
        [ mkBatch (PredicateRef "z.R" 1)
            [ [s| { "key" : "abc" } |] ]
        , mkBatch (PredicateRef "x.P" 1)
            [ [s| { "key" : { "a" : "xyz" } } |] ]
        , mkBatch (PredicateRef "y.Q" 1)
            [ [s| { "key" : { "p" : { "key" : { "a" : "xyz" } }}} |] ]
        , mkBatch (PredicateRef "evolved.P" 1)
            [ [s| { "key" : 0 } |] ]
        ]

      v1_facts =
        [ mkBatch (PredicateRef "x.P" 1)
            [ [s| { "key" : { "a" : "xyz", "b": 3 } } |] ]
        , mkBatch (PredicateRef "y.Q" 1)
            [ [s| { "key" : { "p" : { "key" : { "a" : "xyz", "b" : 3 } }}} |] ]
        , mkBatch (PredicateRef "evolved.P" 2)
            [ [s| { "key" : 0 } |] ]
        , mkBatch (PredicateRef "evolved.P" 1)
            [ [s| { "key" : 0 } |] ]
        ]

    -- create a DB using v0
    repo0 <- mkRepo schema_index_file_0 "0" id $ \env repo -> do
      void $ syncWriteJsonBatch env repo v0_facts Nothing

    -- create a DB using v1
    repo1 <- mkRepo schema_index_file_1 "1" id $ \env repo -> do
      void $ syncWriteJsonBatch env repo v1_facts Nothing

    -- create a DB using v0 by setting glean.schema_id
    let set kickOff = kickOff {
          kickOff_properties = HashMap.insert "glean.schema_id" "v0"
            (kickOff_properties kickOff) }
    repo2 <- mkRepo schema_index_file_1 "2" set $ \env repo -> do
      void $ syncWriteJsonBatch env repo v0_facts Nothing

    -- query repo0 with index 1, explicitly ask for schema v0
    testQuery "multi 0a" repo0 schema_index_file_1 (Just "v0")
      "z.R _" (Just 1)
    testQuery "multi 0b" repo0 schema_index_file_1 (Just "v0")
      "x.P _" (Just 1)
    testQuery "multi 0c" repo0 schema_index_file_1 (Just "v0")
      "derived.D _" (Just 1)
    -- only takes into account evolutions in the schema used plus
    -- the schema stored in the db. Given none of the two have
    -- evolutions for evolved.3, this should return no results.
    testQuery "multi 0d" repo0 schema_index_file_1 (Just "v0")
      "evolved.P.3 _" (Just 0)

    -- query repo0 with index 1, explicitly ask for schema v1
    testQuery "multi 2a" repo0 schema_index_file_1 (Just "v1") "z.R _" Nothing
      -- we can't see z.R now
    testQuery "multi 2b" repo0 schema_index_file_1 (Just "v1") "x.P _" (Just 1)
      -- we can see x.P, but we'll get the new version
    testQuery "multi 2c" repo0 schema_index_file_1 (Just "v1")
      "y.Q { p = { b = 0 }}" (Just 1)
      -- this should give us 1 fact, because we're matching on b's default
    testQuery "multi 2d" repo0 schema_index_file_1 (Just "v1")
      "derived.D _" (Just 0)
    -- the evolution in v1 should be taken into account.
    testQuery "multi 2e" repo0 schema_index_file_1 (Just "v1")
      "evolved.P.3 _" (Just 1)

    -- query repo1 with index 1, explicitly ask for schema v0
    testQuery "multi 3a" repo1 schema_index_file_1 (Just "v0") "x.P _" (Just 1)
    testQuery "multi 3b" repo1 schema_index_file_1 (Just "v0") "y.Q _" (Just 1)
    testQuery "multi 3c" repo1 schema_index_file_1 (Just "v0")
      "y.Q { p = { b = 0 }}" Nothing
    testQuery "multi 3d" repo1 schema_index_file_1 (Just "v0")
      "derived.D _" (Just 1)
      -- this should be a type error, P doesn't have the b field
    -- check that the results gets transformed according to the schema_id.
    testQueryFacts "multi 3d" repo1 schema_index_file_1 (Just "v0") "x.P _"
      (RecordTy [FieldDef "a" StringTy])
      [RTS.Tuple [RTS.String "xyz"]]
    -- the evolution in v1 should be taken into account.
    testQuery "multi 3e" repo1 schema_index_file_1 (Just "v0")
      "evolved.P.3 _" (Just 1)

    -- query repo1 with index 1, don't ask for v0
    testQuery "multi 4a" repo1 schema_index_file_1 Nothing "x.P _" (Just 1)
    testQuery "multi 4b" repo1 schema_index_file_1 Nothing "y.Q _" (Just 1)
    testQuery "multi 4c" repo1 schema_index_file_1 Nothing
      "y.Q { p = { b = 3 }}" (Just 1)
    testQuery "multi 4d" repo1 schema_index_file_1 Nothing
      "derived.D _" (Just 1)
    testQueryFacts "multi 4e" repo1 schema_index_file_1 Nothing "x.P _"
      (RecordTy [FieldDef "a" StringTy, FieldDef "b" NatTy ])
      [RTS.Tuple [RTS.String "xyz", RTS.Nat 3]]

    -- query repo1 with index 1, explicitly ask for schema v1
    testQuery "multi 5a" repo1 schema_index_file_1 (Just "v1") "x.P _" (Just 1)
    testQuery "multi 5b" repo1 schema_index_file_1 (Just "v1") "y.Q _" (Just 1)
    testQuery "multi 5c" repo1 schema_index_file_1 (Just "v1")
      "y.Q { p = { b = 3 }}" (Just 1)
    testQuery "multi 5d" repo1 schema_index_file_1 (Just "v1")
      "derived.D _" (Just 1)
    testQueryFacts "multi 4e" repo1 schema_index_file_1 (Just "v1") "x.P _"
      (RecordTy [FieldDef "a" StringTy, FieldDef "b" NatTy ])
      [RTS.Tuple [RTS.String "xyz", RTS.Nat 3]]

    -- if we downgraded the schema, so the repo has a later version
    -- than the global schema, these queries should still work
    testQuery "multi 6a" repo1 schema_index_file_0 Nothing "x.P _" (Just 1)
    testQuery "multi 6b" repo1 schema_index_file_0 Nothing "y.Q _" (Just 1)
    testQuery "multi 6c" repo1 schema_index_file_0 Nothing
      "y.Q { p = { b = 3 }}" Nothing
    testQuery "multi 6d" repo1 schema_index_file_0 Nothing
      "derived.D _" (Just 1)

    -- asking for a schema that doesn't exist, we will default to the
    -- latest global schema if config_strict_query_schema_id = False
    testQuerySet "multi 7a" [disableStrictSchemaId] repo1 schema_index_file_0
      (Just "na") "x.P _" (Just 1)
    testQuerySet "multi 7b" [disableStrictSchemaId] repo1 schema_index_file_0
      (Just "na") "y.Q _" (Just 1)
    testQuerySet "multi 7c" [disableStrictSchemaId] repo1 schema_index_file_0
      (Just "na") "y.Q { p = { b = 3 }}" Nothing
    testQuerySet "multi 7d" [disableStrictSchemaId] repo1 schema_index_file_0
      (Just "na") "derived.D _" (Just 1)

    -- with config_strict_query_schema_id = True, should fail
    when (config_strict_query_schema_id def) $
      testEnv schema_index_file_0 (Just "na") [] $ \env -> do
        r <- try $ angleQuery env repo1 "some query"
        assertBool "multi 7e" $ case r of
          Left UnknownSchemaId{} -> True
          _ -> False

    -- Test that glean.schema_id works. repo2 should be using schema v0
    testQuery "multi 8a" repo2 schema_index_file_1 (Just "v0")
      "z.R _" (Just 1)
    testQuery "multi 8b" repo2 schema_index_file_1 (Just "v0")
      "x.P _" (Just 1)
    testQuery "multi 8c" repo2 schema_index_file_1 (Just "v0")
      "derived.D _" (Just 1)

    -- create a stacked DB on top of v0
    let
      stacked (Thrift.Repo name hash) =
        Thrift.Dependencies_stacked $ Thrift.Stacked name hash Nothing
      set kickOff = kickOff {
          kickOff_dependencies = Just (stacked repo0) }
    repo3 <- mkRepo schema_index_file_1 "3" set $ \env repo -> do
      void $ syncWriteJsonBatch env repo v0_facts Nothing

    testQuery "multi 9a" repo3 schema_index_file_1 Nothing "x.P _" (Just 1)

validateSchemaChanges :: Test
validateSchemaChanges =
    let
      schema_v0 =
        [s|
          schema x.1 {
            predicate P : { a: string }
            predicate Q : { a: string | b: nat }
            predicate R : { a : string } A where P { a = A }
          }

          schema all.1 : x.1 {}
        |]

      schema_v1 =
        [s|
          schema x.1 {
            predicate P : { a: string, b : { c: nat, d: bool } }
          }

          schema all.1 : x.1 {}
        |]

      schema_v2 =
        [s|
          schema x.1 {
            predicate R : string
            predicate P : { a: string, b : R }
          }

          schema all.1 : x.1 {}
        |]

      schema_v3 =
        [s|
          schema x.1 {
            predicate Q : { a: string | b: nat | c: [nat] }
          }

          schema all.1 : x.1 {}
        |]

      schema_v4 =
        [s|
          schema x.1 {
            predicate P : { a: string }
          }

          schema all.1 : x.1 {}
        |]

      schema_v5 =
        [s|
          schema x.1 {
            type T = string
            predicate P : { a: T }
          }

          schema all.1 : x.1 {}
        |]

      schema_v6 =
        [s|
          schema x.1 {
            predicate P : { a: bool }
            predicate Q : { a: string | b: nat }
          }

          schema all.1 : x.1 {}
        |]

      schema_v7 =
        [s|
          schema x.1 {
            predicate P : { a: bool }
          }

          schema y.1 {
            import x.1
            predicate Q : { a: x.P }
          }

          schema all.1 : x.1, y.1 {}
        |]

      schema_v8 =
        [s|
          schema x.1 {
            predicate P : { a: bool }
          }

          schema x.2 {
            predicate P : { a: bool, b: bool }
          }

          schema x.2 evolves x.1

          schema y.1 {
            import x.2
            predicate Q : { a: x.P }
          }

          schema all.1 : x.1, x.2, y.1 {}
        |]

      schema_v9 =
        [s|
          schema x.1 {
            predicate P : { a: string }
            predicate Q : { a: string | b: nat }
            predicate R : { a : nat } A where Q { b = A }
          }

          schema all.1 : x.1 {}
        |]

      schema_v10 =
        [s|
          schema x.1 {
            predicate P : { a: string }
            predicate Q : { a: string | b: nat }
            predicate R : { a : nat } stored A where Q { b = A }
          }

          schema all.1 : x.1 {}
        |]

      validate
        :: String
        -> String
        -> String
        -> (Either SomeException () -> Bool)
        -> Test
      validate lbl a b p = TestLabel lbl $ TestCase $ do
        withSystemTempDirectory "glean-dbtest" $ \root -> do
          let fileA = "schemaA"
              fileB = "schemaB"
          writeFile (root </> fileA) a
          writeFile (root </> fileB) b
          withIndex root fileA fileB $ \index_file -> do
            schema <- parseSchemaIndex index_file
            r <- try $ validateNewSchemaInstance schema
            print r
            assertBool "validate" $ p (r :: Either SomeException ())

      validateOK lbl a b = TestList
        [ validate (lbl <> " (forwards)") a b isRight
        , validate (lbl <> " (backwards)") b a isRight
        ]
      validateFAIL lbl a b = TestList
        [ validate (lbl <> " (forwards)") a b isLeft
        , validate (lbl <> " (backwards)") b a isLeft
        ]
    in
    TestList
      [ validateOK "adding a defaultable field" schema_v0 schema_v1
      , validateFAIL "adding a non-defaultable field" schema_v0 schema_v2
      , validateOK "adding an alternative" schema_v0 schema_v3
      , validateOK "adding a predicate" schema_v0 schema_v4
      , validateOK "adding a type synonym" schema_v0 schema_v5
      , validateFAIL "changing the type of a field of a stored predicate"
          schema_v0 schema_v6
      , validateOK "changing the type of a field of a derived predicate"
          schema_v0 schema_v9
      , validateOK "stored and derived predicates don't conflict between versions"
          schema_v0 schema_v10
      -- this tests only going from schema_v7 to schema_v8 because the
      -- evolves relationship doesn't exist the other way around.
      , validate "change evolved field" schema_v8 schema_v7 isRight
      ]

withIndex :: FilePath -> String -> String -> (FilePath -> IO b) -> IO b
withIndex root a b f = do
  saveJSON schema_index_file schema_index
  f schema_index_file
  where
  schema_index_file = root </> "schema_index_" <> a <> b
  schema_index = Internal.SchemaIndex
    { schemaIndex_current = Internal.SchemaInstance
      { schemaInstance_versions = Map.fromList [ ("v1", 1) ]
      , schemaInstance_file = Text.pack a
      }
    , schemaIndex_older = [
        Internal.SchemaInstance
          { schemaInstance_versions = Map.fromList [ ("v0", 1) ]
          , schemaInstance_file = Text.pack b
          }
    ]
    }

schemaMismatch :: Test
schemaMismatch = TestCase $
    let
      schema_v0 =
        [s|
          schema x.1 {
            predicate P : { a: string }
          }

          schema all.1 : x.1 {}
        |]

      schema_v1 =
        [s|
          schema x.1 {
            predicate P : { a: string, b : nat }
          }

          schema all.1 : x.1 {}
        |]
    in
    withSystemTempDirectory "glean-dbtest" $ \root -> do

      let dbRoot = root </> "db"
      createDirectory dbRoot

      let file0 = "schema0"
          file1 = "schema1"
      writeFile (root </> file0) schema_v0
      writeFile (root </> file1) schema_v1

      withIndex root file0 file1 $ \index_file -> do
        withTestEnv [
          setRoot dbRoot,
          setSchemaIndex index_file ] $ \env -> do
        let repo = Repo "test" "0"
        kickOffTestDB env repo id
        let
          facts =
            [ mkBatch (PredicateRef "x.P" 1)
                [ [s| { "key" : { "a" : "abc" }} |] ]
            ]
        -- don't use syncWriteJsonBatch, it doesn't check the schema ID
        r <- try $ sendJsonBatch env repo facts (Just def {
          sendJsonBatchOptions_schema_id = Just (SchemaId "v0") })
        assertBool "schema mismatch" $ case r of
          Left (e :: SomeException) ->
            "does not match schema ID" `isInfixOf` show e
          _ -> False

main :: IO ()
main = withUnitTest $ testRunner $ TestList
  [
    TestLabel "multiSchemaTest" multiSchemaTest,
    TestLabel "validateSchemaChanges" validateSchemaChanges,
    TestLabel "schemaMismatch" schemaMismatch
  ]