// See also: file:///usr/share/doc/python/html/reference/grammar.html?highlight=grammar
#![allow(unknown_lints,clippy)]

use super::ast;
use super::lexer;
use std::iter::FromIterator;
use std::str::FromStr;

grammar;

pub 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,
};

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

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

SmallStatement: ast::LocatedStatement = {
  // <e:Expression> => ast::Statement::Expression { expression: e },
  ExpressionStatement,
  <loc:@L> "pass" => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::Pass,
    }
  },
  FlowStatement,
  ImportStatement,
  AssertStatement,
};

ExpressionStatement: ast::LocatedStatement = {
  <loc:@L> <expr:TestList> <suffix:AssignSuffix*> => {
    // Just an expression, no assignment:
    if suffix.is_empty() {
      if expr.len() > 1 {
        ast::LocatedStatement {
          location: loc.clone(),
          node: ast::Statement::Expression { expression: ast::Expression::Tuple { elements: expr } }
        }
      } else {
        ast::LocatedStatement {
          location: loc.clone(),
          node: ast::Statement::Expression { expression: expr[0].clone() },
        }
      }
    } else {
      let mut targets = vec![if expr.len() > 1 {
        ast::Expression::Tuple { elements: expr }
      } else {
        expr[0].clone()
      }];
      let mut values : Vec<ast::Expression> = suffix.into_iter().map(|test_list| if test_list.len() > 1 { ast::Expression::Tuple { elements: test_list }} else { test_list[0].clone() }).collect();

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

      let value = values[0].clone();

      ast::LocatedStatement {
        location: loc.clone(),
        node: ast::Statement::Assign { targets, value },
      }
    }
  },
  <loc:@L> <e1:Test> <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: e1, op: op, value: rhs },
    }
  },
};

AssignSuffix: Vec<ast::Expression> = {
    "=" <e:TestList> => e,
};

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> "raise" <t:Test?> => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::Raise { expression: t },
    }
  },
  <loc:@L> <y:YieldExpr> => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::Expression { expression: y },
    }
  },
};

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
  },
};

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" <i1:WithItem> <i2:("," WithItem)*> ":" <s:Suite> => {
    let mut items = vec![i1];
    for item in i2 {
        items.push(item.1);
    }
    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: Vec<(String, Option<ast::Expression>)> = {
  "(" <a: TypedArgsList> ")" => a,
};

// parameters are (String, None), kwargs are (String, Some(Test)) where Test is
// the default
TypedArgsList: Vec<(String, Option<ast::Expression>)> = {
  <a: Comma<Parameter>> => a
};

Parameter: (String, Option<ast::Expression>) = {
    <i:Identifier> => (i.clone(), None),
    <i:Identifier> "=" <e:Expression> => (i.clone(), Some(e.clone())),
};

ClassDef: ast::LocatedStatement = {
  <d:Decorator*> <loc:@L> "class" <n:Identifier> <a:Parameters?> ":" <s:Suite> => {
    ast::LocatedStatement {
      location: loc,
      node: ast::Statement::ClassDef {
          name: n,
          args: a.unwrap_or(vec![]),
          body: s,
          decorator_list: d,
      },
    }
  },
};

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

YieldExpr: ast::Expression = {
    "yield" => {
        ast::Expression::Yield {
            expression: None,
        }
    }
};

Test: ast::Expression = {
  <e:OrTest> => e,
  <e:LambdaDef> => e,
};

LambdaDef: ast::Expression = {
    "lambda" <p:TypedArgsList> ":" <b:Expression> =>
        ast::Expression::Lambda {
            args:p,
            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,
};

pub 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) },
  <e:AndExpression> => e,
};

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

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> => e,
  "-" <e:Factor> => ast::Expression::Unop { a: Box::new(e), op: ast::UnaryOperator::Neg },
  <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:FunctionArguments> ")" => ast::Expression::Call { function: Box::new(f), args: a },
  <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 = {
  <s:String> => ast::Expression::String { value: s },
  <n:Number> => ast::Expression::Number { value: n },
  <i:Identifier> => ast::Expression::Identifier { name: i },
  "[" <e:TestList?> <_trailing_comma:","?> "]" => {
    match e {
      None => ast::Expression::List { elements: Vec::new() },
      Some(elements) => ast::Expression::List { elements },
    }
  },
  "(" <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:TestDict?> "}" => ast::Expression::Dict { elements: e.unwrap_or(Vec::new()) },
  "True" => ast::Expression::True,
  "False" => ast::Expression::False,
  "None" => ast::Expression::None,
};

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

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

ExpressionList: Vec<ast::Expression> = {
  <e: Comma<Expression>> => e,
};

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

FunctionArguments: Vec<(Option<String>, ast::Expression)> = {
  <e: Comma<FunctionArgument>> => e,
};

FunctionArgument: (Option<String>, ast::Expression) = {
    <e:Expression> => (None, e.clone()),
    <i:Identifier> "=" <e:Expression> => (Some(i.clone()), e.clone()),
};

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

Number: ast::Number = {
  <s:number> => {
    if s.contains(".") {
      ast::Number::Float { value: f64::from_str(&s).unwrap() }
    } else {
      ast::Number::Integer { value: i32::from_str(&s).unwrap() }
    }
  }
};

String: String = {
  <s:string+> => {
    s.join("")
  },
};
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,
        "+" => lexer::Tok::Plus,
        "-" => lexer::Tok::Minus,
        ":" => 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::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::Break,
        "def" => lexer::Tok::Def,
        "elif" => lexer::Tok::Elif,
        "else" => lexer::Tok::Else,
        "except" => lexer::Tok::Except,
        "finally" => lexer::Tok::Finally,
        "for" => lexer::Tok::For,
        "if" => lexer::Tok::If,
        "in" => lexer::Tok::In,
        "is" => lexer::Tok::Is,
        "import" => lexer::Tok::Import,
        "from" => lexer::Tok::From,
        "lambda" => lexer::Tok::Lambda,
        "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,
        number => lexer::Tok::Number { value: <String> },
        string => lexer::Tok::String { value: <String> },
        name => lexer::Tok::Name { name: <String> },
        "\n" => lexer::Tok::Newline,
        ";" => lexer::Tok::Semi,
    }
}