{-# OPTIONS_GHC -fno-warn-unused-do-bind -fno-warn-missing-signatures #-}

module Mathista.Parser (parse_str) where
import Data.List.Split hiding (sepBy)
import Data.String.Utils
import Text.Regex
import Text.Parsec
import Text.Parsec.Expr
import Text.Parsec.String
import Text.Parsec.Language (emptyDef)
import qualified Text.Parsec.Token as Token
import Mathista.AST


lexer :: Token.TokenParser ()
lexer  = Token.makeTokenParser emptyDef
              { Token.commentLine    = "//"
              , Token.reservedNames  = ["if", "then", "else", "elseif", "function",
                                        "end", "for", "break", "continue", "return",
                                        "while", "to", "step"]
              , Token.caseSensitive  = False
              , Token.identStart     = letter
              , Token.identLetter    = alphaNum <|> char '_' <|> char '\''
              , Token.reservedOpNames = ["+", "-", "*", "/", "==", "!=", ">", "<",
                                         ">=", "<=", "and", "or", "not", "=", "@"]
              }

semi        = Token.semi       lexer
float       = Token.float      lexer
symbol      = Token.symbol     lexer
parens      = Token.parens     lexer
comma       = Token.comma      lexer
identifier  = Token.identifier lexer
integer     = Token.integer    lexer
reserved    = Token.reserved   lexer
reservedOp  = Token.reservedOp lexer
whiteSpace  = Token.whiteSpace lexer


--
--  Var
--
var :: Parser Var
var = do
    name <- identifier
    is <- optionMaybe indexes
    return (name, is)
    where indexes = do
              char '['
              is <- index `sepBy` comma
              char ']'
              return is
          index = do
              st <- choice [expr, return $ Number $ 0]
              end <- optionMaybe $ do
                         char ':'
                         end <- choice [expr, return $ Number $ -1]
                         return end
              return (st, end)
    

--
--  Statements
--
stmt :: Parser Stmt
stmt =  try assignStmt
    <|> debugprintStmt
    <|> ifStmt
    <|> whileStmt
    <|> forStmt
    <|> breakStmt
    <|> continueStmt
    <|> returnStmt
    <|> funcDeclStmt
    <|> exprStmt
    <?> "stmt"

stmts :: Parser [Stmt]
stmts = many stmt

assignStmt :: Parser Stmt
assignStmt = do
    lhs <- var `sepBy` comma
    whiteSpace
    reservedOp "="
    e <- expr
    semi
    return $ Assign lhs e

debugprintStmt :: Parser Stmt
debugprintStmt = do
    reservedOp "@"
    e <- expr
    semi
    return $ ExprStmt (FuncCall "debugprint" [e])

ifStmt :: Parser Stmt
ifStmt = do
        between (reserved "if") (reserved "end") inIfBlock
    where
        inIfBlock = do
            ifBlock <- block
            elseifBlocks <- many elseifStmt
            elseBlock <- optionMaybe elseStmt
            return $ If (ifBlock:elseifBlocks) elseBlock
        elseifStmt = reserved "elseif" >> block
        elseStmt   = reserved "else"   >> stmts
        block = do
            e <- expr
            whiteSpace
            reserved "then"
            st <- stmts
            return (e, st)

whileStmt :: Parser Stmt
whileStmt = do
    reserved "while"
    e  <- expr
    whiteSpace
    st <- between (reserved "then") (reserved "end") stmts
    return $ While e st

forStmt :: Parser Stmt
forStmt = do
    reserved "for"
    v <- identifier
    reservedOp "="
    from <- integer
    char ':'
    to <- integer
    whiteSpace
    st <- between (reserved "then") (reserved "end") stmts
    return $ For v from to st

breakStmt :: Parser Stmt
breakStmt = do
    reserved "break"
    return $ Break

continueStmt :: Parser Stmt
continueStmt = do
    reserved "continue"
    return $ Continue

returnStmt :: Parser Stmt
returnStmt = do
    reserved "return"
    e <- expr `sepBy` comma
    return $ Return e

exprStmt :: Parser Stmt
exprStmt = do
    e <- optionMaybe expr
    semi
    case e of
        Just x  -> return $ ExprStmt x
        Nothing -> return $ DoNothing   -- semicolon only lines

arg :: Parser (Id, Id)
arg = do
    name <- identifier
    symbol ":"
    type_ <- identifier
    return (name, type_)

funcDeclStmt = do
    reserved "function"
    name <- identifier
    args <- parens (arg `sepBy` comma)
    symbol "->"
    returns <- parens (identifier `sepBy` comma)
    st <- stmts
    reserved "end"
    return $ FuncDecl name args returns st


--
--  Expressions
--

integerNumber :: Parser Expr
integerNumber = do
    ds <- many1 digit
    return $ Number (read ds :: Double)

floatNumber :: Parser Expr
floatNumber = do
    n <- float
    return $ Number n

number :: Parser Expr
number =  try floatNumber
      <|> integerNumber

matrixElems :: Parser Expr
matrixElems = do
    elems <- literal `sepBy` whiteSpace
    return $ Matrix elems

matrix1d :: Parser Expr
matrix1d = do
    char '['
    elems <- matrixElems
    char ']'
    return elems

matrix2d :: Parser Expr
matrix2d = do
    char '['
    rows <- matrixElems `sepBy` semi
    char ']'
    return $ Matrix rows

matrix :: Parser Expr
matrix =  try matrix1d
      <|> matrix2d

literal :: Parser Expr
literal =  try number
       <|> matrix

funcCall :: Parser Expr
funcCall = do
    name <- identifier
    exprs <- parens (expr `sepBy` comma)
    return $ FuncCall name exprs

varRef :: Parser Expr
varRef = do       
    v <- var
    return $ VarRef v


term = parens expr
    <|> try funcCall
    <|> literal
    <|> varRef
    <?> "term"

ops = [ [unaryOp  "-" Minus, unaryOp "+" Plus ]
      , [unaryOp "not" Not]
      , [binaryOp "*"   Mul AssocLeft, binaryOp "/"  Div AssocLeft ]
      , [binaryOp "+"   Add AssocLeft, binaryOp "-"  Sub AssocLeft ]
      , [binaryOp "=="  Eq  AssocLeft, binaryOp "!=" Neq AssocLeft ]
      , [binaryOp ">"   Gt  AssocLeft, binaryOp "<"  Lt  AssocLeft ]
      , [binaryOp ">="  Gte AssocLeft, binaryOp "<=" Lte AssocLeft ]
      , [binaryOp "and" And AssocLeft, binaryOp "or" Or  AssocLeft ]
      ]
    where
        unaryOp name f = Prefix (do{ reservedOp name; return f })
        binaryOp name f assoc = Infix (do{ reservedOp name; return f }) assoc

expr :: Parser Expr
expr = buildExpressionParser ops term


--
--  A full featured parser
--
program :: Parser [Stmt]
program = do
    whiteSpace
    decls <- many stmt
    eof
    return decls

parse_str :: [Char] -> Either ParseError [Stmt]
parse_str s = parse program "" s'
              where
                  stripComment x = subRegex (mkRegex "-- .*$") x ""
                  s' = join "\n" $ map stripComment $ splitOn "\n" s