{-# LANGUAGE DeriveDataTypeable, FlexibleContexts #-}
module QuasiArith
    ( eval
    , expr
    ) where

import Data.Functor.Identity (Identity)
import Data.Generics
import Language.Haskell.TH.Quote
import Text.Parsec.Prim (ParsecT, Stream)
import Text.ParserCombinators.Parsec

import qualified Language.Haskell.TH as TH


data Expr
    =  IntExpr Integer
    |  AntiIntExpr String
    |  BinopExpr BinOp Expr Expr
    |  AntiExpr String
    deriving (Show, Typeable, Data)

data BinOp
    =  AddOp
    |  SubOp
    |  MulOp
    |  DivOp
    deriving (Show, Typeable, Data)


eval :: Expr -> Integer
eval (IntExpr n)        = n
eval (BinopExpr op x y) = (opToFun op) (eval x) (eval y)
  where
    opToFun AddOp = (+)
    opToFun SubOp = (-)
    opToFun MulOp = (*)
    opToFun DivOp = div
eval (AntiIntExpr s)    = read s
eval (AntiExpr s)       = read s


small, large, idchar :: ParsecT [Char] u Identity Char
small = lower <|> char '_'
large = upper
idchar = small <|> large <|> digit <|> char '\''

lexeme :: Stream s m Char => ParsecT s u m b -> ParsecT s u m b
lexeme p = do { x <- p; spaces; return x  }

symbol :: Stream s m Char => String -> ParsecT s u m String
symbol name  = lexeme (string name)

parens :: Stream s m Char => ParsecT s u m a -> ParsecT s u m a
parens p = between (symbol "(") (symbol ")") p

pExpr :: CharParser st Expr
pExpr = term    `chainl1` addop

term :: CharParser st Expr
term = factor `chainl1` mulop

factor :: CharParser st Expr
factor = parens pExpr <|> integer <|> try antiIntExpr <|> antiExpr

mulop :: ParsecT [Char] u Identity (Expr -> Expr -> Expr)
mulop =      do { _ <- symbol "*"; return $ BinopExpr MulOp }
        <|>  do { _ <- symbol "/"; return $ BinopExpr DivOp }

addop :: ParsecT [Char] u Identity (Expr -> Expr -> Expr)
addop =      do { _ <- symbol "+"; return $ BinopExpr AddOp }
        <|>  do { _ <- symbol "-"; return $ BinopExpr SubOp }

integer  ::  CharParser st Expr
integer = lexeme $ do { ds <- many1 digit ; return $ IntExpr (read ds) }

ident :: CharParser s String
ident =  do { c <- small; cs <- many idchar; return (c:cs) }

antiIntExpr :: ParsecT [Char] u Identity Expr
antiIntExpr = lexeme $ do { _ <- symbol "$int:"; i <- ident; return $ AntiIntExpr i }

antiExpr :: ParsecT [Char] u Identity Expr
antiExpr = lexeme $ do { _ <- symbol "$"; i <- ident; return $ AntiExpr i }

parseExpr :: Monad m => (String, Int, Int) -> String -> m Expr
parseExpr (file, line, col) s =
    case runParser p () "" s of
      Left err  -> fail $ show err
      Right e   -> return e
  where
    p = do
        pos <- getPosition
        setPosition $
              (flip setSourceName) file $
              (flip setSourceLine) line $
              (flip setSourceColumn) col $
              pos
        spaces
        e <- pExpr
        eof
        return e

quoteExprExp :: String -> TH.ExpQ
quoteExprExp s =  do
    loc <- TH.location
    let pos =  (TH.loc_filename loc,
                    fst (TH.loc_start loc),
                    snd (TH.loc_start loc))
    e <- parseExpr pos s
    dataToExpQ (const Nothing `extQ` antiExprExp) e

antiExprExp :: Expr -> Maybe (TH.Q TH.Exp)
antiExprExp (AntiIntExpr v)  = Just $ TH.appE (TH.conE (TH.mkName "IntExpr"))
                                                (TH.varE (TH.mkName v))
antiExprExp (AntiExpr v)     = Just $ TH.varE (TH.mkName v)
antiExprExp _                = Nothing


quoteExprPat :: String -> TH.PatQ
quoteExprPat s = do
    loc <- TH.location
    let pos =  (TH.loc_filename loc,
                    fst (TH.loc_start loc),
                    snd (TH.loc_start loc))
    e <- parseExpr pos s
    dataToPatQ (const Nothing `extQ` antiExprPat) e

antiExprPat :: Expr -> Maybe (TH.Q TH.Pat)
antiExprPat (AntiIntExpr v) = Just $ TH.conP (TH.mkName "IntExpr")
                                                [TH.varP (TH.mkName v)]
antiExprPat (AntiExpr v) = Just $ TH.varP (TH.mkName v)
antiExprPat _            = Nothing

quoteExprType :: String -> TH.Q TH.Type
quoteExprType = undefined

quoteExprDec :: String -> TH.Q [TH.Dec]
quoteExprDec _ = return []

expr  :: QuasiQuoter
expr  =  QuasiQuoter quoteExprExp quoteExprPat quoteExprType quoteExprDec