use std::fmt;

#[derive(Debug,PartialEq,Clone)]
pub enum Stmt {
    Assignment(String, Expr),
    Funcall(Expr),
    Conditional(Vec<Expr>, Option<Expr>),
    ReplPrint(Expr),
}

#[derive(Debug,PartialEq,Clone)]
pub enum Expr {
    Id(String),
    Atom(Atom),
    Funcall(String, Vec<Expr>),
    Funcdef(Option<String>, Box<Expr>),
    UnaryMinus(Box<Expr>),
    Plus(Box<Expr>, Box<Expr>),
    Minus(Box<Expr>, Box<Expr>),
    Mult(Box<Expr>, Box<Expr>),
    Div(Box<Expr>, Box<Expr>),
    Block(Vec<Stmt>),
    GuardedBlock(Box<GuardedBlock>),
    Object(Vec<Stmt>),
}

#[derive(Debug,PartialEq,Clone)]
pub enum Atom {
    String(String),
    Num(f64),
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Atom::String(a) => write!(f, "\"{}\"", a),
            Atom::Num(a) => write!(f, "{}", a),
        }
    }
}

#[derive(Debug,PartialEq,Clone)]
pub struct GuardedBlock {
    pub guard: Expr,
    pub block: Vec<Stmt>,
}

peg::parser! {
    grammar deelang_parser() for str {
        pub rule program() -> Vec<Stmt>
            = __*  s:stmt()* { s }
        pub rule stmt() -> Stmt
            = a:assignment() { a } /
            f:funcall() stop() { Stmt::Funcall(f) } /
            c:conditional() { c } /
            e:expr() stop() { Stmt::ReplPrint(e) }
        rule expr() -> Expr = precedence! {
            "-" _ e1:@ { Expr::UnaryMinus(Box::new(e1)) }
            --
            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 }
            ['"'] s:$((!['"'] [_] / r#"\""#)*) ['"'] { Expr::Atom(Atom::String(s.to_string())) }
            f:funcall() { f }
            f:funcdef() { f }
            o:object() { o }
            i:id() _ { Expr::Id(i) }
            n:num() _ { Expr::Atom(Atom::Num(n)) }

        }

        rule id() -> String
            = i:$(letter() (letter() / digit() / ['?'|'.'|'-'])*) _ { i.to_string() }
        rule assignment() -> Stmt
            = i:id() "<-" _ e:expr() stop() { Stmt::Assignment(i, e) }
        rule num() -> f64
            = n:$(digit()+ "."? digit()* / "." digit()+) _ { n.parse().unwrap() }
        rule funcall() -> Expr
            = i:id() "(" _ e:(expr() ** ("," _)) ")" _ { Expr::Funcall(i, e) }
        rule funcdef() -> Expr
            = i:id()? "->" _ e:(expr() / block()) { Expr::Funcdef(i, Box::new(e)) }
        rule conditional() -> Stmt
            = i:_if() __* ei:elif()* __* e:_else()? __* { Stmt::Conditional([vec![i], ei].concat(), e) }
        rule object() -> Expr
            = "{" _ stop() indent() __* a:assignment()+ dedent() __* "}" _ { Expr::Object(a) }
        
        rule _if() -> Expr
            = "if" _ g:expr() b:indented_block() {
                Expr::GuardedBlock(Box::new(GuardedBlock {
                    guard: g,
                    block: b,
                }))
            }
        rule elif() -> Expr
            = "elif" _ g:expr() b:indented_block() {
                Expr::GuardedBlock(Box::new(GuardedBlock {
                    guard: g,
                    block: b
                }))
            }
        rule _else() -> Expr
            = "else" _ b:block() { b }
        rule block() -> Expr
            = i:indented_block() { Expr::Block(i) }
        rule indented_block() -> Vec<Stmt>
            = stop() indent() __* s:stmt()+ dedent() { s }

        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 {
                use std::cmp::Ordering::*;
                let curr = stack.last().unwrap();
                match curr.cmp(&count) {
                    Less => {
                        stack.push(count);
                        output.push_str(">>>");                  
                    },
                    Greater => {
                        while stack.last().unwrap() > &count {
                            output.push_str("<<<");
                            output.push('\n');
                            stack.pop();
                        }
                    }
                    Equal => (),
                }
                break;
            }
        }
        output.push_str(line.trim());
        output.push('\n');
    }
    output
}

