module ProgramParser(parseProgram) where

import Control.Monad
import Data.Char
import Text.Parsec
import qualified Text.Parsec.Expr as E

import AST
import Defs


type Parser = Parsec String ()


parseProgram :: String -> Either String Program
parseProgram s = case parse pProgram "" s of
    Left err -> Left $ show err
    Right p -> Right p

pProgram :: Parser Program
pProgram = do
    pWhiteComment
    decls <- many pDecl
    eof
    return $ Program (lefts decls) (rights decls)

pDecl :: Parser (Either DVar DFunc)
pDecl = (Right <$> pDFunc) <|> (Left <$> pDVar)

pDFunc :: Parser DFunc
pDFunc = do
    symbol "func"
    rt <- (Just <$> pType) <|> return Nothing
    n <- pName
    symbol "("
    args <- sepBy pTypeAndName (symbol ",")
    symbol ")"
    body <- pBlock
    return $ DFunc rt n args body

pDVar :: Parser DVar
pDVar = do
    t <- pType
    n <- pName
    symbol ":="
    e <- pExpression
    symbol ";"
    return $ DVar t n e

pTypeAndName :: Parser (Type, Name)
pTypeAndName = (,) <$> pType <*> pName

pType :: Parser Type
pType = do
    t <- pBasicType
    (do
        symbol "["
        msz <- optionMaybe pInteger
        symbol "]"
        return $ TArr t msz) <|> return t

pBasicType :: Parser Type
pBasicType = (symbol "int" >> return TInt) <|> (symbol "char" >> return TChar)

pBlock :: Parser Block
pBlock = do
    symbol "{"
    body <- many pStatement
    symbol "}"
    return $ Block body

pStatement :: Parser Statement
pStatement = pSIf <|> pSWhile <|> pSReturn <|> pSDecl <|> pSAs <|> pSExpr

pSDecl :: Parser Statement
pSDecl = do
    (t, n) <- try $ do
        t <- pType
        n <- pName
        symbol ":="
        return (t, n)
    e <- pExpression
    symbol ";"
    return $ SDecl t n e

pSAs :: Parser Statement
pSAs = do
    n <- try $ pAsExpression <* symbol "="
    e <- pExpression
    symbol ";"
    return $ SAs n e

pSIf :: Parser Statement
pSIf = do
    symbol "if"
    symbol "("
    cond <- pExpression
    symbol ")"
    bl1 <- pBlock
    bl2 <- try (symbol "else" >> pBlock) <|> return (Block [])
    return $ SIf cond bl1 bl2

pSWhile :: Parser Statement
pSWhile = do
    symbol "while"
    symbol "("
    cond <- pExpression
    symbol ")"
    bl <- pBlock
    return $ SWhile cond bl

pSReturn :: Parser Statement
pSReturn = do
    symbol "return"
    SReturn <$> ((symbol ";" >> return Nothing) <|>
                 ((Just <$> pExpression) <* symbol ";"))

pSExpr :: Parser Statement
pSExpr = do
    e <- pExpression
    symbol ";"
    return $ SExpr e

pExpression :: Parser Expression
pExpression = E.buildExpressionParser optable litparser
  where
    optable =
        [[E.Infix (symbol "**" >> return (mkEBin BOPow)) E.AssocRight],
         [E.Infix (symbol "*" >> return (mkEBin BOMul)) E.AssocLeft,
          E.Infix (symbol "/" >> return (mkEBin BODiv)) E.AssocLeft,
          E.Infix (symbol "%" >> return (mkEBin BOMod)) E.AssocLeft],
         [E.Infix (symbol "+" >> return (mkEBin BOAdd)) E.AssocLeft,
          E.Infix (symbol "-" >> return (mkEBin BOSub)) E.AssocLeft],
         [E.Infix (symbol ">=" >> return (mkEBin BOGeq)) E.AssocNone,
          E.Infix (symbol "<=" >> return (mkEBin BOLeq)) E.AssocNone,
          E.Infix (symbol ">" >> return (mkEBin BOGt)) E.AssocNone,
          E.Infix (symbol "<" >> return (mkEBin BOLt)) E.AssocNone,
          E.Infix (symbol "==" >> return (mkEBin BOEq)) E.AssocNone,
          E.Infix (symbol "!=" >> return (mkEBin BONeq)) E.AssocNone],
         [E.Infix (symbol "&&" >> return (mkEBin BOAnd)) E.AssocLeft],
         [E.Infix (symbol "||" >> return (mkEBin BOOr)) E.AssocLeft]]

    mkEBin :: BinaryOp -> Expression -> Expression -> Expression
    mkEBin bo a b = EBin bo a b Nothing

    mkELit :: Literal -> Expression
    mkELit l = ELit l Nothing

    litparser :: Parser Expression
    litparser = do
        pops <- many pPrefixOp
        e <- pParenExpr <|> pENew <|> pCastExpr <|> (mkELit <$> pLiteral)
        subs <- many $ between (symbol "[") (symbol "]") pExpression
        let e' = foldl (\ex sub -> ESubscript ex sub Nothing) e subs
            e'' = foldl (\ex pop -> EUn pop ex Nothing) e' pops
        return e''

