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 <|> do
    t <- pType
    n <- pName
    pDecFunction' t n <|> pDecVariable' t n

pDecTypedef :: Parser Declaration
pDecTypedef = do
    symbol "type"
    n <- pName
    symbol "="
    t <- pType
    symbol ";"
    return $ DecTypedef 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 ">" Greater E.AssocNone,
      binary "<" Less E.AssocNone,
      binary ">=" GEqual E.AssocNone,
      binary "<=" LEqual 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 pExLit

pExLit :: Parser Expression
pExLit = ExLit <$> pLiteral

pLiteral :: Parser Literal
pLiteral = (LitInt <$> pInteger) <|> (LitString <$> pString)
           <|> try pLitCall <|> (LitVar <$> pName)

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 <|> 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"
    e <- pExpression
    symbol ";"
    return $ StReturn e


primitiveTypes :: Map.Map String Type
primitiveTypes = Map.fromList
    [("i8", TypeInt 8), ("i16", TypeInt 16), ("i32", TypeInt 32), ("i64", TypeInt 64),
     ("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 <|> pPtrType <|> pTypeName

pPrimType :: Parser Type
pPrimType = findPrimType <$> choice (map symbol' $ Map.keys primitiveTypes)

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

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


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

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

pString :: Parser String
pString = do
    void $ char '"'
    s <- many (pEscape <|> anyChar)
    symbol "\""
    return s
  where
    pEscape :: Parser Char
    pEscape = char '\\' >> (pEscapeN <|> pEscapeR <|> pEscapeT <|> pEscapeHex)

    pEscapeN, pEscapeR, pEscapeT :: Parser 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

    pHexChar :: Parser Int
    pHexChar = (liftM (\c -> ord c - ord '0') (satisfy isDigit))
               <|> (liftM (\c -> ord c - ord 'a' + 10) (oneOf "abcdef"))
               <|> (liftM (\c -> ord c - ord 'A' + 10) (oneOf "ABCDEF"))

symbol :: String -> Parser ()
symbol s = try (string s) >> pWhiteComment

symbol' :: String -> Parser String
symbol' s = try (string s) << 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