{-
  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 Angle.ArrayTest (main) where

import Data.List
import Data.Text (Text)
import Data.Word
import Test.HUnit

import TestRunner
import Util.String.Quasi

import Glean.Init
import Glean.Query.Thrift as Thrift
import qualified Glean.Schema.GleanTest.Types as Glean.Test
import Glean.Types

import TestDB

main :: IO ()
main = withUnitTest $ withDbTests $ \dbTestCase -> testRunner $ TestList
  [ TestLabel "array" $ angleArray dbTestCase id
  , TestLabel "array/page" $ angleArray dbTestCase (limit 1)
  ]

angleArray :: (WithDB () -> Test) -> (forall a . Query a -> Query a) -> Test
angleArray dbTestCase modify = TestList
  [ TestLabel "generators" $ angleArrayGenerator dbTestCase modify
  , TestLabel "prefix" $ angleArrayPrefix dbTestCase modify
  ]

angleArrayGenerator
  :: (WithDB () -> Test)
  -> (forall a . Query a -> Query a) -> Test
angleArrayGenerator dbTestCase modify = TestList
  [ TestLabel "array of pred" $
    dbTestCase $ \env repo -> do
    -- fetch all elements of an array
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
      [s|
        glean.test.Predicate { array_of_pred = Arr };
        Arr [..]
      |]
    print results
    assertEqual "angle - array generator 1" 2 (length results)

  , TestLabel "array of nat" $
    dbTestCase $ \env repo -> do

    -- match on elements of an array
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
      [s|
        P where
        P = glean.test.Predicate { array_of_nat = Arr };
        3 = Arr [..]  # any P with a 3 in array_of_nat
     |]
    print results
    assertEqual "angle - array generator 2" 1 (length results)

  , TestLabel "array of pred with match" $
    dbTestCase $ \env repo -> do
      -- test that a generator on the left gets compiled correctly
    results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
      [s|
        P where
        P = glean.test.Predicate { array_of_pred = Arr };
        { nat = 42 } = Arr [..]  # any P with a nat = 4
      |]
    print results
    assertEqual "angle - array generator 3" 1 (length results)

  , TestLabel "array of records" $
    dbTestCase $ \env repo -> do
    results <- runQuery_ env repo $ modify $
      angleData @(Text, Nat)
     [s|
        [ { "a",1 }, { "b",2 } ] [..]
      |]
    print results
    assertEqual "angle - array generator 4"
      [ ("a", Nat 1), ("b", Nat 2) ] results
  , TestLabel "array of bytes" $
    dbTestCase $ \env repo -> do
    results <- runQuery_ env repo $ modify $
      angleData @Byte
      [s| [1 : byte, 255][..] |]
    print results
    assertEqual "angle - array generator 5"
      (sort [Byte 1, Byte (fromIntegral (255 :: Word8))]) (sort results)
  ]

angleArrayPrefix
  :: (WithDB () -> Test)
  -> (forall a . Query a -> Query a) -> Test
angleArrayPrefix dbTestCase modify = TestList
  [ TestLabel "nat" $ TestList
    [ TestLabel "nested" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_nat = [3,4, ..] }
          |]
        assertEqual "angle - array prefix" 1 (length results)
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_nat = [3, ..] }
          |]
        assertEqual "angle - array prefix" 1 (length results)
    , TestLabel "flat" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_nat = A };
            [3, ..] = A
          |]
        assertEqual "angle - array prefix" 1 (length results)
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_nat = A };
            [3,4, ..] = A
          |]
        assertEqual "angle - array prefix" 1 (length results)
    ]
  , TestLabel "pred" $ TestList
    [ TestLabel "nested" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_pred = [glean.test.Predicate _, ..] }
          |]
        assertEqual "angle - array prefix" 2 (length results)
    , TestLabel "flat" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_pred = A };
            [_, ..] = A
          |]
        assertEqual "angle - array prefix" 2 (length results)
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_pred = A };
            [glean.test.Predicate _, ..] = A
          |]
        assertEqual "angle - array prefix" 2 (length results)
    ]
  , TestLabel "string" $ TestList
    [ TestLabel "nested" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_string = ["abba","baba", ..] }
          |]
        assertEqual "angle - array prefix" 2 (length results)
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_string = ["abba", ..] }
          |]
        assertEqual "angle - array prefix" 2 (length results)
    ]
  , TestLabel "nat and string" $ TestList
    [ TestLabel "nested" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
          [s|
            P where
            P = glean.test.Predicate { array_of_nat = [3, ..], array_of_string = ["abba", ..] }
          |]
        assertEqual "angle - array prefix" 1 (length results)
    ]

  -- regression test for a bug in MatchArrayPrefix handling
  , TestLabel "buildTerm" $ dbTestCase $ \env repo -> do
        results <- runQuery_ env repo $ modify $ angleData @[Text]
          [s|
            X where X = ["a","b"]; X = ["a",..]
          |]
        assertEqual "angle - array prefix" 1 (length results)
  ]