// See also: file:///usr/share/doc/python/html/reference/grammar.html?highlight=grammar
// See also: https://github.com/antlr/grammars-v4/blob/master/python3/Python3.g4
// See also: file:///usr/share/doc/python/html/reference/compound_stmts.html#function-definitions
// See also: https://greentreesnakes.readthedocs.io/en/latest/nodes.html#keyword
#![allow(unknown_lints,clippy)]

use super::ast;
use super::lexer;
use std::iter::FromIterator;
use num_bigint::BigInt;

grammar;

// This is a hack to reduce the amount of lalrpop tables generated:
// For each public entry point, a full parse table is generated.
// By having only a single pub function, we reduce this to one.
pub Top: ast::Top = {
    StartProgram <p:Program> => ast::Top::Program(p),
    StartStatement <s:Statement> => ast::Top::Statement(s),
    StartExpression <e:Expression> => ast::Top::Expression(e),
};

Program: ast::Program = {
    <lines:FileLine*> => ast::Program {
        statements: Vec::from_iter(lines.into_iter().filter_map(|e| e))
    },
};

// A file line either has a declaration, or an empty newline:
FileLine: Option<ast::LocatedStatement> = {
    <s:Statement> => Some(s),
    "\n" => None,
};

Suite: Vec<ast::LocatedStatement> = {
    <s:SimpleStatement> => vec![s],
    "\n" indent <s:Statement+> dedent => s,
};

Statement: ast::LocatedStatement = {
    SimpleStatement,
    CompoundStatement,
};

SimpleStatement: ast::LocatedStatement = {
    <s:SmallStatement> "\n" => s,
    <s:SmallStatement> ";" => s,
};

SmallStatement: ast::LocatedStatement = {
    ExpressionStatement,
    PassStatement,
    DelStatement,
    FlowStatement,
    ImportStatement,
    GlobalStatement,
    NonlocalStatement,
    AssertStatement,
};

PassStatement: ast::LocatedStatement = {
    <loc:@L> "pass" => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Pass,
        }
    },
};

DelStatement: ast::LocatedStatement = {
    <loc:@L> "del" <e:ExpressionList2> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Delete { targets: e },
        }
    },
};

ExpressionStatement: ast::LocatedStatement = {
    <loc:@L> <expr:TestOrStarExprList> <suffix:AssignSuffix*> => {
        // Just an expression, no assignment:
        if suffix.is_empty() {
            ast::LocatedStatement {
                location: loc.clone(),
                node: ast::Statement::Expression { expression: expr }
            }
        } else {
          let mut targets = vec![expr];
          let mut values = suffix;

          while values.len() > 1 {
                targets.push(values.remove(0));
          }

          let value = values.into_iter().next().unwrap();

          ast::LocatedStatement {
            location: loc.clone(),
            node: ast::Statement::Assign { targets, value },
          }
        }
    },
    <loc:@L> <expr:TestOrStarExprList> <op:AugAssign> <e2:TestList> => {
        // TODO: this works in most cases:
        let rhs = e2.into_iter().next().unwrap();
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::AugAssign { target: expr, op: op, value: rhs },
        }
    },
};

AssignSuffix: ast::Expression = {
    "=" <e:TestList> => {
        if e.len() > 1 {
            ast::Expression::Tuple {
                elements: e
            }
        } else {
            e.into_iter().next().unwrap()
        }
    },
    "=" <e:YieldExpr> => e,
};

TestOrStarExprList: ast::Expression = {
    <e:TestOrStarExpr> <e2:("," TestOrStarExpr)*> <comma:","?> => {
        let mut res = vec![e];
        res.extend(e2.into_iter().map(|x| x.1));

        // First build tuple from first item:
        let expr = if (res.len() > 1) || comma.is_some() {
            ast::Expression::Tuple { elements: res }
        } else {
            res.into_iter().next().unwrap()
        };

        expr
    }
};

TestOrStarExpr: ast::Expression = {
    Test,
    StarExpr,
};

