{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# HLINT ignore "Use const" #-}
{-# HLINT ignore "Unused LANGUAGE pragma" #-}
{-# HLINT ignore "Avoid lambda" #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

module Test.Datalog.InMemoryDBSpec where

import Data.Map qualified as Map
import Data.Set qualified as Set
import Datalog.DatalogParser
import Datalog.InMemoryDB
import Datalog.InMemoryDB qualified as InMemoryDB
import Test.Hspec
import Datalog.DatalogDB
import Data.Text

spec :: Spec
spec = do
  describe "InMemoryDB operations" do
    it "can ingest facts into relations & a universe" $ do
      let db =
            InMemoryDB.withFacts
              [ "parent(\"alice\", \"bob\")."
              , "parent(\"bob\", \"carol\")."
              ]
      allConstants db
        `shouldBe` Set.fromList (Sym <$> ["alice", "bob", "carol"])
      _relations db
        `shouldBe` Map.fromList
          [
            ( "parent"
            , Relation "parent" 2 (Set.fromList $ Sym <<$>> [["alice", "bob"], ["bob", "carol"]]) []
            )
          ]
      relationNames db `shouldBe` [ "parent" ]
    it "can ingest facts and rules" do
      let db =
            InMemoryDB.withFactsAndRules
              [ "parent(\"alice\", \"bob\")."
              , "parent(\"bob\", \"carol\")."
              ]
              [ "ancestor(X,Y) :- parent(X,Z), ancestor(Z,Y)."
              , "ancestor(X,Y) :- parent(X,Y)."
              ]
          parentRelation =
            Relation
              { _name = "parent"
              , _arity = 2
              , _tuples =
                  Set.fromList $
                    Sym <<$>> [["alice", "bob"], ["bob", "carol"]]
              , _rules = []
              }
          ancestorRule = RelationRule
              { headVariables = ["X", "Y", "Z"]
              , bodyElements =
                  [ ruleBody "parent"   [0, 2]
                  , ruleBody "ancestor" [2, 1]
                  ]
              }
          ancestorRule2 = RelationRule
              { headVariables = ["X", "Y"]
              , bodyElements =
                  [ ruleBody "parent" [0, 1] ]
              }
          ancestorRelation =
            Relation
              { _arity = 2
              , _name = "ancestor"
              , _tuples = Set.empty
              , _rules = [ancestorRule, ancestorRule2]
              }

      allConstants db
        `shouldBe` Set.fromList (Sym <$> ["alice", "bob", "carol"])

      _relations db
        `shouldBe` Map.fromList
          [ ("ancestor", ancestorRelation)
          , ("parent", parentRelation)
          ]
      Set.fromList (relationNames db) `shouldBe` Set.fromList [ "parent", "ancestor" ]

    it "can ingest facts and rules with constants" do
      let db =
            InMemoryDB.withFactsAndRules
              []
              ["ancestor(X,\"patriarch\") :- ."]
          ancestorRule =
            RelationRule
              { headVariables = ["X"]
              , bodyElements = []
              }
          ancestorRelation =
            Relation
              { _arity = 2
              , _name = "ancestor"
              , _tuples = Set.empty
              , _rules = [ancestorRule]
              }
      _relations db
        `shouldBe` Map.singleton "ancestor" ancestorRelation

      allConstants db
        `shouldBe` Set.fromList (Sym <$> ["patriarch"])
      Set.fromList (relationNames db) `shouldBe` Set.fromList [ "ancestor" ]

    it "can ingest facts and rules with duplicate head entries" do
      let db =
            InMemoryDB.withFactsAndRules
              []
              ["equivalent(Q,Q) :- ."]
          equivalentRule =
            RelationRule
              { headVariables = ["Q"]
              , bodyElements = []
              }
          equivalentRelation =
            Relation
              { _arity = 2
              , _name = "equivalent"
              , _tuples = Set.empty
              , _rules = [equivalentRule]
              }
      _relations db
        `shouldBe` Map.singleton "equivalent" equivalentRelation

      allConstants db
        `shouldBe` Set.empty
      relationNames db `shouldBe` [ "equivalent" ]

    it "can ingest a theory of equivalence relations" do
      let db =
            InMemoryDB.withFactsAndRules
              []
              [ "equivalent(Q,Q) :- ."
              , "equivalent(R,Q) :- equivalent(Q,R)."
              , "equivalent(Q,S) :- equivalent(Q,R), equivalent(R,S)."
              ]
          rule1 =
            RelationRule
              { headVariables = ["Q"]
              , bodyElements = []
              }
          rule2 =
            RelationRule
              { headVariables = ["R", "Q"]
              , bodyElements =
                  [ ruleBody "equivalent" [1, 0] ]
              }
          rule3 =
            RelationRule
              { headVariables = ["Q", "S", "R"]
              , bodyElements =
                  [ ruleBody "equivalent" [0, 2]
                  , ruleBody "equivalent" [2, 1]
                  ]
              }
          equivalentRelation =
            Relation
              { _arity = 2
              , _name = "equivalent"
              , _tuples = Set.empty
              , _rules = [rule1, rule2, rule3]
              }
      _relations db
        `shouldBe` Map.singleton "equivalent" equivalentRelation

      allConstants db
        `shouldBe` Set.empty
      Set.fromList (relationNames db) `shouldBe` Set.fromList [ "equivalent" ]

(<<$>>) :: (Functor f1, Functor f2) => (a -> b) -> f1 (f2 a) -> f1 (f2 b)
(<<$>>) = fmap fmap fmap

(<<<$>>>) :: (Functor f1, Functor f2, Functor f3) => (a -> b) -> f1 (f2 (f3 a)) -> f1 (f2 (f3 b))
(<<<$>>>) = fmap fmap fmap fmap fmap fmap fmap fmap

ruleBody :: Text -> [Int] -> RuleBodyElement
ruleBody subRelationId indices =
  RuleBodyElement
    { _subRelationId = subRelationId
    , _ruleElements =
      RuleElementVariable <$> indices
    }