import ParseLib

-- Code needed for exercise 3 and 4:

data Expr = Add Expr Expr
          | Mul Expr Expr
          | Lam String Expr
          | If Expr Expr Expr
          | App Expr Expr
          | Var String
          | Nat Int
          | Bool Bool
    deriving Show

lang :: Parser Expr
lang = expr <* eof

expr :: Parser Expr
expr = Lam <$ symbol "\\" <*> identifier <* symbol "->" <*> expr
        <|> If <$ symbol "if" <*> expr <* symbol "then" <*> expr <* symbol "else" <*> expr
        <|> App <$> expra <*> expra
        <|> expra

expra :: Parser Expr
expra = e <$> term <*> try (symbol "+" *> expra)
  where
   e term Nothing = term
   e term (Just expr) = Add term expr

term :: Parser Expr
term = t <$> factor <*> try (symbol "*" *> term)
  where
   t fac Nothing = fac
   t fac (Just term) = Mul fac term

factor :: Parser Expr
factor = val <|> Var <$> identifier <|> paren expr

val :: Parser Expr
val = Bool <$> bool <|> Nat <$> natural

bool :: Parser Bool
bool = True <$ symbol "True" <|> False <$ symbol "False"

-- Technically we only want real values back not lambdas (ints, bools, tuples of vals)
-- but for simplicity we just go with printing an ast evaluated as far as possible
interpreter :: String -> String
interpreter = show . ((eval []) . fst . head ) . (parse lang)

eval :: [(String, Expr)] -> Expr -> Expr
eval env (Add e1 e2) = case (eval env e1, eval env e2) of
                    (Nat i1, Nat i2) -> Nat $ i1 + i2
                    _                -> error "Arguments for addition should be integers"
eval env (Mul e1 e2) = case (eval env e1, eval env e2) of
                    (Nat i1, Nat i2) -> Nat $ i1 + i2
                    _                -> error "Arguments for multiplication should be integers"
eval env (Lam x e) = (Lam x e)
eval env (If e1 e2 e3) = case eval env e1 of
                            (Bool True)  -> eval env e2
                            (Bool False) -> eval env e3
                            _            -> error "Conditional of if expression must be a boolean"
eval env (Var x) = case lookup x env of
                   Nothing -> error "Cannot reference unbound variables"
                   (Just e) -> eval env e 
eval env (App e1 e2) = case eval env e1 of
                        (Lam x e3) -> eval ((x,e2):env) e3
                        _          -> error "Can only apply to functions"
eval _   (Nat i) = Nat i
eval _   (Bool b) = Bool b