AugAssign: ast::Operator = {
    "+=" => ast::Operator::Add,
    "-=" => ast::Operator::Sub,
    "*=" => ast::Operator::Mult,
    "@=" => ast::Operator::MatMult,
    "/=" => ast::Operator::Div,
    "%=" => ast::Operator::Mod,
    "&=" => ast::Operator::BitAnd,
    "|=" => ast::Operator::BitOr,
    "^=" => ast::Operator::BitXor,
    "<<=" => ast::Operator::LShift,
    ">>=" => ast::Operator::RShift,
    "**=" => ast::Operator::Pow,
    "//=" => ast::Operator::FloorDiv,
};

FlowStatement: ast::LocatedStatement = {
    <loc:@L> "break" => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Break,
        }
    },
    <loc:@L> "continue" => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Continue,
        }
    },
    <loc:@L> "return" <t:TestList?> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Return { value: t },
        }
    },
    <loc:@L> <y:YieldExpr> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Expression { expression: y },
        }
    },
    RaiseStatement,
};

RaiseStatement: ast::LocatedStatement = {
    <loc:@L> "raise" => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Raise { exception: None, cause: None },
        }
    },
    <loc:@L> "raise" <t:Test> <c:("from" Test)?> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Raise { exception: Some(t), cause: c.map(|x| x.1) },
        }
    },
};

ImportStatement: ast::LocatedStatement = {
  <loc:@L> "import" <i: Comma<ImportPart<<DottedName>>>> => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::Import {
        import_parts: i
            .iter()
            .map(|(n, a)|
                ast::SingleImport {
                    module: n.to_string(),
                    symbol: None,
                    alias: a.clone()
                })
            .collect()
      },
    }
  },
  <loc:@L> "from" <n:ImportFromLocation> "import" <i: ImportAsNames> => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::Import {
        import_parts: i
            .iter()
            .map(|(i, a)|
                ast::SingleImport {
                    module: n.to_string(),
                    symbol: Some(i.to_string()),
                    alias: a.clone()
                })
            .collect()
      },
    }
  },
};

ImportFromLocation: String = {
    <dots: "."*> <name:DottedName> => {
       let mut r = "".to_string();
       for _dot in dots {
         r.push_str(".");
       }
      r.push_str(&name);
      r
    },
    <dots: "."+> => {
       let mut r = "".to_string();
       for _dot in dots {
         r.push_str(".");
       }
      r
    },
};

ImportAsNames: Vec<(String, Option<String>)> = {
    <i:Comma<ImportPart<Identifier>>> => i,
    "(" <i:Comma<ImportPart<Identifier>>> ")" => i,
    "*" => {
        // Star import all
        vec![("*".to_string(), None)]
    },
};


#[inline]
ImportPart<I>: (String, Option<String>) = {
    <i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
};

// A name like abc or abc.def.ghi
DottedName: String = {
    <n:name> => n,
    <n:name> <n2: ("." Identifier)+>  => {
        let mut r = n.to_string();
        for x in n2 {
            r.push_str(".");
            r.push_str(&x.1);
        }
        r
    },
};

GlobalStatement: ast::LocatedStatement = {
    <loc:@L> "global" <names:OneOrMore<Identifier>> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Global { names }
        }
    },
};

NonlocalStatement: ast::LocatedStatement = {
    <loc:@L> "nonlocal" <names:OneOrMore<Identifier>> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Nonlocal { names }
        }
    },
};

AssertStatement: ast::LocatedStatement = {
    <loc:@L> "assert" <t:Test> <m: ("," Test)?> => {
        ast::LocatedStatement {
          location: loc,
          node: ast::Statement::Assert {
            test: t,
            msg: match m {
              Some(e) => Some(e.1),
              None => None,
            }
          }
        }
    },
};

CompoundStatement: ast::LocatedStatement = {
    IfStatement,
    WhileStatement,
    ForStatement,
    TryStatement,
    WithStatement,
    FuncDef,
    ClassDef,
};

