{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}

-- | Parser for our STLC + GADTs, whose AST is defined in "AST".
-- | Base on the parser from the HW2.
module Parse (parse, tryParse) where

import qualified AST                            as A
import           Control.Monad                  (unless, void)
import qualified Control.Monad.Combinators.Expr as Ex
import           Data.Bifunctor                 (Bifunctor (..))
import           Data.Void                      (Void)
import           Text.Megaparsec                (MonadParsec (..), Parsec,
                                                 between, errorBundlePretty,
                                                 many, runParser, some, (<?>),
                                                 (<|>))
import           Text.Megaparsec.Char           (alphaNumChar, lowerChar,
                                                 space1, upperChar)
import qualified Text.Megaparsec.Char.Lexer     as L


-- * Exposed functions

-- | Parse a program; throw an exception over an error
parse :: String -> A.Program
parse = either error id . tryParse

-- | Parse some code 'String' into an 'A.Program' or an error message.
tryParse :: String -> Either String A.Program
tryParse = first errorBundlePretty .
           runParser (pSpace >> pProgram <* eof) "<input>"

pProgram :: Parser A.Program
pProgram = do
  gadts <- many pGadt
  unless (null gadts) $ pToken ";"
  term <- pExpr
  return $ A.Program (gadts, term)

pGadt :: Parser A.GADT
pGadt = do
  pToken "data"
  name <- pConstr
  args <- many pIdent
  pToken "where"
  constructors <- some auxConstr
  return $ A.GADT (name, length args, constructors)
  where
    auxConstr = do
      pToken "|"
      name <- pConstr
      pToken "::"
      typ <- pMAnn
      return (name, typ)

-- | Entry point for expression parser.
pTerm :: Parser A.Term
pTerm = pBody <?> "expression"

-- | Parse expressions at the lowest level of precedence, i.e., lambdas.
pBody :: Parser A.Term
pBody = pLam <|> pLetAnn
 where
  pLam = do
    pToken "\\"
    bs <- some pIdent <?> "lambda binders"
    pToken "."
    body <- pBody <?> "lambda body"
    return $ foldr A.TLam body bs

pLetAnn :: Parser A.Term
pLetAnn = try aux <|> pLet
  where
    aux = do
      pToken "let"
      x <- pIdent <?> "let binder"
      pToken "::"
      ann <- pPAnn <?> "let annotation"
      pToken "="
      m1 <- pExpr <?> "let value"
      pToken "in"
      m2 <- pLet <?> "let body"
      return $ A.TLetAnn x ann m1 m2

pLet :: Parser A.Term
pLet = aux <|> pIf
  where
    aux = do
      pToken "let"
      x <- pIdent <?> "let binder"
      pToken "="
      m1 <- pExpr <?> "let value"
      pToken "in"
      m2 <- pExpr <?> "let body"
      return $ A.TLet x m1 m2

pIf :: Parser A.Term
pIf = aux <|> pCase
  where
    aux = do
      pToken "if"
      t1 <- pExpr <?> "if condition"
      pToken "then"
      t2 <- pExpr <?> "if then"
      pToken "else"
      t3 <- pExpr <?> "if else"
      return $ A.TIf t1 t2 t3

pCase :: Parser A.Term
pCase = aux <|> pApp
  where
    aux = do
      pToken "case"
      t <- pExpr <?> "case expr"
      pToken "of"
      cases <- some pCaseAux <?> "case choices"
      return $ A.TCase t cases
    pCaseAux = do
      pToken "|"
      p <- pPattern <?> "case pattern"
      pToken "=>"
      t <- pExpr <?> "case expr"
      return (p, t)

-- | Parse juxtaposition as application.
pApp :: Parser A.Term
pApp = foldl1 A.TApp <$> some pAtom <?> "term application"

-- | Parse expressions at the highest precedence, including parenthesized terms
pAtom :: Parser A.Term
pAtom = pInt <|> pBool <|> A.TVar <$> auxVar <|> A.TConstr <$> auxConstr <|> pParens pExpr
 where
  auxVar = pIdent <?> "variable"
  auxConstr = pConstr <?> "constructor"

