// 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

use crate::ast;
use crate::fstring::parse_located_fstring;
use crate::function::{parse_args, parse_params};
use crate::error::LexicalError;
use crate::lexer;
use alloc::{boxed::Box, string::{String, ToString}, vec, vec::Vec};

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:Test> ("\n")* => ast::Top::Expression(e),
};

Program: ast::Program = {
    <lines:FileLine*> => ast::Program {
        statements: lines.into_iter().flatten().collect(),
    },
};

// A file line either has a declaration, or an empty newline:
FileLine: ast::Suite = {
    Statement,
    "\n" => vec![],
};

Suite: ast::Suite = {
    SimpleStatement,
    "\n" Indent <s:Statement+> Dedent => s.into_iter().flatten().collect(),
};

Statement: ast::Suite = {
    SimpleStatement,
    <s:CompoundStatement> => vec![s],
};

SimpleStatement: ast::Suite = {
    <s1:SmallStatement> <s2:(";" SmallStatement)*> ";"? "\n" => {
        let mut statements = vec![s1];
        statements.extend(s2.into_iter().map(|e| e.1));
        statements
    }
};

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

PassStatement: ast::Statement = {
    <location:@L> "pass" => {
        ast::Statement {
            location,
            custom: (),
            node: ast::StatementType::Pass,
        }
    },
};

DelStatement: ast::Statement = {
    <location:@L> "del" <targets:ExpressionList2> => {
        ast::Statement {
            location,
            custom: (),
            node: ast::StatementType::Delete { targets },
        }
    },
};

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

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

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

            ast::Statement {
                custom: (),
                location,
                node: ast::StatementType::Assign { targets, value },
            }
        }
    },
    <location:@L> <target:TestOrStarExprList> <op:AugAssign> <rhs:TestListOrYieldExpr> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::AugAssign {
                target: Box::new(target),
                op,
                value: Box::new(rhs)
            },
        }
    },
    <location:@L> <target:Test> ":" <annotation:Test> <rhs:AssignSuffix?> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::AnnAssign {
                target: Box::new(target),
                annotation: Box::new(annotation),
                value: rhs
            },
        }
    },
};

AssignSuffix: ast::Expression = {
    "=" <e:TestListOrYieldExpr> => e
};

TestListOrYieldExpr: ast::Expression = {
    TestList,
    YieldExpr
}

TestOrStarExprList: ast::Expression = {
    <location:@L> <elements:OneOrMore<TestOrStarExpr>> <comma:","?> => {
        if elements.len() == 1 && comma.is_none() {
            elements.into_iter().next().unwrap()
        } else {
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Tuple { elements }
            }
        }
    }
};

TestOrStarNamedExprList: ast::Expression = {
    <location:@L> <elements:OneOrMore<TestOrStarNamedExpr>> <comma:","?> => {
        if elements.len() == 1 && comma.is_none() {
            elements.into_iter().next().unwrap()
        } else {
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Tuple { elements }
            }
        }
    }
};

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

TestOrStarNamedExpr: ast::Expression = {
    NamedExpressionTest,
    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::Statement = {
    <location:@L> "break" => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Break,
        }
    },
    <location:@L> "continue" => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Continue,
        }
    },
    <location:@L> "return" <value:TestList?> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Return { value },
        }
    },
    <location:@L> <expression:YieldExpr> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Expression { expression },
        }
    },
    RaiseStatement,
};

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

ImportStatement: ast::Statement = {
    <location:@L> "import" <names: OneOrMore<ImportAsAlias<DottedName>>> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Import { names },
        }
    },
    <location:@L> "from" <source:ImportFromLocation> "import" <names: ImportAsNames> => {
        let (level, module) = source;
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::ImportFrom {
                level,
                module,
                names
            },
        }
    },
};

ImportFromLocation: (usize, Option<String>) = {
    <dots: ImportDots*> <name:DottedName> => {
        (dots.iter().sum(), Some(name))
    },
    <dots: ImportDots+> => {
        (dots.iter().sum(), None)
    },
};