IfStatement: ast::LocatedStatement = {
    <loc:@L> "if" <t:Test> ":" <s1:Suite> <s2:(@L "elif" Test ":" Suite)*> <s3:("else" ":" Suite)?> => {
        // Determine last else:
        let mut last = match s3 {
            Some(s) => Some(s.2),
            None => None,
        };

        // handle elif:
        for i in s2.into_iter().rev() {
          let x = ast::LocatedStatement {
            location: i.0,
            node: ast::Statement::If { test: i.2, body: i.4, orelse: last },
          };
          last = Some(vec![x]);
        }

        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::If { test: t, body: s1, orelse: last }
        }
    },
};

WhileStatement: ast::LocatedStatement = {
    <loc:@L> "while" <e:Test> ":" <s:Suite> <s2:("else" ":" Suite)?> => {
        let or_else = match s2 {
            Some(s) => Some(s.2),
            None => None,
        };
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::While { test: e, body: s, orelse: or_else },
        }
    },
};

ForStatement: ast::LocatedStatement = {
    <loc:@L> "for" <e:ExpressionList> "in" <t:TestList> ":" <s:Suite> <s2:("else" ":" Suite)?> => {
        let or_else = match s2 {
            Some(s) => Some(s.2),
            None => None,
        };
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::For {
                target: e,
                iter: t, body: s, orelse: or_else
            },
        }
    },
};

TryStatement: ast::LocatedStatement = {
    <loc:@L> "try" ":" <body:Suite> <handlers:ExceptClause*> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> => {
        let or_else = match else_suite {
            Some(s) => Some(s.2),
            None => None,
        };
        let finalbody = match finally {
            Some(s) => Some(s.2),
            None => None,
        };
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::Try {
                body: body,
                handlers: handlers,
                orelse: or_else,
                finalbody: finalbody,
            },
        }
    },
};

ExceptClause: ast::ExceptHandler = {
    "except" <typ:Test?> ":" <body:Suite> => {
        ast::ExceptHandler {
            typ: typ,
            name: None,
            body: body,
        }
    },
    "except" <x:(Test "as" Identifier)> ":" <body:Suite> => {
        ast::ExceptHandler {
            typ: Some(x.0),
            name: Some(x.2),
            body: body,
        }
    },
};

WithStatement: ast::LocatedStatement = {
    <loc:@L> "with" <items:OneOrMore<WithItem>> ":" <s:Suite> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::With { items: items, body: s },
        }
    },
};

WithItem: ast::WithItem = {
    <t:Test> <n:("as" Expression)?> => {
        let optional_vars = match n {
            Some(val) => Some(val.1),
            None => None,
        };
        ast::WithItem { context_expr: t, optional_vars }
    },
};

FuncDef: ast::LocatedStatement = {
    <d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::FunctionDef {
                name: i,
                args: a,
                body: s,
                decorator_list: d,
            }
        }
    },
};

Parameters: ast::Parameters = {
    "(" <a: (TypedArgsList)?> ")" => {
        match a {
            Some(a) => a,
            None => Default::default(),
        }
    },
};

// parameters are (String, None), kwargs are (String, Some(Test)) where Test is
// the default
TypedArgsList: ast::Parameters = {
    <param1:TypedParameters> <args2:("," ParameterListStarArgs)?> => {
        let (names, default_elements) = param1;

        // Now gather rest of parameters:
        let (vararg, kwonlyargs, kw_defaults, kwarg) = match args2 {
            Some((_, x)) => x,
            None => (None, vec![], vec![], None),
        };

        ast::Parameters {
            args: names,
            kwonlyargs: kwonlyargs,
            vararg: vararg,
            kwarg: kwarg,
            defaults: default_elements,
            kw_defaults: kw_defaults,
        }
    },
    <param1:TypedParameters> <kw:("," KwargParameter)> => {
        let (names, default_elements) = param1;

        // Now gather rest of parameters:
        let vararg = None;
        let kwonlyargs = vec![];
        let kw_defaults = vec![];
        let kwarg = Some(kw.1);

        ast::Parameters {
            args: names,
            kwonlyargs: kwonlyargs,
            vararg: vararg,
            kwarg: kwarg,
            defaults: default_elements,
            kw_defaults: kw_defaults,
        }
    },
    <params:ParameterListStarArgs> => {
        let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
        ast::Parameters {
            args: vec![],
            kwonlyargs: kwonlyargs,
            vararg: vararg,
            kwarg: kwarg,
            defaults: vec![],
            kw_defaults: kw_defaults,
        }
    },
    <kw:KwargParameter> => {
        ast::Parameters {
            args: vec![],
            kwonlyargs: vec![],
            vararg: None,
            kwarg: Some(kw),
            defaults: vec![],
            kw_defaults: vec![],
        }
    },
};

