{-# HLINT ignore "Redundant flip" #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

module Datalog.DatalogDB where

import Control.Exception.Base
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Datalog.DatalogParser (Literal (..), Term (..), Statement)
import Text.Megaparsec (ParseErrorBundle)
import Data.Void

data Relation = Relation
  { _name :: RelationId
  , _arity :: Int
  , _tuples :: Set [Constant]
  , _rules :: [RelationRule]
  }
  deriving (Show, Eq)

data RelationRule = RelationRule
  { headVariables :: [Text]
  , bodyElements :: [RuleBodyElement]
  }
  deriving (Show, Eq)

data RuleBodyElement = RuleBodyElement
  { _subRelationId :: RelationId
  , _ruleElements :: [RuleElement]
  }
  deriving (Show, Eq)

data RuleElement -- entry occurring in a head or body relation - constant or variable
  = RuleElementConstant Constant
  | RuleElementVariable Int
  deriving (Show, Eq)

data DatalogDBException
  = CannotParseStatementException Text (ParseErrorBundle Text Void)
  | NonFactException Text Statement
  | NonRuleException Text Statement
  | NonQueryException Text Statement
  | BadArityException Text Int
  | VariableLookupException Text [Text]
  | UnexpectedConstantException Constant
  deriving (Show)

instance Exception DatalogDBException

-- Our constants will be the terms of the Datalog grammar - ints/variables/symbols
type Constant = Term
type RelationId = Text

class DatalogDB db where
  emptyDB :: db
  lookupRelation :: db -> Text -> Maybe Relation
  insertRelation :: db -> Relation -> db
  addConstants :: db -> Set Constant -> db

lookupRelationArity :: DatalogDB db => Text -> db -> Int -> Set [Constant] -> Relation
lookupRelationArity relationName db newArity tuples =
  case lookupRelation db relationName of
    Nothing -> Relation relationName newArity tuples []
    Just relation ->
      if _arity relation == newArity then
          let newTuples = Set.union tuples $ _tuples relation
           in relation { _tuples = newTuples }
        else throw $ BadArityException relationName newArity

addFact :: (DatalogDB db) => Literal -> db -> db
addFact (Literal neg relationName terms) db =
  insertRelation (addConstants db extraConstants) newRelation
  where
  newArity = length terms
  newRelation = lookupRelationArity relationName db newArity (Set.singleton terms)
  extraConstants = Set.fromList terms