import Data.Char
import Data.List
import Control.Applicative

newtype Parser a = Parser { parse :: String -> Maybe (a,String) }

instance Functor Parser where
  fmap f p = Parser $ \s ->
    case parse p s of
      Nothing -> Nothing
      Just (v,s') -> Just (f v,s')

instance Applicative Parser where
  pure a = Parser $ \s -> Just (a,s)
  f <*> a = Parser $ \s ->
    case parse f s of
      Nothing -> Nothing
      Just (g,s') -> parse (fmap g a) s'

eof :: Parser ()
eof = Parser $ \s -> if null s then Just ((),"") else Nothing

satisfy :: (Char -> Bool) -> Parser Char
satisfy p = Parser f
  where f [] = Nothing
        f (x:xs) = if p x then Just (x,xs) else Nothing

ensure :: (a -> Bool) -> Parser a -> Parser a
ensure pred p = Parser $ \s ->
  case parse p s of
    Nothing -> Nothing
    Just (a,s') -> if pred a then Just (a,s') else Nothing

lookahead :: Parser (Maybe Char)
lookahead = Parser f
  where f [] = Just (Nothing,[])
        f (c:s) = Just (Just c,c:s)

instance Alternative Parser where
  empty = Parser $ \s -> Nothing
  p1 <|> p2 = Parser $ \s ->
    case parse p1 s of
      Just (a,s') -> Just (a,s')
      Nothing -> parse p2 s

spaces :: Parser ()
spaces = many (satisfy isSpace) *> pure ()

char :: Char -> Parser Char
char c = spaces *> satisfy (==c)

int :: Parser Int
int = read <$> (spaces *> some (satisfy isDigit))

parens p = char '(' *> (p <* char ')')

keywords :: [String]
keywords = ["WHILE", "DO", "END", "SKIP"]

kw :: String -> Parser ()
kw s = pure () <* spaces <* (sequenceA $ map (satisfy . (==)) s)

type VarName = String

data ArithExp =
    Var VarName
  | Num Int
  | Plus ArithExp ArithExp
  | Times ArithExp ArithExp
  | Neg ArithExp
  deriving Show

term,factor,atom :: Parser ArithExp
term = (Plus <$> (factor <* char '+') <*> term) <|>
       ((\a b -> Plus a (Neg b)) <$> (factor <* char '-') <*> term) <|>
       factor
factor = (Times <$> neg <* char '*' <*> factor) <|> neg
neg = Neg <$> (char '-' *> atom) <|> atom
atom = Num <$> int <|>
       Var <$> var <|>
       parens term

var :: Parser VarName
var = ensure (not . (`elem` keywords)) (spaces *> id)
  where id = (:) <$> satisfy isAlpha <*> many (satisfy isAlphaNum)

data WhileNZ =
    Skip
  | Assign VarName ArithExp
  | Seq WhileNZ WhileNZ
  | WhileNZ ArithExp WhileNZ
  deriving Show

sepBy1 :: Parser a -> Parser b -> Parser [a]
sepBy1 p sep = (:) <$> p <*> many (sep *> p)

statement, statements :: Parser WhileNZ
statements = foldl1 Seq <$> sepBy1 statement (char ';')
statement =     const Skip <$> kw "SKIP"
            <|> Assign <$> (var <* kw ":=") <*> term
            <|> WhileNZ <$>
                  (kw "WHILE" *> (term <* kw "DO")) <*>
                  (statements <* kw "END")
            <|> (\c b -> Seq c (WhileNZ b c)) <$>
                  (kw "DO" *> (statements <* kw "WHILE")) <*>
                  (term <* kw "END")

-- use statements as the top-level parser

parseProgram :: String -> Maybe WhileNZ
parseProgram s = fst <$> parse statements s

prog1 = parseProgram "x := 0"
prog2 = parseProgram "x := 0;\nWHILE x - 5 DO x := x + 1 END; SKIP;"