// Use inline here to make sure the "," is not creating an ambiguity.
#[inline]
TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
    <param1:TypedParameterDef> <param2:("," TypedParameterDef)*> => {
        // Combine first parameters:
        let mut args = vec![param1];
        args.extend(param2.into_iter().map(|x| x.1));

        let mut names = vec![];
        let mut default_elements = vec![];

        for (name, default) in args.into_iter() {
            names.push(name.clone());
            if let Some(default) = default {
                default_elements.push(default);
            } else {
                if default_elements.len() > 0 {
                    // Once we have started with defaults, all remaining arguments must
                    // have defaults
                    panic!(
                        "non-default argument follows default argument: {}",
                        name
                    );
                }
            }
        }

        (names, default_elements)
    }
};

TypedParameterDef: (String, Option<ast::Expression>) = {
    <i:TypedParameter> => (i, None),
    <i:TypedParameter> "=" <e:Test> => (i, Some(e)),
};

// TODO: add type annotations here:
TypedParameter: String = {
    Identifier,
};

ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Expression>>, Option<Option<String>>) = {
    "*" <va:Identifier?> <kw:("," TypedParameterDef)*> <kwarg:("," KwargParameter)?> => {
        // Extract keyword arguments:
        let mut kwonlyargs = vec![];
        let mut kw_defaults = vec![];
        for (name, value) in kw.into_iter().map(|x| x.1) {
            kwonlyargs.push(name);
            kw_defaults.push(value);
        }

        let kwarg = match kwarg {
            Some((_, name)) => Some(name),
            None => None,
        };

        (Some(va), kwonlyargs, kw_defaults, kwarg)
    }
};

KwargParameter: Option<String> = {
    "**" <kwarg:Identifier?> => {
        kwarg
    }
};

ClassDef: ast::LocatedStatement = {
    <d:Decorator*> <loc:@L> "class" <n:Identifier> <a:("(" ArgumentList ")")?> ":" <s:Suite> => {
        let (bases, keywords) = match a {
            Some((_, args, _)) => args,
            None => (vec![], vec![]),
        };
        ast::LocatedStatement {
            location: loc,
            node: ast::Statement::ClassDef {
              name: n,
              bases: bases,
              keywords: keywords,
              body: s,
              decorator_list: d,
            },
        }
    },
};

// Decorators:
Decorator: ast::Expression = {
    "@" <n:DottedName> <a: ("(" ArgumentList ")")?> "\n" => {
        let name = ast::Expression::Identifier { name: n };
        match a {
            Some((_, args, _)) => ast::Expression::Call {
                function: Box::new(name),
                args: args.0,
                keywords: args.1,
            },
            None => name,
        }
    },
};

YieldExpr: ast::Expression = {
    "yield" <ex:TestList?> => {
        ast::Expression::Yield {
            value: ex.map(|expr| Box::new(
                if expr.len() > 1 {
                    ast::Expression::Tuple { elements: expr }
                } else {
                    expr.into_iter().next().unwrap()
                })
            )
        }
    },
    "yield" "from" <e:Test> => {
        ast::Expression::YieldFrom {
            value: Box::new(e),
        }
    },
};

Test: ast::Expression = {
    <e:OrTest> <c: ("if" OrTest "else" Test)?> => {
        match c {
            Some(c) => {
                ast::Expression::IfExpression {
                    test: Box::new(c.1),
                    body: Box::new(e),
                    orelse: Box::new(c.3),
                }
            },
            None => e,
        }
    },
    <e:LambdaDef> => e,
};

