{-# LANGUAGE TypeOperators #-}

import Web.Zwaluw
import Prelude hiding (id, (.))
import Control.Category
import Data.Char


data Expr
  = Variable String
  | Literal Int
  | BinOp Expr Op Expr
  | IfZero Expr Expr Expr
  deriving (Eq, Show)

data Op
  = AddOp
  | MulOp
  deriving (Eq, Show)

variable :: Router r (Expr :- r)
variable = (constr1 Variable $ \a -> do Variable v <- a; return v) . somer (consP . satisfy (\c -> c >= 'a' && c <= 'z')) . nilP

literal :: Router r (Expr :- r)
literal  = (constr1 Literal  $ \a -> do Literal i <- a; return i) . int

binOp :: Router (Expr :- Op :- Expr :- r) (Expr :- r)
binOp    = constr3 BinOp    $ \a -> do BinOp e1 o e2 <- a; return (e1, o, e2)

ifZero :: Router (Expr :- Expr :- Expr :- r) (Expr :- r)
ifZero   = constr3 IfZero   $ \a -> do IfZero i t e <- a; return (i, t, e)

addOp, mulOp :: Router r (Op :- r)
addOp = constr0 AddOp $ \a -> do AddOp <- a; return ()
mulOp = constr0 MulOp $ \a -> do MulOp <- a; return ()

skipSpace, optSpace, sepSpace :: Router r r
skipSpace = lit " " <> lit ""
optSpace = lit "" <> lit " "
sepSpace = somer (lit " ")

keywords = ["ifzero", "else"]

letter = satisfy isLetter
digit = satisfy isDigit

identifier = (consP . letter . listP (letter <> digit)) `having` (`notElem` keywords)

parens p = lit "(" . p . lit ")"

ifz = ifZero . lit "ifzero" . optSpace . parens expr . optSpace . parens expr . optSpace . lit "else" . optSpace . parens expr
atoms = parens (skipSpace . expr . skipSpace) <> literal <> variable <> ifz
op s r = binOp . duck r . optSpace . lit s . optSpace

expr :: Router r (Expr :- r)
expr = atoms `chainl1` op "*" mulOp `chainl1` op "+" addOp