{-# LANGUAGE TupleSections #-}

module Parser where

import Control.Monad.Except
import Data.Char
import Data.List (foldl1)
import qualified Data.Text as T
import Errors (Error (..))
import IR
import Preprocessor
import Text.Megaparsec (Parsec, ShowErrorComponent (..), between, choice, chunk, errorBundlePretty, label, runParser, satisfy, try)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

newtype TypeError = TE Error
    deriving (Eq, Ord)

type Parser = Parsec TypeError Text

instance ShowErrorComponent TypeError where
    showErrorComponent (TE e) = toString e

skipSpace :: Parser ()
skipSpace =
    L.space
        space1
        (L.skipLineComment "--")
        (L.skipBlockCommentNested "[*" "*]")

lexeme :: Parser a -> Parser a
lexeme = L.lexeme skipSpace

eat :: Text -> Parser ()
eat = void . lexeme . chunk

keywords :: [Text]
keywords = ["forall", "let", "in", "end", "fun", "def", "axiom"]

reservedChars :: [Char]
reservedChars = "();:"

pIdentifier :: Parser Text
pIdentifier = try $ label "identifier" $ lexeme $ do
    chars <- many $ satisfy isAllowed
    let ident = T.pack chars
    when (ident `elem` keywords) $ fail $ "Reserved word: " ++ T.unpack ident
    pure ident
  where
    isAllowed c = isLetter c || isNumber c || c == '_' || isSymbol c && c `notElem` reservedChars

pVar :: Parser IRExpr
pVar = label "variable" $ lexeme $ Var <$> pIdentifier

defChoice :: NonEmpty Text -> Parser ()
defChoice options = lexeme $ label (T.unpack $ head options) $ void $ choice $ fmap chunk options

pParamGroup :: Parser [Param]
pParamGroup = lexeme $ label "parameter group" $ between (char '(') (char ')') $ do
    idents <- some pIdentifier
    eat ":"
    ty <- pIRExpr
    pure $ map (,ty) idents

pSomeParams :: Parser [Param]
pSomeParams = lexeme $ concat <$> some pParamGroup

pManyParams :: Parser [Param]
pManyParams = lexeme $ concat <$> many pParamGroup

pLAbs :: Parser IRExpr
pLAbs = lexeme $ label "λ-abstraction" $ do
    _ <- defChoice $ "λ" :| ["fun"]
    params <- pSomeParams
    _ <- defChoice $ "=>" :| ["⇒"]
    body <- pIRExpr
    pure $ foldr (uncurry Abs) body params

pPAbs :: Parser IRExpr
pPAbs = lexeme $ label "Π-abstraction" $ do
    _ <- defChoice $ "∏" :| ["forall", "∀"]
    params <- pSomeParams
    eat ","
    body <- pIRExpr
    pure $ foldr (uncurry Pi) body params

pBinding :: Parser (Text, Maybe IRExpr, IRExpr)
pBinding = lexeme $ label "binding" $ do
    eat "("
    ident <- pIdentifier
    params <- pManyParams
    ascription <- optional pAscription
    eat ":="
    value <- pIRExpr
    eat ")"
    pure
        ( ident
        , flip (foldr (uncurry Pi)) params <$> ascription
        , foldr (uncurry Abs) value params
        )

pLet :: Parser IRExpr
pLet = lexeme $ label "let expression" $ do
    eat "let"
    bindings <- some pBinding
    eat "in"
    body <- try pIRExpr
    eat "end"
    pure $ foldr letTuple body bindings
  where
    letTuple :: (Text, Maybe IRExpr, IRExpr) -> IRExpr -> IRExpr
    letTuple (name, ascription, value) = Let name ascription value

pArrow :: Parser IRExpr
pArrow = lexeme $ label "->" $ do
    a <- pAppTerm
    _ <- defChoice $ "->" :| ["→"]
    Pi "" a <$> pIRExpr

pApp :: Parser IRExpr
pApp = lexeme $ foldl1 App <$> some pTerm

pStar :: Parser IRExpr
pStar = lexeme $ Level 0 <$ eat "*"

pNumSort :: Parser IRExpr
pNumSort = lexeme $ label "sort" $ eat "*" >> Level <$> L.decimal

pSort :: Parser IRExpr
pSort = try pNumSort <|> pStar

pAxiom :: Parser IRDef
pAxiom = lexeme $ label "axiom" $ do
    skipSpace
    eat "def"
    ident <- pIdentifier
    params <- pManyParams
    ascription <- fmap (flip (foldr (uncurry Pi)) params) pAscription
    eat ";"
    pure $ Axiom ident ascription

pDef :: Parser IRDef
pDef = lexeme $ label "definition" $ do
    skipSpace
    eat "def"
    ident <- pIdentifier
    params <- pManyParams
    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> optional pAscription
    eat ":="
    body <- pIRExpr
    eat ";"
    pure $ Def ident ascription $ foldr (uncurry Abs) body params

pIRDef :: Parser IRDef
pIRDef = pDef <|> pAxiom

pTerm :: Parser IRExpr
pTerm =
    lexeme $
        label "term" $
            choice [between (char '(') (char ')') pIRExpr, pSort, pVar]

pAppTerm :: Parser IRExpr
pAppTerm = lexeme $ choice [pLAbs, pPAbs, pLet, pApp]

pIRExpr :: Parser IRExpr
pIRExpr = lexeme $ try pArrow <|> pAppTerm

pAscription :: Parser IRExpr
pAscription = lexeme $ try $ eat ":" >> label "type" pIRExpr

pIRProgram :: Parser IRProgram
pIRProgram = many pIRDef

parserWrapper :: Parser a -> String -> Text -> Either String a
parserWrapper p filename input = first errorBundlePretty $ runParser p filename input

parseProgram :: String -> Text -> Either String IRProgram
parseProgram = parserWrapper pIRProgram

parseDef :: String -> Text -> Either String IRDef
parseDef = parserWrapper pIRDef

parseExpr :: String -> Text -> Either String IRExpr
parseExpr = parserWrapper pIRExpr

handleFile :: String -> ExceptT String IO IRProgram
handleFile filename = toString `withExceptT` runPreprocessor filename >>= hoistEither . parseProgram filename