ImportDots: usize = {
    "..." => 3,
    "." => 1,
};

ImportAsNames: Vec<ast::ImportSymbol> = {
    <i:OneOrMore<ImportAsAlias<Identifier>>> => i,
    "(" <i:OneOrMore<ImportAsAlias<Identifier>>> ","? ")" => i,
    "*" => {
        // Star import all
        vec![ast::ImportSymbol { symbol: "*".to_string(), alias: None }]
    },
};


#[inline]
ImportAsAlias<I>: ast::ImportSymbol = {
    <symbol:I> <a: ("as" Identifier)?> => ast::ImportSymbol { symbol, alias: 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::Statement = {
    <location:@L> "global" <names:OneOrMore<Identifier>> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Global { names }
        }
    },
};

NonlocalStatement: ast::Statement = {
    <location:@L> "nonlocal" <names:OneOrMore<Identifier>> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Nonlocal { names }
        }
    },
};

AssertStatement: ast::Statement = {
    <location:@L> "assert" <test:Test> <msg: ("," Test)?> => {
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Assert {
                test, msg: msg.map(|e| e.1)
            }
        }
    },
};

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

IfStatement: ast::Statement = {
    <location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(@L "elif" NamedExpressionTest ":" Suite)*> <s3:("else" ":" Suite)?> => {
        // Determine last else:
        let mut last = s3.map(|s| s.2);

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

        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::If { test, body, orelse: last }
        }
    },
};

WhileStatement: ast::Statement = {
    <location:@L> "while" <test:NamedExpressionTest> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
        let orelse = s2.map(|s| s.2);
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::While {
                test,
                body,
                orelse
            },
        }
    },
};

ForStatement: ast::Statement = {
    <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
        let is_async = is_async.is_some();
        let orelse = s2.map(|s| s.2);
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::For {
                is_async,
                target: Box::new(target),
                iter: Box::new(iter),
                body,
                orelse
            },
        }
    },
};

TryStatement: ast::Statement = {
    <location:@L> "try" ":" <body:Suite> <handlers:ExceptClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> => {
        let orelse = else_suite.map(|s| s.2);
        let finalbody = finally.map(|s| s.2);
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Try {
                body,
                handlers,
                orelse,
                finalbody,
            },
        }
    },
    <location:@L> "try" ":" <body:Suite> <finally:("finally" ":" Suite)> => {
        let handlers = vec![];
        let orelse = None;
        let finalbody = Some(finally.2);
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::Try {
                body,
                handlers,
                orelse,
                finalbody,
            },
        }
    },
};

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

WithStatement: ast::Statement = {
    <location:@L> <is_async:"async"?> "with" <items:OneOrMore<WithItem>> ":" <body:Suite> => {
        let is_async = is_async.is_some();
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::With { is_async, items, body },
        }
    },
};

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

FuncDef: ast::Statement = {
    <decorator_list:Decorator*> <location:@L> <is_async:"async"?> "def" <name:Identifier> <args:Parameters> <r:("->" Test)?> ":" <body:Suite> => {
        let is_async = is_async.is_some();
        ast::Statement {
            custom: (),
            location,
            node: ast::StatementType::FunctionDef {
                is_async,
                name,
                args: Box::new(args),
                body,
                decorator_list,
                returns: r.map(|x| x.1),
            }
        }
    },
};

Parameters: ast::Parameters = {
    "(" <a: (ParameterList<TypedParameter>)?> ")" => a.unwrap_or_default(),
};