LambdaDef: ast::Expression = {
    "lambda" <p:TypedArgsList?> ":" <b:Expression> =>
        ast::Expression::Lambda {
            args: p.unwrap_or(Default::default()),
            body:Box::new(b)
       }
}

OrTest: ast::Expression = {
    <e:AndTest> => e,
    <e1:OrTest> "or" <e2:AndTest> => ast::Expression::BoolOp { a: Box::new(e1), op: ast::BooleanOperator::Or, b: Box::new(e2) },
};

AndTest: ast::Expression = {
    <e:NotTest> => e,
    <e1:AndTest> "and" <e2:NotTest> => ast::Expression::BoolOp { a: Box::new(e1), op: ast::BooleanOperator::And, b: Box::new(e2) },
};

NotTest: ast::Expression = {
    "not" <e:NotTest> => ast::Expression::Unop { a: Box::new(e), op: ast::UnaryOperator::Not },
    <e:Comparison> => e,
};

Comparison: ast::Expression = {
    <e1:Comparison> <op:CompOp> <e2:Expression> => ast::Expression::Compare { a: Box::new(e1), op: op, b: Box::new(e2) },
    <e:Expression> => e,
};

CompOp: ast::Comparison = {
    "==" => ast::Comparison::Equal,
    "!=" => ast::Comparison::NotEqual,
    "<" => ast::Comparison::Less,
    "<=" => ast::Comparison::LessOrEqual,
    ">" => ast::Comparison::Greater,
    ">=" => ast::Comparison::GreaterOrEqual,
    "in" => ast::Comparison::In,
    "not" "in" => ast::Comparison::NotIn,
    "is" => ast::Comparison::Is,
    "is" "not" => ast::Comparison::IsNot,
};

Expression: ast::Expression = {
    <e1:Expression> "|" <e2:XorExpression> => ast::Expression::Binop { a: Box::new(e1), op: ast::Operator::BitOr, b: Box::new(e2) },
    <e:XorExpression> => e,
};

XorExpression: ast::Expression = {
    <e1:XorExpression> "^" <e2:AndExpression> => ast::Expression::Binop { a: Box::new(e1), op: ast::Operator::BitXor, b: Box::new(e2) },
    AndExpression,
};

AndExpression: ast::Expression = {
    <e1:AndExpression> "&" <e2:ShiftExpression> => ast::Expression::Binop { a: Box::new(e1), op: ast::Operator::BitAnd, b: Box::new(e2) },
    ShiftExpression,
};

ShiftExpression: ast::Expression = {
    <e1:ShiftExpression> <op:ShiftOp> <e2:ArithmaticExpression> => ast::Expression::Binop { a: Box::new(e1), op: op, b: Box::new(e2) },
    ArithmaticExpression,
};

ShiftOp: ast::Operator = {
    "<<" => ast::Operator::LShift,
    ">>" => ast::Operator::RShift,
};

ArithmaticExpression: ast::Expression = {
    <a:ArithmaticExpression> <op:AddOp> <b:Term> => ast::Expression::Binop { a: Box::new(a), op: op, b: Box::new(b) },
    Term,
};

AddOp: ast::Operator = {
    "+" => ast::Operator::Add,
    "-" => ast::Operator::Sub,
};

Term: ast::Expression = {
    <a:Term> <op:MulOp> <b:Factor> => ast::Expression::Binop { a: Box::new(a), op: op, b: Box::new(b) },
    Factor,
};

MulOp: ast::Operator = {
    "*" => ast::Operator::Mult,
    "/" => ast::Operator::Div,
    "//" => ast::Operator::FloorDiv,
    "%" => ast::Operator::Mod,
    "@" => ast::Operator::MatMult,
};

Factor: ast::Expression = {
    "+" <e:Factor> => ast::Expression::Unop { a: Box::new(e), op: ast::UnaryOperator::Pos },
    "-" <e:Factor> => ast::Expression::Unop { a: Box::new(e), op: ast::UnaryOperator::Neg },
    "~" <e:Factor> => ast::Expression::Unop { a: Box::new(e), op: ast::UnaryOperator::Inv },
    <e:Power> => e,
};