pub fn parse(prgm: &str) -> Vec<Stmt> {
    let prgm = preprocess(prgm);
    deelang_parser::program(&prgm).unwrap()
}

pub fn parse_stmt(stmt: &str) -> Stmt {
    let stmt = preprocess(stmt);
    deelang_parser::stmt(&stmt).unwrap()
}


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

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

    #[test]
    fn test_funcall() {
        let expected = vec![
            Stmt::Funcall(Expr::Funcall("pear".to_string(), vec![])),
            Stmt::Funcall(Expr::Funcall("pear".to_string(),
                                        vec![Expr::Id("x".to_string()), Expr::Id("y".to_string())],
            ))
        ];
        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".to_string(),
                             Expr::Atom(Atom::Num(1.0))),
            Stmt::Assignment("apple".to_string(),
                             Expr::Funcall("pear".to_string(),
                                           vec![
                                               Expr::Id("x".to_string()),
                                               Expr::Id("y".to_string()),
                                           ])),
        ];
        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".to_string(),
                             Expr::Plus(
                                 Box::new(Expr::Atom(Atom::Num(1.0))),
                                 Box::new(Expr::Atom(Atom::Num(2.0))),
                             )),
            Stmt::Assignment("one".to_string(),
                             Expr::Minus(
                                 Box::new(Expr::Atom(Atom::Num(3.0))),
                                 Box::new(Expr::Atom(Atom::Num(2.0))),
                             )),
            Stmt::Assignment("four".to_string(),
                             Expr::Mult(
                                 Box::new(Expr::Minus(
                                     Box::new(Expr::Atom(Atom::Num(3.0))),
                                     Box::new(Expr::Atom(Atom::Num(1.0))),
                                 )),
                                 Box::new(Expr::Atom(Atom::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".to_string(),
                             Expr::Plus(
                                 Box::new(Expr::Funcall(
                                     "pear".to_string(),
                                     vec![Expr::Id("x".to_string()), Expr::Id("y".to_string())],
                                 )),
                                 Box::new(Expr::Id("z".to_string()))
                             )
            ),
        ];
        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".to_string(),
                Expr::Funcdef(
                    None,
                    Box::new(Expr::Funcall("bar".to_string(), vec![])),
                )
            ),
            Stmt::Assignment(
                "foo".to_string(),
                Expr::Funcdef(
                    None,
                    Box::new(Expr::Block(vec![
                        Stmt::Funcall(Expr::Funcall("bar".to_string(), vec![])),
                        Stmt::Funcall(Expr::Funcall("baz".to_string(), vec![])),
                    ]))
                )
            ),
            Stmt::Assignment(
                "foo".to_string(),
                Expr::Funcdef(
                    Some("x".to_string()),
                    Box::new(
                        Expr::Funcdef(
                            Some("y".to_string()),
                            Box::new(
                                Expr::Mult(
                                    Box::new(Expr::Id("x".to_string())),
                                    Box::new(Expr::Id("y".to_string())),
                                )
                            )
                        )
                    )
                )
            )
        ];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }
    #[test]
    fn test_conditional() {
        let prgm = r"if foo
>>>bar()
<<<
elif baz
>>>foobar()
<<<
else
>>>quux()
<<<";
        let expected = vec![Stmt::Conditional(
            vec![
                Expr::GuardedBlock(Box::new(GuardedBlock {
                    guard: Expr::Id("foo".to_string()),
                    block: vec![Stmt::Funcall(Expr::Funcall("bar".to_string(), vec![]))]
                })),
                Expr::GuardedBlock(Box::new(GuardedBlock {
                    guard: Expr::Id("baz".to_string()),
                    block: vec![Stmt::Funcall(Expr::Funcall("foobar".to_string(), vec![]))]
                })),
            ],
            Some(Expr::Block(vec![Stmt::Funcall(Expr::Funcall("quux".to_string(), vec![]))])),
        )];
        assert_eq!(deelang_parser::program(prgm).unwrap(), expected);
    }
    #[test]
    fn test_object() {
        let prgm = r"fruit <- {
>>>apple <- 1
pear <- 2
<<<
}";
        let expected = vec![Stmt::Assignment(
            "fruit".to_string(),
            Expr::Object(vec![
                Stmt::Assignment("apple".to_string(), Expr::Atom(Atom::Num(1.0))),
                Stmt::Assignment("pear".to_string(), Expr::Atom(Atom::Num(2.0))),
            ]),
        )];
        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);
    }
}