module Main where
import Text.Parsec
import Text.Parsec.String
import Text.Parsec.Expr
import Control.Applicative((<$>), (<*>))
import System.Console.Haskeline
import System.Environment (getArgs)
import Data.Fixed (mod')

data Exp = Add Exp Exp           {- e + e -}
         | Ifnz Exp Exp Exp      {- ifnz e1 then e2 else e3 -}
         | K Int                 {- Number representation -}
  deriving Eq

-- evaluate a function to a value
eval :: Exp -> Exp

eval (Add e1 e2) = K (n1 + n2) where
  K n1 = eval e1
  K n2 = eval e2

eval (Ifnz e0 e1 e2) = if n0 /= 0 then v1 else v2 where
  K n0 = eval e0
  v1 = eval e1
  v2 = eval e2

eval (K n) = K n

-- main and the REPL
main :: IO ()
main = do
  args <- getArgs
  if null args then
    runInputT defaultSettings readEvalPrintLoop
  else do
    input <- readFile $ head args
    case parse expr "input" input of
      Right e -> do
        putStrLn $ show $ eval e
      Left err ->
        putStrLn $ show err

readEvalPrintLoop :: InputT IO ()
readEvalPrintLoop = do
  maybeLine <- getInputLine "> "
  case maybeLine of
    Nothing -> return ()
    Just line ->
      case parse command "input" line of
        Right (Eval e)  -> do
          outputStrLn $ show $ eval e
          readEvalPrintLoop
        Right Quit -> return ()
        Left err -> do
          outputStrLn $ show err
          readEvalPrintLoop

-- REPL commands
data Cmd = Eval Exp | Quit

-- A simple parser
command :: Parser Cmd
command = (do
    spaces
    c <- (do { string ":quit"; return Quit}) <|> eval
    spaces
    eof
    return c) <?> "expression, :quit"
  where
    eval  = do { x <- expr; return (Eval x) }

expr :: Parser Exp
expr = (do
  spaces
  string "ifnz"
  e1 <- binary
  spaces
  string "then"
  spaces
  e2 <- binary
  spaces
  string "else"
  spaces
  e3 <- binary
  spaces
  return $ Ifnz e1 e2 e3) <|> binary <?> "expression"

binary :: Parser Exp
binary = buildExpressionParser table factor <?> "expression" where
  table = [
      [Infix (do { char '+'; return $ Add }) AssocLeft]
    ]

factor :: Parser Exp
factor =  (do
              spaces
              f <- do { char '(';  x <- expr; char ')'; return x } <|> number
              spaces
              return f
              )
          <?> "simple expression"

sign :: Parser Char
sign = char '+' <|> char '-'

makeNumber '+' ds = ds
makeNumber '-' ds = '-':ds

number :: Parser Exp
number = (K . read) <$> (do { s <- option '+' sign;
                              ds <- many1 digit;
                              return $ makeNumber s ds})

-- A simple pretty-printer for Exps
instance Show Exp where
  show (K n) = show n
  show (Add e1 e2) = (paren e1) ++ " + " ++ (paren e2)
  show (Ifnz e0 e1 e2) = "ifnz " ++ (paren e0) ++ " then " ++ (paren e1) ++ " else " ++ (paren e2)

paren :: Exp -> String
paren e @ (Add _ _) = "(" ++ show e ++ ")"
paren e @ _           = show e