Power: ast::Expression = {
    <e:AtomExpr> <e2:("**" Factor)?> => {
        match e2 {
            None => e,
            Some(x) => ast::Expression::Binop { a: Box::new(e), op: ast::Operator::Pow, b: Box::new(x.1) },
        }
    }
};

AtomExpr: ast::Expression = {
    <e:Atom> => e,
    <f:AtomExpr>  "(" <a:ArgumentList> ")" => ast::Expression::Call { function: Box::new(f), args: a.0, keywords: a.1 },
    <e:AtomExpr>  "[" <s:Subscript> "]" => ast::Expression::Subscript { a: Box::new(e), b: Box::new(s) },
    <e:AtomExpr> "." <n:Identifier> => ast::Expression::Attribute { value: Box::new(e), name: n },
};

Subscript: ast::Expression = {
    <e:Test> => e,
    <e1:Test?> ":" <e2:Test?> <e3:SliceOp?>  => {
        let s1 = e1.unwrap_or(ast::Expression::None);
        let s2 = e2.unwrap_or(ast::Expression::None);
        let s3 = e3.unwrap_or(ast::Expression::None);
        ast::Expression::Slice { elements: vec![s1, s2, s3] }
    }
};

SliceOp: ast::Expression = {
    ":" <e:Test?> => e.unwrap_or(ast::Expression::None)
}

Atom: ast::Expression = {
    StringConstant,
    <n:Number> => ast::Expression::Number { value: n },
    <i:Identifier> => ast::Expression::Identifier { name: i },
    "[" <e:TestListComp?> "]" => {
        let elements = e.unwrap_or(Vec::new());
        ast::Expression::List { elements }
    },
    "[" <e:TestListComp2> "]" => {
        // List comprehension:
        e
    },
    "(" <e:TestList?> <trailing_comma:","?> ")" => {
         match e {
             None => ast::Expression::Tuple { elements: Vec::new() },
             Some(elements) => {
                 if elements.len() == 1 && trailing_comma.is_none() {
                   // This is "(e)", which is equivalent to "e"
                   elements.into_iter().next().unwrap()
                 } else {
                   ast::Expression::Tuple { elements }
                 }
             }
         }
    },
    "(" <e:Test> <c:CompFor> ")" => {
        ast::Expression::Comprehension {
            kind: Box::new(ast::ComprehensionKind::GeneratorExpression { element: e }),
            generators: c,
        }
    },
    "{" <e:TestDict?> "}" => ast::Expression::Dict { elements: e.unwrap_or(Vec::new()) },
    "{" <e:TestDictComp> "}" => e,
    "{" <e:TestSet> "}" => ast::Expression::Set { elements: e },
    "{" <e:TestSetComp> "}" => e,
    "True" => ast::Expression::True,
    "False" => ast::Expression::False,
    "None" => ast::Expression::None,
};

TestListComp: Vec<ast::Expression> = {
    <e:OneOrMore<TestOrStarExpr>> <_trailing_comma:","?> => {
        e
    },
};

TestListComp2: ast::Expression = {
    <e:TestOrStarExpr> <c:CompFor> => {
        ast::Expression::Comprehension {
            kind: Box::new(ast::ComprehensionKind::List { element: e }),
            generators: c,
        }
    },
};

TestDict: Vec<(ast::Expression, ast::Expression)> = {
    <e1:OneOrMore<DictEntry>> <_trailing_comma:","?> => {
        e1
    }
};

TestDictComp: ast::Expression = {
    <e1:DictEntry> <c:CompFor> => {
        ast::Expression::Comprehension {
            kind: Box::new(ast::ComprehensionKind::Dict { key: e1.0, value: e1.1 }),
            generators: c,
        }
    }
};

DictEntry: (ast::Expression, ast::Expression) = {
    <e1: Test> ":" <e2: Test> => (e1, e2),
};

TestSet: Vec<ast::Expression> = {
    <e1:OneOrMore<Test>> ","? => {
        e1
    }
};

TestSetComp: ast::Expression = {
    <e1:Test> <c:CompFor> => {
        ast::Expression::Comprehension {
            kind: Box::new(ast::ComprehensionKind::Set { element: e1 }),
            generators: c,
        }
    }
};