// Note that this is a macro which is used once for function defs, and
// once for lambda defs.
ParameterList<ArgType>: ast::Parameters = {
    <param1:ParameterDefs<ArgType>> <args2:("," ParameterListStarArgs<ArgType>)?> ","? =>? {
        let posonlyargs_count = param1.0.len();
        let (names, defaults) = parse_params(param1)?;

        // Now gather rest of parameters:
        let (vararg, kwonlyargs, kw_defaults, kwarg) = args2.map_or((None, vec![], vec![], None), |x| x.1);

        Ok(ast::Parameters {
            posonlyargs_count,
            args: names,
            kwonlyargs,
            vararg: vararg.into(),
            kwarg: kwarg.into(),
            defaults,
            kw_defaults,
        })
    },
    <param1:ParameterDefs<ArgType>> <kw:("," KwargParameter<ArgType>)> ","? =>? {
        let posonlyargs_count = param1.0.len();
        let (names, defaults) = parse_params(param1)?;

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

        Ok(ast::Parameters {
            posonlyargs_count,
            args: names,
            kwonlyargs,
            vararg: vararg.into(),
            kwarg: kwarg.into(),
            defaults,
            kw_defaults,
        })
    },
    <params:ParameterListStarArgs<ArgType>> ","? => {
        let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
        ast::Parameters {
            posonlyargs_count: 0,
            args: vec![],
            kwonlyargs,
            vararg: vararg.into(),
            kwarg: kwarg.into(),
            defaults: vec![],
            kw_defaults,
        }
    },
    <kw:KwargParameter<ArgType>> ","? => {
        ast::Parameters {
            posonlyargs_count: 0,
            args: vec![],
            kwonlyargs: vec![],
            vararg: ast::Varargs::None,
            kwarg: Some(kw).into(),
            defaults: vec![],
            kw_defaults: vec![],
        }
    },
};

// Use inline here to make sure the "," is not creating an ambiguity.
#[inline]
ParameterDefs<ArgType>: (Vec<(ast::Parameter, Option<ast::Expression>)>, Vec<(ast::Parameter, Option<ast::Expression>)>) = {
    <args:OneOrMore<ParameterDef<ArgType>>> => {
        (vec![], args)
    },
    <pos_args:OneOrMore<ParameterDef<ArgType>>> "," "/" <args:("," ParameterDef<ArgType>)*>  => {
        (pos_args, args.into_iter().map(|e| e.1).collect())
    },
};

ParameterDef<ArgType>: (ast::Parameter, Option<ast::Expression>) = {
    <i:ArgType> => (i, None),
    <i:ArgType> "=" <e:Test> => (i, Some(e)),
};

UntypedParameter: ast::Parameter = {
    <location:@L> <arg:Identifier> => ast::Parameter { location, arg, annotation: None },
};

TypedParameter: ast::Parameter = {
    <location:@L> <arg:Identifier> <a:(":" Test)?>=> {
        let annotation = a.map(|x| Box::new(x.1));
        ast::Parameter { location, arg, annotation }
    },
};

// Use inline here to make sure the "," is not creating an ambiguity.
// TODO: figure out another grammar that makes this inline no longer required.
#[inline]
ParameterListStarArgs<ArgType>: (Option<Option<ast::Parameter>>, Vec<ast::Parameter>, Vec<Option<ast::Expression>>, Option<Option<ast::Parameter>>) = {
    "*" <va:ArgType?> <kw:("," ParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> => {
        // 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 = kwarg.map(|n| n.1);

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

KwargParameter<ArgType>: Option<ast::Parameter> = {
    "**" <kwarg:ArgType?> => {
        kwarg
    }
};

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

Path: ast::Expression = {
    <location:@L> <n:Identifier> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Identifier { name: n }
    },
    <p:Path> <location:@L> "." <n:name> => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Attribute {
                value: Box::new(p),
                name: n,
            }
        }
    },
};

// Decorators:
Decorator: ast::Expression = {
    <location:@L>"@" <p:Test> "\n" => {
        p
    },
};

YieldExpr: ast::Expression = {
    <location:@L> "yield" <value:TestList?> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Yield { value: value.map(Box::new) }
    },
    <location:@L> "yield" "from" <e:Test> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::YieldFrom { value: Box::new(e) }
    },
};

