{-# LANGUAGE NamedFieldPuns #-}

module Parser (parseDef, parseDefEmpty, parseExpr, parseProgram, handleFile) where

import Check
import Control.Monad.Except
import Data.List (elemIndex, foldl, foldl1)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Errors (Error (..))
import Eval
import Expr (Expr (..), incIndices)
import Preprocessor
import Relude.Extra.Lens
import Text.Megaparsec (ParsecT, ShowErrorComponent (..), between, choice, chunk, customFailure, errorBundlePretty, label, runParserT, try)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

data InnerState = IS {_binders :: [Text], _env :: Env}

bindsL :: Lens' InnerState [Text]
bindsL = lens _binders setter
  where
    setter (IS{_env}) new = IS{_env, _binders = new}

envL :: Lens' InnerState Env
envL = lens _env setter
  where
    setter (IS{_binders}) new = IS{_env = new, _binders}

newtype TypeError = TE Error
    deriving (Eq, Ord)

type Parser = ParsecT TypeError Text (State InnerState)

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"]

pIdentifier :: Parser Text
pIdentifier = try $ label "identifier" $ lexeme $ do
    firstChar <- letterChar <|> char '_'
    rest <- many $ alphaNumChar <|> char '_'
    let ident = T.pack (firstChar : rest)
    when (ident `elem` keywords) $
        fail $
            "Reserved word: " ++ T.unpack ident
    pure ident

pVar :: Parser Expr
pVar = label "variable" $ lexeme $ do
    name <- pIdentifier
    binders <- view bindsL <$> get
    pure $ Var name . fromIntegral <$> elemIndex name binders ?: Free name

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

pParamGroup :: Parser [(Text, Expr)]
pParamGroup = lexeme $ label "parameter group" $ between (char '(') (char ')') $ do
    idents <- some pIdentifier
    eat ":"
    ty <- pExpr
    modify $ over bindsL $ flip (foldl $ flip (:)) idents
    pure $ zip idents (iterate incIndices ty)

pSomeParams :: Parser [(Text, Expr)]
pSomeParams = lexeme $ concat <$> some pParamGroup

pManyParams :: Parser [(Text, Expr)]
pManyParams = lexeme $ concat <$> many pParamGroup

withBinders :: Parser a -> Parser a
withBinders parser = do
    oldBinders <- view bindsL <$> get
    result <- parser
    modify $ set bindsL oldBinders
    pure result

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

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

pBinding :: Parser (Text, Expr)
pBinding = lexeme $ label "binding" $ do
    env <- get
    eat "("
    ident <- pIdentifier
    params <- pManyParams
    eat ":="
    value <- pExpr
    eat ")"
    put env
    modify $ over bindsL (ident :)
    pure (ident, foldr (uncurry Abs) value params)

pLet :: Parser Expr
pLet = lexeme $ label "let expression" $ withBinders $ do
    eat "let"
    bindings <- some pBinding
    eat "in"
    body <- try pExpr
    eat "end"
    pure $ foldr (uncurry Let) body bindings

pArrow :: Parser Expr
pArrow = lexeme $ label "->" $ do
    a <- pAppTerm
    _ <- defChoice $ "->" :| ["→"]
    Pi "" a . incIndices <$> pExpr

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

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

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

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

checkAscription :: Text -> Expr -> Maybe Expr -> Parser ()
checkAscription ident value massert = do
    env <- get
    ty <- either (customFailure . TE) pure (checkType (env ^. envL) value)
    case massert of
        Nothing -> updateStateDefinition ident ty value
        Just assert -> do
            equiv <- either (customFailure . TE) pure (checkBeta (env ^. envL) ty assert)
            unless equiv $ customFailure $ TE $ NotEquivalent ty assert value
            updateStateDefinition ident assert value

updateStateDefinition :: Text -> Expr -> Expr -> Parser ()
updateStateDefinition ident ty value = do
    env <- get
    when (M.member ident (env ^. envL)) (customFailure $ TE $ DuplicateDefinition ident)
    modify $ over envL $ M.insert ident $ makeDef ty value

pAxiom :: Text -> Maybe Expr -> Parser ()
pAxiom ident Nothing = customFailure $ TE $ PNMissingType ident
pAxiom ident (Just ascription) = do
    eat "axiom"
    eat ";"
    updateStateDefinition ident ascription (Axiom ident)

pBody :: [(Text, Expr)] -> Text -> Maybe Expr -> Parser ()
pBody params ident ascription = do
    value <- flip (foldr (uncurry Abs)) params <$> pExpr
    checkAscription ident value ascription
    eat ";"

pDef :: Parser ()
pDef = lexeme $ label "definition" $ withBinders $ do
    skipSpace
    ident <- pIdentifier
    params <- pManyParams
    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> pAscription
    eat ":="
    choice [pAxiom ident ascription, pBody params ident ascription]

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

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

pExpr :: Parser Expr
pExpr = lexeme $ try pArrow <|> pAppTerm

pAscription :: Parser (Maybe Expr)
pAscription = lexeme $ optional $ try $ defChoice (pure ":") >> label "type" pExpr

pProgram :: Parser Env
pProgram = lexeme $ skipSpace >> many pDef >> _env <$> get

emptyBinders :: Env -> InnerState
emptyBinders env = IS{_binders = [], _env = env}

parseDef :: Text -> State Env (Either String ())
parseDef input = do
    env <- get
    let (output, IS{_env}) = runState (runParserT pDef "" input) (emptyBinders env)
    put _env
    pure $ first errorBundlePretty output

parseExpr :: Env -> Text -> Either String Expr
parseExpr env input = first errorBundlePretty $ evalState (runParserT pExpr "" input) $ emptyBinders env

parseDefEmpty :: Env -> Text -> (Either String (), Env)
parseDefEmpty env input = runState (parseDef input) env

parseProgram :: Env -> Text -> Either String Env
parseProgram initial input = first errorBundlePretty $ evalState (runParserT pProgram "" input) $ emptyBinders initial

handleFile :: Env -> String -> ExceptT String IO Env
handleFile initial filename = (toString `withExceptT` runPreprocessor filename) >>= liftEither . parseProgram initial