ExpressionList: ast::Expression = {
    <e: ExpressionList2> => {
        if e.len() == 1 {
            e.into_iter().next().unwrap()
        } else {
            ast::Expression::Tuple { elements: e }
        }
    },
};

ExpressionList2: Vec<ast::Expression> = {
    <e1:Expression> <e2:("," Expression)*> ","? => {
        let mut l = vec![e1];
        l.extend(e2.into_iter().map(|x| x.1));
        l
    },
};

#[inline]
TestList: Vec<ast::Expression> = {
  <e1:Test> <e2: ("," Test)*> => {
    let mut l = vec![e1];
    l.extend(e2.into_iter().map(|x| x.1));
    l
  }
};

// Test
StarExpr: ast::Expression = {
    "*" <e:Expression> => ast::Expression::Starred { value: Box::new(e) },
};

// Comprehensions:
CompFor: Vec<ast::Comprehension> = {
    <c:SingleForComprehension+> => c,
};

SingleForComprehension: ast::Comprehension = {
    "for" <e:ExpressionList> "in" <i:OrTest> <c2:ComprehensionIf*> => {
        ast::Comprehension {
            target: e,
            iter: i,
            ifs: c2,
        }
    }
};

ExpressionNoCond: ast::Expression = {
    OrTest,
};

ComprehensionIf: ast::Expression = {
    "if" <c:ExpressionNoCond> => c,
};

ArgumentList: (Vec<ast::Expression>, Vec<ast::Keyword>) = {
    <e: Comma<FunctionArgument>> => {
        let mut args = vec![];
        let mut keywords = vec![];
        for (name, value) in e {
            match name {
                Some(n) => {
                    keywords.push(ast::Keyword { name: n, value: value });
                },
                None => {
                    if keywords.len() > 0 {
                        panic!("positional argument follows keyword argument");
                    };
                    args.push(value);
                },
            }
        }
        (args, keywords)
    }
};

FunctionArgument: (Option<Option<String>>, ast::Expression) = {
    <e:Test> <c:CompFor?> => {
        let expr = match c {
            Some(c) => ast::Expression::Comprehension {
                kind: Box::new(ast::ComprehensionKind::GeneratorExpression { element: e }),
                generators: c,
            },
            None => e,
        };
        (None, expr)
    },
    <i:Identifier> "=" <e:Test> => (Some(Some(i.clone())), e),
    "*" <e:Test> => (None, ast::Expression::Starred { value: Box::new(e) }),
    "**" <e:Test> => (Some(None), e),
};

Comma<T>: Vec<T> = {
    <items: (<T> ",")*> <last: T?> => {
        let mut items = items;
        items.extend(last);
        items
    }
};

#[inline]
OneOrMore<T>: Vec<T> = {
    <i1: T> <i2:("," T)*> => {
        let mut items = vec![i1];
        items.extend(i2.into_iter().map(|e| e.1));
        items
    }
};

Number: ast::Number = {
    <s:int> => { ast::Number::Integer { value: s } },
    <s:float> => { ast::Number::Float { value: s } },
    <s:complex> => { ast::Number::Complex { real: s.0, imag: s.1 } },
};

StringConstant: ast::Expression = {
    <s:string+> => {
        let glued = s.join("");
        ast::Expression::String { value: glued }
    },
    <s:bytes+> => {
        let glued = s.into_iter().flatten().collect::<Vec<u8>>();
        ast::Expression::Bytes { value: glued }
    },
};

Identifier: String = <s:name> => s;

