use std::str::FromStr;
use ast::{self, WithLocation};

grammar;

pub Program: ast::Program = {
    Expression => ast::Program(<>),
};

RecordFields: Vec<WithLocation<ast::RecordField>> = {
    <l: @L> <n:Name> ":" <v:Name> <r: @R> "," <rest: RecordFields> => {
        let mut rest = rest;
        rest.push(WithLocation::new(ast::RecordField::new(n, v), l, r));
        rest
    },
    <l: @L> <n:Name> ":" <v:Name> <r: @R> "," => vec![WithLocation::new(ast::RecordField::new(n, v), l, r)],
    <l: @L> <n:Name> ":" <v:Name> <r: @R> => vec![WithLocation::new(ast::RecordField::new(n, v), l, r)],
    "" => vec![],
};

Ty: WithLocation<ast::Ty> = {
    <Name> => <>.map(|v| ast::Ty::Name(v)),
    "array" "of" <Name> => <>.map(|v| ast::Ty::Array(Box::new(ast::Ty::Name(v)))),
    <l: @L> "{" <v: RecordFields> "}" <r: @R> => WithLocation::new(ast::Ty::Record(v), l, r),
};

Declaration: Box<WithLocation<ast::Declaration>> = {
    <l: @L> "var" <name: Name> ":" <ty: Ty> "=" <exp: Expression> <r: @R> =>
        Box::new(WithLocation::new(ast::Declaration::new_var(name, Some(ty), exp), l, r)),
    <l: @L> "var" <name: Name> "=" <exp: Expression> <r: @R> =>
        Box::new(WithLocation::new(ast::Declaration::new_var(name, None, exp), l, r)),
    <l: @L> "type" <name: Name> "=" <ty: Ty> <r: @R> =>
        Box::new(WithLocation::new(ast::Declaration::new_ty(name, ty), l, r)),
    <l: @L> "function" <name: Name> "(" <params: RecordFields> ")" "=" <body: Expression> <r: @R> =>
        Box::new(WithLocation::new(ast::Declaration::new_fun(name, params, None, body), l, r)),
    <l: @L> "function" <name: Name> "(" <params: RecordFields> ")" ":" <ty: Ty> "=" <body: Expression> <r: @R> =>
        Box::new(WithLocation::new(ast::Declaration::new_fun(name, params, Some(ty), body), l, r)),
};

DeclarationsList: Vec<WithLocation<ast::Declaration>> = {
    <h: DeclarationsList> <t: Declaration> => {
        let mut h = h;
        h.push(*t);
        h
    },
    <Declaration> => vec![*<>],
};

Expression: Box<ast::WithLocation<ast::Expression>> = {
    <e1:Expression> "&&" <e2:ExpressionEq> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::And, Box::new(v1), Box::new(v2))
    })),
    <e1:Expression> "||" <e2:ExpressionEq> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Or, Box::new(v1), Box::new(v2))
    })),
    <l: @L> "!" <e: ExpressionEq> <r: @R> => Box::new(WithLocation::new(
        ast::Expression::UnaryOp(ast::UnaryOp::Not, e),
        l, r
    )),
    <ExpressionEq> => <>,
};

ExpressionEq: Box<ast::WithLocation<ast::Expression>> = {
    <e1:ExpressionEq> "=" <e2:ExpressionCmp> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Eq, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionEq> "!=" <e2:ExpressionCmp> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Neq, Box::new(v1), Box::new(v2))
    })),
    <ExpressionCmp> => <>,
};

ExpressionCmp: Box<ast::WithLocation<ast::Expression>> = {
    <e1:ExpressionCmp> ">" <e2:ExpressionAdd> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Gt, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionCmp> "<" <e2:ExpressionAdd> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Lt, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionCmp> ">=" <e2:ExpressionAdd> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Ge, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionCmp> "<=" <e2:ExpressionAdd> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Le, Box::new(v1), Box::new(v2))
    })),
    <ExpressionAdd> => <>,
};

ExpressionAdd: Box<ast::WithLocation<ast::Expression>> = {
    <e1:ExpressionAdd> "+" <e2:ExpressionMul> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Add, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionAdd> "-" <e2:ExpressionMul> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Sub, Box::new(v1), Box::new(v2))
    })),
    <ExpressionMul> => <>,
};

ExpressionMul: Box<ast::WithLocation<ast::Expression>> = {
    <e1:ExpressionMul> "*" <e2:ExpressionSign> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Mul, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionMul> "/" <e2:ExpressionSign> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::Div, Box::new(v1), Box::new(v2))
    })),
    <e1:ExpressionMul> "//" <e2:ExpressionSign> => Box::new(e1.join_map(*e2, |v1, v2| {
        ast::Expression::BinOp(ast::BinOp::FloorDiv, Box::new(v1), Box::new(v2))
    })),
    <ExpressionSign> => <>,
};

ExpressionSign: Box<ast::WithLocation<ast::Expression>> = {
    <l: @L> "+" <e: ExpressionBase> <r: @R> => Box::new(WithLocation::new(
        ast::Expression::UnaryOp(ast::UnaryOp::Pos, e),
        l, r
    )),
    <l: @L> "-" <e: ExpressionBase> <r: @R> => Box::new(WithLocation::new(
        ast::Expression::UnaryOp(ast::UnaryOp::Neg, e),
        l, r
    )),
    <ExpressionBase> => <>,
};

ExpressionBase: Box<ast::WithLocation<ast::Expression>> = {
    Num => Box::new(<>.map(|v| ast::Expression::Number(v))),
    String => Box::new(<>.map(|v| ast::Expression::String(v))),
    Name => Box::new(<>.map(|v| ast::Expression::Name(v))),
    Spanned<r"nil"> => Box::new(<>.map(|v| ast::Expression::Nil)),
    <l: @L> "let" <d: DeclarationsList> "in" <e: Expression> "end" <r: @R> =>
        Box::new(WithLocation::new(ast::Expression::Let(d, e), l, r)),
    "(" <Expression> ")" => <>,
};

Name: WithLocation<String> = <e: Spanned<r"[[:alpha:]_][[:alpha:]_0-9]*">> => e.map(|v| v.to_owned());
Num: WithLocation<u64> = <e: Spanned<r"[0-9]+">> => e.map(|v| u64::from_str(v).unwrap());
String: WithLocation<String> = <l: @L> <v:r##""(?:[^"\\]|\\\\)*""##> <r: @R> => {
    WithLocation::new(v[1..v.len() - 1].to_owned(), l, r)
};

Spanned<T>: WithLocation<T> = {
    <l: @L> <v: T> <r: @R> => WithLocation::new(v, l, r)
};

match {
    "nil",
    "var",
    "let",
    "in",
    "end",
    "type",
    "function",
}
else {
    r"[[:alpha:]_][[:alpha:]_0-9]*",
}
else {
    _,
}