pPAnn :: Parser A.PolyType
pPAnn = do
  pToken "forall"
  vars <- many pIdent <?> "poly ann variables"
  pToken "."
  body <- pMAnn <?> "poly ann body"
  return $ A.PTForall vars body

pMAnn :: Parser A.MonoType
pMAnn = try auxFunc <|> auxConstr
  where
    auxFunc = do
      arg <- auxConstr <?> "mono ann func arg"
      pToken "->"
      res <- pMAnn <?> "mono ann func res"
      return $ A.MTFunc arg res
    auxConstr = try (do
      constr <- pConstr <?> "mono ann type constructor"
      args <- many auxConstr <?> "mono ann type args"
      return $ A.MTConstr constr args ) <|> auxVar
    auxVar = A.MTVar <$> aux <|> pParens pMAnn
      where
        aux = pIdent <?> "mono ann type variable"

pPattern :: Parser A.Pattern
pPattern = auxAnn
  where
    auxAnn = try (do
      p <- auxConstr <?> "pattern body"
      pToken "::"
      ann <- pMAnn <?> "pattern annotation"
      return $ A.PAnn p ann) <|> auxConstr
    auxConstr = try (do
      c <- pConstr <?> "pattern constructor"
      ps <- many auxVar <?> "pattern constructor args"
      return $ A.PConstr c ps) <|> auxVar
    auxVar = (do
      v <- pIdent <?> "pattern variable"
      return $ A.PVar v) <|> pParens pPattern

-- * Megaparsec boilerplate and helpers

-- | Parsing monad.
type Parser = Parsec Void String

-- | keywords
keywords :: [String]
keywords = ["let", "in", "case", "of", "forall", "data", "where", "if", "then", "else"]
cKeywords :: [String]
cKeywords = []
-- cKeywords = ["Int", "Bool"]

-- | Parse an identifier, possibly surrounded by spaces, excluding keywords
pIdent :: Parser String
pIdent = (L.lexeme pSpace . try) (p >>= check)
  where
    p       = (:) <$> lowerChar <*> many alphaNumChar
    check x = if x `elem` keywords
              then fail $ "keyword " ++ show x ++ " cannot be an identifier"
              else return x

pConstr :: Parser String
pConstr = (L.lexeme pSpace . try) (p >>= check)
  where
    p       = (:) <$> upperChar <*> many alphaNumChar
    check x = if x `elem` cKeywords
              then fail $ "keyword " ++ show x ++ " cannot be a constructor"
              else return x

pInt :: Parser A.Term
pInt = do
  n <- L.lexeme pSpace L.decimal
  return $ A.TLit (A.LInt n)

pBool :: Parser A.Term
pBool = (pToken "true" >> return (A.TLit (A.LBool True)))
  <|> (pToken "false" >> return (A.TLit (A.LBool False)))

-- | Consume a token defined by a string, possibly surrounded by spaces
pToken :: String -> Parser ()
pToken = void . L.symbol pSpace

-- | Parse some element surrounded by parentheses.
pParens :: Parser a -> Parser a
pParens = between (pToken "(") (pToken ")")

-- | Consumes whitespace and comments.
pSpace :: Parser ()
pSpace = label "whitespace" $ L.space
    space1
    (L.skipLineComment "--")
    (L.skipBlockCommentNested "{-" "-}")

pExpr :: Parser A.Term
pExpr = Ex.makeExprParser ( pTerm <?> "term") table <?> "expression"

table :: [[Ex.Operator Parser A.Term]]
table = [
    [ binary A.Mul $ pToken "*"],
    [ binary A.Add $ pToken "+"
    , binary A.Sub $ pToken "-"],
    [ binary A.Equal $ pToken "=="]
  ]
    where
      binary op symP = Ex.InfixL $ A.TBinop op <$ symP
