module Parser where

import Ast
import Data.Maybe (isJust)
import Lexer
import Text.Parsec
import qualified Text.Parsec.Expr as E
import Text.Parsec.String (Parser)

binary s f assoc = E.Infix (reservedOp s >> return (BinOp f)) assoc

opTable =
  [ [ binary "*" Mul E.AssocLeft,
      binary "/" Div E.AssocLeft
    ],
    [ binary "+" Plus E.AssocLeft,
      binary "-" Minus E.AssocLeft
    ]
  ]

expr :: Parser Expr
expr = E.buildExpressionParser opTable factor

number :: Parser Expr
number = do
  pos <- getPosition
  n <- integer
  return $ Num n pos

variable :: Parser Expr
variable = do
  pos <- getPosition
  var <- identifier
  return $ Var var pos

typ :: Parser Type
typ = do
  p <- identifier
  return $ Type p

typFn :: Parser Type
typFn = do
  p <- parens $ commaSep typ
  r <- optionMaybe $ do
    reserved "->"
    typ
  return $ FuncType p r

type' :: Parser Type
type' = do
  try typFn
    <|> typ <?> "type"

arg :: Parser (String, Type, SourcePos)
arg = do
  pos <- getPosition
  i <- identifier
  _ <- colonSep <?> "\":\" before type"
  t <- type' <?> "type"
  return $ (i, t, pos)

fdef :: Ident -> Bool -> Bool -> SourcePos -> Parser Expr
fdef ident priv anon pos = do
  args <- parens $ commaSep arg
  rtyp <-
    optionMaybe
      ( do
          _ <- colonSep <?> "\":\" before type"
          rtyp <- type' <?> "return type"
          return $ rtyp
      )
  body <- braces $ many statement
  return $ Func ident args rtyp body priv anon pos

function :: Parser Expr
function = do
  pos <- getPosition
  priv <- optionMaybe $ reserved "private"
  reserved "def"
  ident <- identifier
  fdef ident (isJust priv) False pos

lambdaId :: SourcePos -> Ident
lambdaId s =
  "lambda" ++ "@" ++ show (sourceLine s) ++ "," ++ show (sourceColumn s)

lambda :: Parser Expr
lambda = do
  pos <- getPosition
  fdef (lambdaId pos) True True pos

return' :: Parser Expr
return' = do
  pos <- getPosition
  reserved "return"
  value <- optionMaybe expr
  reservedOp ";"
  return $ Return value pos

call :: Parser Expr
call = do
  pos <- getPosition
  ident <- try lambda <|> variable
  args <- parens $ commaSep expr
  return $ Call ident args pos

factor :: Parser Expr
factor =
  try number
    <|> try call
    <|> try lambda
    <|> try variable
    <|> parens expr

exprStmt :: Parser Expr
exprStmt = do
  e <- expr <|> factor
  reservedOp ";"
  return $ e

statement :: Parser Expr
statement = do
  try exprStmt
    <|> return'

module' :: Parser Expr
module' = do
  pos <- getPosition
  reserved "module"
  ident <- identifier
  return $ Module ident pos

program :: Parser [Expr]
program = do
  m <- module'
  n <-
    many $ do
      try function
        -- TODO: variable decl
        <|> statement <?> "statement"
  return $ [m] ++ n

parse :: Parser [Expr]
parse = program

parseFromFile p fname = do
  input <- readFile fname
  return (runParser p () fname input)