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

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

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

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 (Rule () headAtom [])

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

-- parseProgram :: Parser Program
parseProgram :: Parser (Program' SrcLoc)
parseProgram = do
    -- annotateSrcLoc $
    s <- getSourcePos
    c <- many parseRule'
    e <- getSourcePos
    -- Program _ <$> many parseRule
    pure $ Program (SrcLoc s e) c

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

parseThingWithSub :: (Parser c) -> (SrcLoc -> c -> f SrcLoc) -> Parser (f SrcLoc)
parseThingWithSub parseSub f = do
    -- annotateSrcLoc $
    s <- getSourcePos
    -- c <- many parseRule'
    c <- parseSub
    e <- getSourcePos
    -- Program _ <$> many parseRule
    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).
    """