// Hook external lexer:
extern {
    type Location = lexer::Location;
    type Error = lexer::LexicalError;

    enum lexer::Tok {
        indent => lexer::Tok::Indent,
        dedent => lexer::Tok::Dedent,
        StartProgram => lexer::Tok::StartProgram,
        StartStatement => lexer::Tok::StartStatement,
        StartExpression => lexer::Tok::StartExpression,
        "+" => lexer::Tok::Plus,
        "-" => lexer::Tok::Minus,
        "~" => lexer::Tok::Tilde,
        ":" => lexer::Tok::Colon,
        "." => lexer::Tok::Dot,
        "," => lexer::Tok::Comma,
        "*" => lexer::Tok::Star,
        "**" => lexer::Tok::DoubleStar,
        "&" => lexer::Tok::Amper,
        "@" => lexer::Tok::At,
        "%" => lexer::Tok::Percent,
        "//" => lexer::Tok::DoubleSlash,
        "^" => lexer::Tok::CircumFlex,
        "|" => lexer::Tok::Vbar,
        "<<" => lexer::Tok::LeftShift,
        ">>" => lexer::Tok::RightShift,
        "/" => lexer::Tok::Slash,
        "(" => lexer::Tok::Lpar,
        ")" => lexer::Tok::Rpar,
        "[" => lexer::Tok::Lsqb,
        "]" => lexer::Tok::Rsqb,
        "{" => lexer::Tok::Lbrace,
        "}" => lexer::Tok::Rbrace,
        "=" => lexer::Tok::Equal,
        "+=" => lexer::Tok::PlusEqual,
        "-=" => lexer::Tok::MinusEqual,
        "*=" => lexer::Tok::StarEqual,
        "@=" => lexer::Tok::AtEqual,
        "/=" => lexer::Tok::SlashEqual,
        "%=" => lexer::Tok::PercentEqual,
        "&=" => lexer::Tok::AmperEqual,
        "|=" => lexer::Tok::VbarEqual,
        "^=" => lexer::Tok::CircumflexEqual,
        "<<=" => lexer::Tok::LeftShiftEqual,
        ">>=" => lexer::Tok::RightShiftEqual,
        "**=" => lexer::Tok::DoubleStarEqual,
        "//=" => lexer::Tok::DoubleSlashEqual,
        "==" => lexer::Tok::EqEqual,
        "!=" => lexer::Tok::NotEqual,
        "<" => lexer::Tok::Less,
        "<=" => lexer::Tok::LessEqual,
        ">" => lexer::Tok::Greater,
        ">=" => lexer::Tok::GreaterEqual,
        "and" => lexer::Tok::And,
        "as" => lexer::Tok::As,
        "assert" => lexer::Tok::Assert,
        "break" => lexer::Tok::Break,
        "class" => lexer::Tok::Class,
        "continue" => lexer::Tok::Continue,
        "def" => lexer::Tok::Def,
        "del" => lexer::Tok::Del,
        "elif" => lexer::Tok::Elif,
        "else" => lexer::Tok::Else,
        "except" => lexer::Tok::Except,
        "finally" => lexer::Tok::Finally,
        "for" => lexer::Tok::For,
        "from" => lexer::Tok::From,
        "global" => lexer::Tok::Global,
        "if" => lexer::Tok::If,
        "in" => lexer::Tok::In,
        "is" => lexer::Tok::Is,
        "import" => lexer::Tok::Import,
        "from" => lexer::Tok::From,
        "lambda" => lexer::Tok::Lambda,
        "nonlocal" => lexer::Tok::Nonlocal,
        "not" => lexer::Tok::Not,
        "or" => lexer::Tok::Or,
        "pass" => lexer::Tok::Pass,
        "raise" => lexer::Tok::Raise,
        "return" => lexer::Tok::Return,
        "try" => lexer::Tok::Try,
        "while" => lexer::Tok::While,
        "with" => lexer::Tok::With,
        "yield" => lexer::Tok::Yield,
        "True" => lexer::Tok::True,
        "False" => lexer::Tok::False,
        "None" => lexer::Tok::None,
        int => lexer::Tok::Int { value: <BigInt> },
        float => lexer::Tok::Float { value: <f64> },
        complex => lexer::Tok::Complex { real: <f64>, imag: <f64> },
        string => lexer::Tok::String { value: <String> },
        bytes => lexer::Tok::Bytes { value: <Vec<u8>> },
        name => lexer::Tok::Name { name: <String> },
        "\n" => lexer::Tok::Newline,
        ";" => lexer::Tok::Semi,
    }
}