Test: ast::Expression = {
    <expr:OrTest> <condition: (@L "if" OrTest "else" Test)?> => {
        if let Some(c) = condition {
            ast::Expression {
                location: c.0,
                custom: (),
                node: ast::ExpressionType::IfExpression {
                    test: Box::new(c.2),
                    body: Box::new(expr),
                    orelse: Box::new(c.4),
                }
            }
        } else {
            expr
        }
    },
    LambdaDef,
};

NamedExpressionTest: ast::Expression = {
    <location:@L><left: (Identifier ":=")?> <right:Test> => {
        if let Some(l) = left {
            ast::Expression {
                location: location,
                custom: (),
                node: ast::ExpressionType::NamedExpression {
                        left: Box::new(
                            ast::Expression {
                                location: location,
                                custom: (),
                                node: ast::ExpressionType::Identifier { name: l.0, },
                            }),
                        right: Box::new(right),
                }
            }
        } else {
            right
        }
    }
}

LambdaDef: ast::Expression = {
    <location:@L> "lambda" <p:ParameterList<UntypedParameter>?> ":" <body:Test> =>
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Lambda {
                args: Box::new(p.unwrap_or_default()),
                body: Box::new(body)
            }
       }
}

OrTest: ast::Expression = {
    <e1:AndTest> <location:@L> <e2:("or" AndTest)*> => {
        if e2.is_empty() {
            e1
        } else {
            let mut values = vec![e1];
            values.extend(e2.into_iter().map(|e| e.1));
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::BoolOp { op: ast::BooleanOperator::Or, values }
            }
        }
    },
};

AndTest: ast::Expression = {
    <e1:NotTest> <location:@L> <e2:("and" NotTest)*> => {
        if e2.is_empty() {
            e1
        } else {
            let mut values = vec![e1];
            values.extend(e2.into_iter().map(|e| e.1));
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::BoolOp { op: ast::BooleanOperator::And, values }
            }
        }
    },
};

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

Comparison: ast::Expression = {
    <e:Expression> <location:@L> <comparisons:(CompOp Expression)+> => {
        let mut vals = vec![e];
        let mut ops = vec![];
        for x in comparisons {
            ops.push(x.0);
            vals.push(x.1);
        }
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Compare { vals, ops }
        }
    },
    Expression,
};

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> <location:@L> "|" <e2:XorExpression> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Binop { a: Box::new(e1), op: ast::Operator::BitOr, b: Box::new(e2) }
    },
    XorExpression,
};

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

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

ShiftExpression: ast::Expression = {
    <e1:ShiftExpression> <location:@L> <op:ShiftOp> <e2:ArithmaticExpression> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Binop { a: Box::new(e1), op, b: Box::new(e2) }
    },
    ArithmaticExpression,
};

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

ArithmaticExpression: ast::Expression = {
    <a:ArithmaticExpression> <location:@L> <op:AddOp> <b:Term> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Binop { a: Box::new(a), op, b: Box::new(b) }
    },
    Term,
};

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

Term: ast::Expression = {
    <a:Term> <location:@L> <op:MulOp> <b:Factor> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Binop { a: Box::new(a), 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 = {
    <location:@L> <op:UnOp> <e:Factor> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Unop { a: Box::new(e), op }
    },
    Power,
};

UnOp: ast::UnaryOperator = {
    "+" => ast::UnaryOperator::Pos,
    "-" => ast::UnaryOperator::Neg,
    "~" => ast::UnaryOperator::Inv,
};

Power: ast::Expression = {
    <e:AtomExpr> <e2:(@L "**" Factor)?> => {
        match e2 {
            None => e,
            Some((location, _, b)) => ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Binop { a: Box::new(e), op: ast::Operator::Pow, b: Box::new(b) }
            },
        }
    }
};

AtomExpr: ast::Expression = {
    <location:@L> <is_await:"await"?> <atom:AtomExpr2> => {
        if is_await.is_some() {
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Await { value: Box::new(atom) }
            }
        } else {
            atom
        }
    }
}

