module Datalog.Parser
    ( parseTerm
    , parseAtom
    , parseRule
    , parseQuery
    , parseProgram
    )
  where

import Datalog.Syntax
import Data.Text (Text)
import Data.Text qualified as T
import Data.Void (Void)
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L

type Parser = Parsec Void Text

lexeme :: (MonadParsec e Text m) => m a -> m a
lexeme = L.lexeme whitespace

symbol :: (MonadParsec e Text m) => Text -> m Text
symbol = L.symbol whitespace

whitespace :: (MonadParsec e Text m) => m ()
whitespace = L.space
  space1
  (L.skipLineComment "--")
  (L.skipBlockComment "{-" "-}")

parens :: (MonadParsec e Text m) => m a -> m a
parens = between (symbol "(") (symbol ")")

comma, period :: (MonadParsec e Text m) => m ()
comma = () <$ symbol ","
period = () <$ symbol "."

parseCon :: (MonadParsec e Text m) => m Term
parseCon = Con . ConId . T.pack <$> lexeme (liftA2 (:) (numberChar <|> lowerChar) (many alphaNumChar))

parseVar :: (MonadParsec e Text m) => m Term
parseVar = Var . VarId . T.pack <$> lexeme (liftA2 (:) upperChar (many alphaNumChar))

parseTerm :: Parser Term
parseTerm = parseVar <|> parseCon

parseAtom :: Parser Atom
parseAtom = do
  rel <- RelId . T.pack <$> lexeme (liftA2 (:) lowerChar (many alphaNumChar))
  args <- parens (parseTerm `sepBy` comma)
  return (Atom rel args)

parseQuery :: Parser [Atom]
parseQuery = parseAtom `sepBy` comma

parseFact :: Parser Rule
parseFact = do
  headAtom <- parseAtom
  period
  return (headAtom :- [])

parseRule :: Parser Rule
parseRule = try parseFact <|> do
  headAtom <- parseAtom <* symbol ":-"
  bodyAtoms <- parseQuery
  period
  return (headAtom :- bodyAtoms)

parseProgram :: Parser Program
parseProgram = Program <$> many parseRule