{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MultilineStrings #-}

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

import Data.Text (Text)
import Data.Text qualified as T
import Data.Void (Void)
import Datalog.Syntax hiding (Atom, Program, Rule, Term)
import Datalog.Syntax (Atom' (..), Program' (..), Rule' (..), Term' (..))
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L
import Text.Pretty.Simple

type Parser = Parsec Void Text

type Atom = Atom' SrcLoc
type Term = Term' SrcLoc
type Rule = Rule' SrcLoc
type Program = Program' SrcLoc

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 :: Parser Term
parseCon = parseThingWithSub Con $ ConId . T.pack <$> lexeme (liftA2 (:) (numberChar <|> lowerChar) (many alphaNumChar))

parseVar :: Parser Term
parseVar = parseThingWithSub Var $ VarId . T.pack <$> lexeme (liftA2 (:) upperChar (many alphaNumChar))

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

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

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

parseRule :: Parser Rule
parseRule = parseThingWithSub (\loc (a, as) -> Rule loc a as) $
        try rule1 <|> rule2
  where
    rule1 = do
        headAtom <- parseAtom
        period
        return (headAtom, [])
    rule2 = do
        headAtom <- parseAtom <* symbol ":-"
        bodyAtoms <- parseQuery
        period
        return (headAtom, bodyAtoms)

parseProgram :: Parser Program
parseProgram = parseThingWithSub Program (many parseRule)

parseThingWithSub :: (SrcLoc -> c -> f SrcLoc) -> Parser c -> Parser (f SrcLoc)
parseThingWithSub f parseSub = do
    s <- getSourcePos
    c <- parseSub
    e <- getSourcePos
    pure $ f (SrcLoc s e) c

annotateSrcLoc :: (Functor f) => Parser (f a) -> Parser (f SrcLoc)
annotateSrcLoc p = do
    s <- getSourcePos
    res <- p
    f <- getSourcePos
    pure (SrcLoc s f <$ res)

data SrcLoc = SrcLoc
    { start :: SourcePos
    , end :: SourcePos
    }
    deriving (Show)

test = do
    let r = runParser parseProgram "???" prog
    pPrint @IO r

prog =
    """
    odd(X,Y) :- r(X,Y).
    odd(X,Y) :- even(X,Z), r(Z,Y).
    even(X,Y) :- odd(X,Z), r(Z,Y).

    r(0,1).
    r(1,2).
    r(2,3).
    r(3,4).
    r(4,5).

    r(X,Y) :- r(Y,X).
    """