AtomExpr2: ast::Expression = {
    Atom,
    <f:AtomExpr2> <location:@L> "(" <a:ArgumentList> ")" => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Call { function: Box::new(f), args: a.args, keywords: a.keywords }
        }
    },
    <e:AtomExpr2> <location:@L> "[" <s:SubscriptList> "]" => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Subscript { a: Box::new(e), b: Box::new(s) }
    },
    <e:AtomExpr2> <location:@L> "." <name:Identifier> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Attribute { value: Box::new(e), name }
    },
};

SubscriptList: ast::Expression = {
    <location:@L> <s1:Subscript> <s2:("," Subscript)*> ","? => {
        if s2.is_empty() {
            s1
        } else {
            let mut dims = vec![s1];
            for x in s2 {
                dims.push(x.1)
            }

            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Tuple { elements: dims },
            }
        }
    }
};

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

SliceOp: ast::Expression = {
    <location:@L> ":" <e:Test?> => e.unwrap_or(ast::Expression {location, custom: (), node: ast::ExpressionType::None})
}

Atom: ast::Expression = {
    <location:@L> <value:StringGroup> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::String { value }
    },
    <location:@L> <value:Bytes> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Bytes { value }
    },
    <location:@L> <value:Number> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Number { value }
    },
    <location:@L> <name:Identifier> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Identifier { name }
    },
    <location:@L> "[" <e:ListLiteralValues?> "]" => {
        let elements = e.unwrap_or_default();
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::List { elements }
        }
    },
    <location:@L> "[" <element:TestOrStarNamedExpr> <generators:CompFor> "]" => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Comprehension {
                kind: Box::new(ast::ComprehensionKind::List { element }),
                generators,
            }
        }
    },
    <location:@L> "(" <elements:TestOrStarNamedExprList?> ")" => {
        elements.unwrap_or(ast::Expression {
             location,
             custom: (),
             node: ast::ExpressionType::Tuple { elements: Vec::new() }
        })
    },
    "(" <e:YieldExpr> ")" => e,
    <location:@L> "(" <element:Test> <generators:CompFor> ")" => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Comprehension {
                kind: Box::new(ast::ComprehensionKind::GeneratorExpression { element }),
                generators,
            }
        }
    },
    <location:@L> "{" <e:DictLiteralValues?> "}" => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Dict { elements: e.unwrap_or_default() }
    },
    <location:@L> "{" <e1:DictEntry> <generators:CompFor> "}" => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Comprehension {
                kind: Box::new(ast::ComprehensionKind::Dict { key: e1.0, value: e1.1 }),
                generators,
            }
        }
    },
    <location:@L> "{" <elements:SetLiteralValues> "}" => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Set { elements }
    },
    <location:@L> "{" <element:Test> <generators:CompFor> "}" => {
        ast::Expression {
            location,
            custom: (),
            node: ast::ExpressionType::Comprehension {
                kind: Box::new(ast::ComprehensionKind::Set { element }),
                generators,
            }
        }
    },
    <location:@L> "True" => ast::Expression { location, custom: (), node: ast::ExpressionType::True },
    <location:@L> "False" => ast::Expression { location, custom: (), node: ast::ExpressionType::False },
    <location:@L> "None" => ast::Expression { location, custom: (), node: ast::ExpressionType::None },
    <location:@L> "..." => ast::Expression { location, custom: (), node: ast::ExpressionType::Ellipsis },
};

ListLiteralValues: Vec<ast::Expression> = {
    <e:OneOrMore<TestOrStarNamedExpr>> ","? => e,
};

DictLiteralValues: Vec<(Option<ast::Expression>, ast::Expression)> = {
    <elements:OneOrMore<DictElement>> ","? => elements,
};

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