pAsExpression :: Parser AsExpression
pAsExpression = do
    n <- pName
    subs <- many $ between (symbol "[") (symbol "]") pExpression
    return $ foldl (\ae expr -> AESubscript ae expr Nothing) (AEVar n Nothing) subs

pPrefixOp :: Parser UnaryOp
pPrefixOp = (symbol "!" >> return UONot) <|>
            (symbol "-" >> return UONeg)

pParenExpr :: Parser Expression
pParenExpr = do
    symbol "("
    e <- pExpression
    symbol ")"
    return e

pCastExpr :: Parser Expression
pCastExpr = do
    t <- try $ pType <* symbol "("
    e <- pExpression
    symbol ")"
    return $ ECast t e

pENew :: Parser Expression
pENew = do
    symbol "new"
    t <- pBasicType
    symbol "["
    e <- pExpression
    symbol "]"
    return $ ENew t e

pLiteral :: Parser Literal
pLiteral = (LInt <$> pInteger) <|> (LChar <$> pCharLit) <|> pLCall <|> (LVar <$> pName)

pCharLit :: Parser Char
pCharLit = do
    void $ char '\''
    c <- pStringChar
    void $ char '\''
    pWhiteComment
    return c

pStringChar :: Parser Char
pStringChar =
    (char '\\' >> ((char 'n' >> return '\n') <|>
                   (char 'r' >> return '\r') <|>
                   (char 't' >> return '\t') <|>
                   (char '0' >> return '\0') <|>
                   (char 'x' >> pHexDigit >>= \a -> pHexDigit >>= \b -> return (chr $ 16 * a + b)))) <|>
    anyToken
  where
    pHexDigit :: Parser Int
    pHexDigit = (subtract 48 . fromEnum <$> digit)
                    <|> ((+ (10 - 97)) . ord <$> oneOf "abcdef")
                    <|> ((+ (10 - 65)) . ord <$> oneOf "ABCDEF")

pLCall :: Parser Literal
pLCall = do
    n <- try $ pName <* symbol "("
    al <- sepBy pExpression (symbol ",")
    symbol ")"
    return $ LCall n al


pName :: Parser Name
pName = do
    c0 <- satisfy (\c -> isAlpha c || c == '_')
    cr <- many $ satisfy (\c -> isAlpha c || isDigit c || c == '_')
    pWhiteComment
    return $ c0 : cr

pInteger :: Parser Integer
pInteger = read <$> many1 (satisfy isDigit) <* pWhiteComment


symbol :: String -> Parser ()
symbol "" = error "symbol \"\""
symbol s = try $ do
    void $ string s
    when (isAlpha (last s)) $ void $ notFollowedBy (satisfy isAlpha)
    when (isDigit (last s)) $ void $ notFollowedBy (satisfy isDigit)
    pWhiteComment

pWhiteComment :: Parser ()
pWhiteComment = void $ pWhite >> endBy pComment pWhite

pWhite :: Parser ()
pWhite = void $ many (oneOf " \t\n")

pComment :: Parser ()
pComment = pLineComment <|> pBlockComment

pLineComment :: Parser ()
pLineComment = do
    void $ try $ string "//"
    void $ many (satisfy (/= '\n'))
    eof <|> void (char '\n')

pBlockComment :: Parser ()
pBlockComment = do
    void $ try $ string "/*"
    void $ sepEndBy (manyTill anyToken (lookAhead $ try (string "/*") <|> try (string "*/")))
                    pBlockComment
    void $ string "*/"

lefts :: [Either a b] -> [a]
lefts = foldr (\e l -> either (:l) (const l) e) []

rights :: [Either a b] -> [b]
rights = foldr (\e l -> either (const l) (:l) e) []