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

-- step function mirroring the small-step operational semantics
step :: Exp -> Exp

-- Note: the order of definition matters because we want
-- more specific patterns to match first

-- computation
step (Add (K x) (K y)) = K (x+y)

step (Ifnz (K 0) e1 e2) = e2
step (Ifnz (K _) e1 e2) = e1

-- congruence
step (Add e1 e2) | value e1  = Add e1 e2'
                 | otherwise = Add e1' e2 where
                      e1' = step e1
                      e2' = step e2

step (Ifnz e0 e1 e2) = Ifnz e0' e1 e2 where
                          e0' = step e0

-- values
step (K n) = K n

-- is the expression a value?
value :: Exp -> Bool
value (K n) = True
value _ = False

-- evaluate an expression given a step function, returning the trace of intermediate terms
trace :: Exp -> [Exp]
trace = takeWhilePlus1 (not . value) . iterate step
  where
    -- like takeWhile, but includes the first element where p is false
    takeWhilePlus1 p xs =
      case span p xs of
        (ys, z:zs) -> ys ++ [z]
        (ys, [])   -> ys

-- evaluate an expression given a step function
eval :: Exp -> Exp
eval e
  | value e   = e
  | otherwise = eval (step e)

-- 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 (Trace e) -> do
          outputStr $ unlines $ map show $ trace e
          readEvalPrintLoop
        Right (Eval e)  -> do
          outputStrLn $ show $ eval e
          readEvalPrintLoop
        Right Quit -> return ()
        Left err -> do
          outputStrLn $ show err
          readEvalPrintLoop

-- REPL commands
data Cmd = Trace Exp | Eval Exp | Quit

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

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