{-# LANGUAGE NamedFieldPuns #-}

module Parser (pAll) where

import Check
import Control.Monad
import Control.Monad.State.Strict
import Data.Bifunctor (first)
import Data.Functor.Identity
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 Eval
import Expr (Expr (..), incIndices)
import Text.Megaparsec hiding (State)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

data InnerState = IS {_binds :: [Text], _defs :: Env}

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

type Parser = ParsecT TypeError Text (State InnerState)

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

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

defsToIS :: (Env -> Env) -> InnerState -> InnerState
defsToIS f x@(IS{_defs}) = x{_defs = f _defs}

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 <- _binds <$> get
    pure $ case elemIndex var binders of
        Just i -> Var (fromIntegral i) var
        Nothing -> Free var

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 $ ":" :| []
    ty <- pExpr
    modify $ bindsToIS $ flip (foldl $ flip (:)) idents
    pure $ zip idents (iterate incIndices ty)

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

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

pPAbs :: Parser Expr
pPAbs = lexeme $ label "Π-abstraction" $ do
    _ <- defChoice $ "∏" :| ["Pi", "forall", "∀"]
    params <- pParams
    _ <- defChoice $ "." :| [","]
    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 = foldl1 App <$> some pTerm

pStar :: Parser Expr
pStar = Star <$ defChoice ("*" :| [])

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

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

pAppTerm :: Parser Expr
pAppTerm = lexeme $ pLAbs <|> pPAbs <|> pApp

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

pDef :: Parser ()
pDef = lexeme $ label "definition" $ do
    ident <- pIdentifier
    _ <- defChoice $ ":=" :| []
    value <- pExpr
    _ <- defChoice $ ";" :| []
    foo <- get
    let ty = checkType (_defs foo) [] value
    case ty of
        Left err -> customFailure $ TE err
        Right _ -> modify $ defsToIS $ M.insert ident value

pProgram :: Parser ()
pProgram = void $ many pDef

pAll :: Text -> Either String ()
pAll input = first errorBundlePretty $ fst $ runIdentity $ runStateT (runParserT pProgram "" input) $ IS{_binds = [], _defs = M.empty}