module Parser(parseProgram) where

import Control.Monad
import Data.Char
import Data.Functor.Identity
import Data.Maybe
import qualified Data.Map.Strict as Map
import Text.Parsec
import qualified Text.Parsec.Expr as E

import AST


type Parser = Parsec String ()


(<<) :: (Monad m) => m a -> m b -> m a
(<<) = (<*)


parseProgram :: String -> String -> Either ParseError Program
parseProgram source fname = parse pProgram fname source

-- parse' :: Parser a -> String -> Either ParseError a
-- parse' p s = parse p "" s

pProgram :: Parser Program
pProgram = pWhiteComment >> (Program <$> many1 pDeclaration)

pDeclaration :: Parser Declaration
pDeclaration = pDecTypedef <|> pDecExtern <|> do
    t <- pTypeVoid <|> pType
    n <- pName
    if t == TypeVoid
      then pDecFunction' t n
      else pDecFunction' t n <|> pDecVariable' t n

pDecTypedef :: Parser Declaration
pDecTypedef = do
    symbol "type"
    n <- pName
    symbol "="
    t <- pType
    symbol ";"
    return $ DecTypedef t n

pDecExtern :: Parser Declaration
pDecExtern = do
    symbol "extern"
    t <- pType
    n <- pName
    symbol ";"
    return $ DecExtern t n

pDecFunction' :: Type -> Name -> Parser Declaration
pDecFunction' t n = do
    symbol "("
    a <- sepBy ((,) <$> pType <*> pName) (symbol ",")
    symbol ")"
    b <- pBlock
    return $ DecFunction t n a b

pDecVariable' :: Type -> Name -> Parser Declaration
pDecVariable' t n = do
    e <- (Just <$> (symbol "=" >> pExpression)) <|> return Nothing
    symbol ";"
    return $ DecVariable t n e

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


exprTable :: (E.OperatorTable String () Identity) Expression
exprTable =
    [[prefix "-" Negate,
      prefix "!" Not,
      prefix "~" Invert,
      prefix "*" Dereference,
      prefix "&" Address],
     [binary "*" Times E.AssocLeft,
      binary "/" Divide E.AssocLeft,
      binary "%" Modulo E.AssocLeft],
     [binary "+" Plus E.AssocLeft,
      binary "-" Minus E.AssocLeft],
     [binary ">=" GEqual E.AssocNone,
      binary "<=" LEqual E.AssocNone,
      binary ">" Greater E.AssocNone,
      binary "<" Less E.AssocNone],
     [binary "==" Equal E.AssocNone,
      binary "!=" Unequal E.AssocNone],
     [binary "&&" BoolAnd E.AssocLeft,
      binary "||" BoolOr E.AssocLeft]]
  where
    binary name op assoc = E.Infix (exBinOp_ op <$ symbol name) assoc
    prefix name op = E.Prefix (exUnOp_ op <$ symbol name)

pExpression :: Parser Expression
pExpression = E.buildExpressionParser exprTable (pExCast <|> pExLit)

pExCast :: Parser Expression
pExCast = do
    symbol "cast"
    symbol "("
    t <- pType
    symbol ")"
    e <- pParenExpr
    return $ ExCast t e

pExLit :: Parser Expression
pExLit = do
    litex <- (exLit_ <$> pLiteral) <|> pParenExpr
    option litex $ do
        symbol "["
        arg <- pExpression
        symbol "]"
        return $ exBinOp_ Index litex arg

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

pLiteral :: Parser Literal
pLiteral = (LitFloat <$> pFloat) <|> pLitInt <|> (LitInt <$> pCharStr)
           <|> (LitString <$> pString) <|> try pLitCall <|> (LitVar <$> pName)

pLitInt :: Parser Literal
pLitInt = do
    i <- pInteger
    liftM (maybe (LitInt i) (const $ LitUInt i)) $ optionMaybe (symbol "U")

pLitCall :: Parser Literal
pLitCall = do
    n <- pName
    symbol "("
    a <- sepBy pExpression (symbol ",")
    symbol ")"
    return $ LitCall n a


pStatement :: Parser Statement
pStatement = pStEmpty <|> pStIf <|> pStWhile <|> pStReturn <|> pStBlock
             <|> try pStAssignment <|> try pStVarDeclaration <|> pStExpr

pStEmpty :: Parser Statement
pStEmpty = symbol ";" >> return StEmpty

pStBlock :: Parser Statement
pStBlock = StBlock <$> pBlock

pStVarDeclaration :: Parser Statement
pStVarDeclaration = do
    t <- pType
    n <- pName
    e <- optionMaybe (symbol "=" >> pExpression)
    symbol ";"
    return $ StVarDeclaration t n e

pStExpr :: Parser Statement
pStExpr = (StExpr <$> pExpression) << symbol ";"

pStAssignment :: Parser Statement
pStAssignment = do
    n <- pName
    symbol "="
    e <- pExpression
    symbol ";"
    return $ StAssignment n e

pStIf :: Parser Statement
pStIf = do
    symbol "if"
    symbol "("
    c <- pExpression
    symbol ")"
    t <- pStatement
    e <- (symbol "else" >> pStatement) <|> return StEmpty
    return $ StIf c t e

