module Datalog.Syntax where

import Data.Char (isUpper)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Set qualified as Set
import Data.Set (Set)

newtype ConId = ConId Text
  deriving (Eq, Ord, Show)

newtype VarId = VarId Text
  deriving (Eq, Ord, Show)

newtype RelId = RelId Text
  deriving (Eq, Ord, Show)

data Term = Con ConId | Var VarId
  deriving (Eq, Ord, Show)

con :: Text -> Term
con = Con . ConId

var :: Text -> Term
var = Var . VarId

term :: Text -> Term
term t = if not (T.null t) && isUpper (T.head t) then var t else con t

data Atom = Atom RelId [Term]
  deriving (Eq, Ord, Show)

atom :: Text -> [Text] -> Atom
atom relName args = Atom (RelId relName) (map term args)

data Rule = Atom :- [Atom]
  deriving (Eq, Ord, Show)

data Program = Program [Rule]
  deriving (Eq, Ord, Show)

class HasConstants a where
  constants :: a -> Set ConId

instance HasConstants Term where
  constants t = case t of
    Con x -> Set.singleton x
    Var _ -> Set.empty

instance HasConstants a => HasConstants [a] where
  constants xs = Set.unions (map constants xs)

instance HasConstants Atom where
  constants (Atom _ ts) = constants ts

instance HasConstants Rule where
  constants (h :- b) = Set.union (constants h) (constants b)

instance HasConstants Program where
  constants (Program rs) = constants rs