module Parser where

import Control.Monad
import Control.Monad.State.Strict
import Data.Functor.Identity
import Data.List (elemIndex)
import Data.List.NonEmpty (NonEmpty((:|)))
import qualified Data.List.NonEmpty as NE
import Expr (Expr (..))
import Text.Megaparsec hiding (State)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L
import Data.Bifunctor (first)

type InnerState = [String]

data CustomErrors = UnboundVariable String [String] deriving (Eq, Ord, Show)

instance ShowErrorComponent CustomErrors where
  showErrorComponent (UnboundVariable var bound) =
    "Unbound variable: " ++ var ++ ". Did you mean one of: " ++ unwords bound ++ "?"

type Parser = ParsecT CustomErrors String (State InnerState)

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

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

pIdentifier :: Parser String
pIdentifier = label "identifier" $ lexeme $ do
  firstChar <- letterChar <|> char '_'
  rest <- many $ alphaNumChar <|> char '_'
  return $ firstChar : rest

pVar :: Parser Expr
pVar = label "variable" $ lexeme $ do
  var <- pIdentifier
  binders <- get
  case elemIndex var binders of
    Just i -> return $ Var $ fromIntegral i
    Nothing -> customFailure $ UnboundVariable var binders

defChoice :: NE.NonEmpty String -> Parser ()
defChoice options = lexeme $ label labelText $ void $ choice $ NE.map string options
  where labelText = NE.head options

pLAbs :: Parser Expr
pLAbs = lexeme $ label "λ-abstraction" $ do
  _ <- defChoice $ "λ" :| ["lambda"]
  ident <- pIdentifier
  _ <- defChoice $ ":" :| []
  ty <- pExpr
  _ <- defChoice $ "." :| []
  modify (ident :)
  Abs ty <$> pExpr

pPAbs :: Parser Expr
pPAbs = lexeme $ label "Π-abstraction" $ do
  _ <- defChoice $ "∏" :| ["Pi"]
  ident <- pIdentifier
  _ <- defChoice $ ":" :| []
  ty <- pExpr
  _ <- defChoice $ "." :| []
  modify (ident :)
  Pi ty <$> 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
        ]

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

pAll :: String -> Either String Expr
pAll input = first errorBundlePretty $ fst $ runIdentity $ runStateT (runParserT pExpr "" input) []