#[derive(Debug,PartialEq)]
pub enum Stmt<'a> {
    Assignment(&'a str, Expr<'a>),
    Funcall(Expr<'a>),
}

#[derive(Debug,PartialEq)]
pub enum Expr<'a> {
    Id(&'a str),
    Num(f64),
    Funcall(&'a str, Vec<Expr<'a>>),
    Funcdef(Option<&'a str>, Box<Expr<'a>>),
    Plus(Box<Expr<'a>>, Box<Expr<'a>>),
    Minus(Box<Expr<'a>>, Box<Expr<'a>>),
    Mult(Box<Expr<'a>>, Box<Expr<'a>>),
    Div(Box<Expr<'a>>, Box<Expr<'a>>),
    Block(Vec<Stmt<'a>>),
}

peg::parser! {
    grammar deelang_parser() for str {
        pub rule program() -> Vec<Stmt<'input>>
            = __*  s:stmt()* { s }
        pub rule stmt() -> Stmt<'input>
            = i:id() "<-" _ e:expr() stop() { Stmt::Assignment(i, e) } /
              f:funcall() stop() { Stmt::Funcall(f) }
        rule expr() -> Expr<'input> = precedence! {
            e1:(@) "+" _ e2:@ { Expr::Plus(Box::new(e1), Box::new(e2)) }
            e1:(@) "-" _ e2:@ { Expr::Minus(Box::new(e1), Box::new(e2)) }
            --
            e1:(@) "*" _ e2:@ { Expr::Mult(Box::new(e1), Box::new(e2)) }
            e1:(@) "/" _ e2:@ { Expr::Div(Box::new(e1), Box::new(e2)) }
            --
            "(" _ e:expr() ")" _ { e }
            f:funcall() { f }
            f:funcdef() { f } 
            i:id() _ { Expr::Id(i) }
            n:num() _ { Expr::Num(n) }
        }
        rule block() -> Expr<'input>
            = stop() indent() __* s:stmt()+ dedent() { Expr::Block(s) }

        rule id() -> &'input str
            = i:$(letter() (letter() / digit() / ['?'|'.'|'-'])*) _ { i }
        rule num() -> f64
            = n:$(digit()+ "."? digit()* / "." digit()+) _ { n.parse().unwrap() }
        rule funcall() -> Expr<'input>
            = i:id() "(" _ e:(expr() ** ("," _)) ")" _ { Expr::Funcall(i, e) }
        rule funcdef() -> Expr<'input>
            = i:id()? "->" _ e:(expr() / block()) { Expr::Funcdef(i, Box::new(e)) }

        rule letter()
            = ['A'..='Z'] / ['a'..='z']
        rule digit()
            = ['0'..='9']
        rule stop()
            = __+ / eof()
        rule indent()
            = ">>>"
        rule dedent()
            = "<<<"
        rule _ // Non-meaningful whitespace
            = ['\t'|' ']*
        rule __ // End Of Statement (comment, newline, eof, TODO semicolon)
            = comment()? newline() / comment() &eof()
        rule comment()
            = "#" (!newline() [_])* &(newline() / eof())
        rule newline()
            = "\r\n" / "\r" / "\n"
        rule eof()
            = ![_]
    }
}

fn preprocess(input: &str) -> String {
    let mut stack = vec![0];
    let mut output = String::new();
    for line in input.lines() {
        let mut count = 0;
        for c in line.chars() {
            if c == ' ' || c == '\t' {
                count += 1;
            } else if c == '#' {
                break;
            } else {
                let curr = stack.last().unwrap();
                if curr < &count {
                    stack.push(count);
                    output.push_str(">>>");                  
                } else if curr > &count {
                    while stack.last().unwrap() > &count {
                        output.push_str("<<<");
                        stack.pop();
                    }
                }
                break;
            }
        }
        output.push_str(line.trim());
        output.push('\n');
    }
    output
}


#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_comments() {
        let prgm = "## This is a comment
apple <- 1 ## This is too
## This comment ends the file";
        let expected = vec![Stmt::Assignment("apple", Expr::Num(1.0))];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_funcall() {
        let expected = vec![
            Stmt::Funcall(Expr::Funcall("pear", vec![])),
            Stmt::Funcall(Expr::Funcall("pear",
                                        vec![Expr::Id("x"), Expr::Id("y")],
            ))
        ];
        let prgm = r"pear()
pear(x, y)";
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_assignment() {
        let prgm = r"apple <- 1
apple <- pear(x, y)";
        let expected = vec![
            Stmt::Assignment("apple",
                             Expr::Num(1.0)),
            Stmt::Assignment("apple",
                             Expr::Funcall("pear",
                                           vec![
                                               Expr::Id("x"),
                                               Expr::Id("y"),
                                           ])),
        ];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_operators() {
        let prgm = r"three <- 1 + 2
one <- 3 - 2
four <- (3 - 1) * 2";
        let expected = vec![
            Stmt::Assignment("three",
                             Expr::Plus(
                                 Box::new(Expr::Num(1.0)),
                                 Box::new(Expr::Num(2.0))
                             )),
            Stmt::Assignment("one",
                             Expr::Minus(
                                 Box::new(Expr::Num(3.0)),
                                 Box::new(Expr::Num(2.0))
                             )),
            Stmt::Assignment("four",
                             Expr::Mult(
                                 Box::new(Expr::Minus(
                                      Box::new(Expr::Num(3.0)),
                                     Box::new(Expr::Num(1.0)),
                                 )),
                                 Box::new(Expr::Num(2.0)),
                             ))
        ];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_compound_expression() {
        let prgm = "apple <- pear(x, y) + z";
        let expected = vec![
            Stmt::Assignment("apple",
                             Expr::Plus(
                                 Box::new(Expr::Funcall(
                                     "pear",
                                     vec![Expr::Id("x"), Expr::Id("y")],
                                 )),
                                 Box::new(Expr::Id("z"))
                             )
            ),
        ];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_funcdef() {
        let prgm = r"foo <- -> bar()
foo <- ->
>>>
bar()
baz()
<<<
foo <- x -> y -> x * y";
        let expected = vec![
            Stmt::Assignment(
                "foo",
                Expr::Funcdef(
                    None,
                    Box::new(Expr::Funcall("bar", vec![])),
                )
            ),
            Stmt::Assignment(
                "foo",
                Expr::Funcdef(
                    None,
                    Box::new(Expr::Block(vec![
                        Stmt::Funcall(Expr::Funcall("bar", vec![])),
                        Stmt::Funcall(Expr::Funcall("baz", vec![])),
                    ]))
                )
            ),
            Stmt::Assignment(
                "foo",
                Expr::Funcdef(
                    Some("x"),
                    Box::new(
                        Expr::Funcdef(
                            Some("y"),
                            Box::new(
                                Expr::Mult(
                                    Box::new(Expr::Id("x")),
                                    Box::new(Expr::Id("y")),
                                )
                            )
                        )
                    )
                )
            )
        ];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }

    #[test]
    fn test_preprocess() {
        let prgm = r"
  .
     .
  .
     .
.
      ## Hello World";
        let expected = r"
>>>.
>>>.
<<<.
>>>.
<<<<<<.
## Hello World
";
        assert_eq!(preprocess(prgm), expected);
    }
}