{-# LANGUAGE NamedFieldPuns #-}

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

import Check
import Control.Monad
import Control.Monad.Except
import Control.Monad.State.Strict
import Data.Bifunctor (first)
import Data.List (elemIndex)
import Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as M
import Data.Text (Text)
import qualified Data.Text as T
import Errors (Error (..))
import Eval
import Expr (Expr (..), incIndices)
import Preprocessor
import Text.Megaparsec hiding (State)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

data TypeDef = TD {_ident :: Text, _type :: Expr}
data DefinitionLine = DL {_td :: TypeDef, _body :: Expr} | PN TypeDef

data InnerState = IS {_binds :: [TypeDef], _env :: Env}

newtype TypeError = TE Error
    deriving (Eq, Ord, Show)

type Parser = ParsecT TypeError Text (State InnerState)

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

bindsToIS :: ([TypeDef] -> [TypeDef]) -> InnerState -> InnerState
bindsToIS f x@(IS{_binds}) = x{_binds = f _binds}

modifyEnv :: (Env -> Env) -> InnerState -> InnerState
modifyEnv f x@(IS{_env}) = x{_env = f _env}

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

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

pIdentifier :: Parser Text
pIdentifier = label "identifier" $ lexeme $ do
    firstChar <- letterChar <|> char '_'
    rest <- many $ alphaNumChar <|> char '_'
    return $ T.pack (firstChar : rest)

pVar :: Parser Expr
pVar = label "variable" $ lexeme $ do
    var <- pIdentifier
    binders <- map _ident . _binds <$> get
    pure $ case elemIndex var binders of
        Just i -> Var (fromIntegral i) var
        Nothing -> Free var

pPN :: Parser ()
pPN = label "primitive notion" $ lexeme $ defChoice $ pure "axiom"

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

pParamGroup :: Parser [(Text, Expr)]
pParamGroup = lexeme $ label "parameter group" $ between (char '(') (char ')') $ do
    idents <- some pIdentifier
    _ <- defChoice $ pure ":"
    ty <- pExpr
    modify $ bindsToIS $ flip (foldl $ flip (:)) (map (\idt -> TD{_ident = idt, _type = ty}) 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

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

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

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 $ Star <$ defChoice (pure "*")

pSquare :: Parser Expr
pSquare = lexeme $ Square <$ defChoice ("□" :| ["[]"])

checkAscription :: Text -> Expr -> Maybe Expr -> Parser DefinitionLine
checkAscription ident value massert = do
    IS{_env} <- get
    case (checkType _env value, massert) of
        (Left err, _) -> customFailure $ TE err
        (Right ty, Nothing) -> pure DL{_td = TD{_ident = ident, _type = ty}, _body = value}
        (Right ty, Just assert) -> case checkBeta _env ty assert of
            Left err -> customFailure $ TE err
            Right equiv -> do
                unless equiv $ customFailure $ TE $ NotEquivalent ty assert value
                pure DL{_td = TD{_ident = ident, _type = assert}, _body = value}

updateStateDefinition :: DefinitionLine -> Parser ()
updateStateDefinition DL{_td, _body} = do
    env <- get
    let ident = _ident _td
    when (M.member ident (_env env)) (customFailure $ TE $ DuplicateDefinition ident)
    modify $
        modifyEnv
            (M.insert ident EL{_ty = _type _td, _val = _body})
updateStateDefinition (PN TD{_type, _ident}) = do
    env <- get
    when (M.member _ident (_env env)) (customFailure $ TE $ DuplicateDefinition _ident)
    modify $
        modifyEnv
            (M.insert _ident EL{_ty = _type, _val = Axiom _ident})

pDef :: Parser DefinitionLine
pDef = lexeme $ label "definition" $ do
    skipSpace
    ident <- pIdentifier
    params <- pManyParams
    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> pAscription
    _ <- defChoice $ pure ":="
    choice
        [ do
            _ <- pPN
            _ <- defChoice $ pure ";"
            case ascription of
                Just ty -> pure $ PN TD{_ident = ident, _type = ty}
                Nothing -> customFailure $ TE $ PNMissingType ident
        , do
            value <- flip (foldr (uncurry Abs)) params <$> pExpr
            res <- checkAscription ident value ascription
            _ <- defChoice $ pure ";"
            pure res
        ]

pDefUpdate :: Parser ()
pDefUpdate = pDef >>= updateStateDefinition

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

pAppTerm :: Parser Expr
pAppTerm = lexeme $ pLAbs <|> pPAbs <|> 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 pDefUpdate >> _env <$> get

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

parseExpr :: Env -> Text -> Either String Expr
parseExpr env input = first errorBundlePretty $ evalState (runParserT pExpr "" input) IS{_binds = [], _env = 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) IS{_binds = [], _env = initial}

handleFile :: Env -> String -> ExceptT String IO Env
handleFile initial filename = do
    text <- show `withExceptT` preprocess filename
    liftEither $ parseProgram initial text