DictElement: (Option<ast::Expression>, ast::Expression) = {
    <e:DictEntry> => (Some(e.0), e.1),
    "**" <e:Expression> => (None, e),
};

SetLiteralValues: Vec<ast::Expression> = {
    <e1:OneOrMore<TestOrStarNamedExpr>> ","? => e1
};

ExpressionOrStarExpression = {
    Expression,
    StarExpr
};

ExpressionList: ast::Expression = {
    <location:@L> <elements: OneOrMore<ExpressionOrStarExpression>> <trailing_comma:","?> => {
        if elements.len() == 1 && trailing_comma.is_none() {
            elements.into_iter().next().unwrap()
        } else {
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Tuple { elements },
            }
        }
    },
};

ExpressionList2: Vec<ast::Expression> = {
    <elements:OneOrMore<Expression>> ","? => elements,
};

// A test list is one of:
// - a list of expressions
// - a single expression
// - a single expression followed by a trailing comma
TestList: ast::Expression = {
    <location:@L> <elements:OneOrMore<TestOrStarExpr>> <trailing_comma: ","?> => {
        if elements.len() == 1 && trailing_comma.is_none() {
            elements.into_iter().next().unwrap()
        } else {
            ast::Expression {
                location,
                custom: (),
                node: ast::ExpressionType::Tuple { elements },
            }
        }
    }
};

// Test
StarExpr: ast::Expression = {
    <location:@L> "*" <e:Expression> => ast::Expression {
        location,
        custom: (),
        node: ast::ExpressionType::Starred { value: Box::new(e) },
    }
};

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

SingleForComprehension: ast::Comprehension = {
    <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:OrTest> <ifs:ComprehensionIf*> => {
        let is_async = is_async.is_some();
        ast::Comprehension { location, target, iter, ifs, is_async }
    }
};

ExpressionNoCond: ast::Expression = OrTest;
ComprehensionIf: ast::Expression = "if" <c:ExpressionNoCond> => c;

ArgumentList: ast::ArgumentList = {
    <e: Comma<FunctionArgument>> =>? {
        let arg_list = parse_args(e)?;
        Ok(arg_list)
    }
};

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

#[inline]
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 = {
    <value:int> => { ast::Number::Integer { value } },
    <value:float> => { ast::Number::Float { value } },
    <s:complex> => { ast::Number::Complex { real: s.0, imag: s.1 } },
};

StringGroup: ast::StringGroup = {
    <loc:@L> <s:string+> =>? {
        let mut values = vec![];
        for (value, is_fstring) in s {
            values.push(if is_fstring {
                parse_located_fstring(&value, loc)?
            } else {
                ast::StringGroup::Constant { value }
            })
        }

        Ok(if values.len() > 1 {
            ast::StringGroup::Joined { values }
        } else {
            values.into_iter().next().unwrap()
        })
    },
};

Bytes: Vec<u8> = {
    <s:bytes+> => {
        s.into_iter().flatten().collect::<Vec<u8>>()
    },
};

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

// Hook external lexer:
extern {
    type Location = ast::Location;
    type Error = 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::Ellipsis,
        "," => 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::ColonEqual,
        "==" => lexer::Tok::EqEqual,
        "!=" => lexer::Tok::NotEqual,
        "<" => lexer::Tok::Less,
        "<=" => lexer::Tok::LessEqual,
        ">" => lexer::Tok::Greater,
        ">=" => lexer::Tok::GreaterEqual,
        "->" => lexer::Tok::Rarrow,
        "and" => lexer::Tok::And,
        "as" => lexer::Tok::As,
        "assert" => lexer::Tok::Assert,
        "async" => lexer::Tok::Async,
        "await" => lexer::Tok::Await,
        "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>, is_fstring: <bool> },
        bytes => lexer::Tok::Bytes { value: <Vec<u8>> },
        name => lexer::Tok::Name { name: <String> },
        "\n" => lexer::Tok::Newline,
        ";" => lexer::Tok::Semi,
    }
}