pStWhile :: Parser Statement
pStWhile = do
    symbol "while"
    symbol "("
    c <- pExpression
    symbol ")"
    b <- pStatement
    return $ StWhile c b

pStReturn :: Parser Statement
pStReturn = do
    symbol "return"
    (symbol ";" >> return (StReturn Nothing)) <|> do
        e <- pExpression
        symbol ";"
        return $ StReturn (Just e)


primitiveTypes :: Map.Map String Type
primitiveTypes = Map.fromList
    [("i8", TypeInt 8), ("i16", TypeInt 16), ("i32", TypeInt 32), ("i64", TypeInt 64),
     ("u1", TypeUInt 1), ("u8", TypeUInt 8), ("u16", TypeUInt 16), ("u32", TypeUInt 32), ("u64", TypeUInt 64),
     ("float", TypeFloat), ("double", TypeDouble)]

findPrimType :: String -> Type
findPrimType s = fromJust $ Map.lookup s primitiveTypes

pType :: Parser Type
pType = pPrimType <|> pTypePtr <|> pTypeFunc <|> pTypeName

pPrimType :: Parser Type
pPrimType = findPrimType <$> choice (map typeParser $ Map.keys primitiveTypes)
    where typeParser t = try $ do
              void $ string t
              void $ lookAhead $ satisfy (\c -> not (isAlphaNum c) && c /= '_')
              pWhiteComment
              return t

pTypeVoid :: Parser Type
pTypeVoid = symbol "void" >> return TypeVoid

pTypePtr :: Parser Type
pTypePtr = do
    symbol "ptr"
    symbol "("
    t <- pType
    symbol ")"
    return $ TypePtr t

pTypeFunc :: Parser Type
pTypeFunc = do
    symbol "func"
    r <- pTypeVoid <|> pType
    symbol "("
    a <- sepBy pType (symbol ",")
    symbol ")"
    return $ TypeFunc r a

pTypeName :: Parser Type
pTypeName = TypeName <$> pName


pName :: Parser Name
pName = ((:) <$> pFirstChar <*> many pOtherChar) << pWhiteComment
    where pFirstChar = satisfy (isAlpha .||. (=='_'))
          pOtherChar = satisfy (isAlpha .||. isDigit .||. (=='_'))

pInteger :: Parser Integer
pInteger = (read <$> many1 digit) << pWhiteComment

pFloat :: Parser Double
pFloat = try $ do
    pre <- many1 digit
    post <- choice [pExponent,
                    (do
                        void $ char '.'
                        s <- many1 digit
                        ex <- option "" pExponent
                        return $ '.' : s ++ ex)]
    pWhiteComment
    return $ read $ pre ++ post
  where
    pExponent = do
        c <- choice [char 'e', char 'E']
        pm <- option "" $ choice [string "+", string "-"]
        val <- many1 digit
        return $ c : pm ++ val

pString :: Parser String
pString = do
    void $ char '"'
    s <- many (pEscape <|> satisfy (/='"'))
    symbol "\""
    return s

pCharStr :: Parser Integer
pCharStr = do
    void $ char '\''
    c <- pEscape <|> satisfy (/='\'')
    symbol "'"
    return $ fromIntegral (ord c)

pEscape :: Parser Char
pEscape = char '\\' >> (pEscapeQuote <|> pEscapeN <|> pEscapeR <|> pEscapeT <|> pEscapeHex)
  where
    pEscapeQuote, pEscapeN, pEscapeR, pEscapeT :: Parser Char
    pEscapeQuote = ('"' <$ char '"') <|> ('\'' <$ char '\'')
    pEscapeN = '\n' <$ char 'n'
    pEscapeR = '\r' <$ char 'r'
    pEscapeT = '\t' <$ char 't'

pEscapeHex :: Parser Char
pEscapeHex = do
    void $ char 'x'
    c1 <- pHexChar
    c2 <- pHexChar
    return $ chr $ 16 * c1 + c2
  where
    pHexChar :: Parser Int
    pHexChar = (liftM (\c -> ord c - ord '0') digit)
               <|> (liftM (\c -> ord c - ord 'a' + 10) (oneOf "abcdef"))
               <|> (liftM (\c -> ord c - ord 'A' + 10) (oneOf "ABCDEF"))

symbol :: String -> Parser ()
symbol s = do
    void $ try (string s)
    when (isAlphaNum (last s)) $ notFollowedBy alphaNum
    pWhiteComment

pWhiteComment :: Parser ()
pWhiteComment = sepBy pWhitespace pComment >> return ()

pWhitespace :: Parser ()
pWhitespace = many (oneOf " \t\n") >> return ()

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

pLineComment :: Parser ()
pLineComment = try (string "//") >> manyTill anyChar (char '\n') >> return ()

pBlockComment :: Parser ()
pBlockComment = try (string "/*") >> manyTill anyChar (try (string "*/")) >> return ()


infixr 2 .||.
(.||.) :: (a -> Bool) -> (a -> Bool) -> a -> Bool
f .||. g = \x -> f x || g x