// See also: file:///usr/share/doc/python/html/reference/grammar.html?highlight=grammar
use super::ast;
use super::lexer;
use std::iter::FromIterator;

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::Statement> = {
    <s:Statement> => Some(s),
    "\n" => None,
};

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

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

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

SmallStatement: ast::Statement = {
  // <e:Expression> => ast::Statement::Expression { expression: e },
  ExpressionStatement,
  "pass" => ast::Statement::Pass,
  FlowStatement,
  ImportStatement,
  AssertStatement,
};

ExpressionStatement: ast::Statement = {
  <e:TestList> <e2:AssignSuffix*> => {
    //match e2 {
    //  None => ast::Statement::Expression { expression: e },
    //  Some(e3) => ast::Statement::Expression { expression: e },
    //}
    if e2.len() > 0 {
      // Dealing with assignment here
      // TODO: for rhs in e2 {
       let rhs = e2.into_iter().next().unwrap();
       // ast::Expression::Tuple { elements: e2.into_iter().next().unwrap()
       let v = rhs.into_iter().next().unwrap();
       let lhs = ast::Statement::Assign { targets: e, value: v };
      lhs
    } else {
      if e.len() > 1 {
        panic!("Not good?");
        // ast::Statement::Expression { expression: e[0] }
      } else {
        ast::Statement::Expression { expression: e.into_iter().next().unwrap() }
      }
    }
  },
  <e1:Test> <op:AugAssign> <e2:TestList> => {
    // TODO: this works in most cases:
    let rhs = e2.into_iter().next().unwrap();
    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::Statement = {
  "break" => ast::Statement::Break,
  "continue" => ast::Statement::Continue,
  "return" <t:TestList?> => ast::Statement::Return { value: t},
  // raise
  // yield
};

ImportStatement: ast::Statement = {
  "import" <n:DottedName> => ast::Statement::Import { name: n },
};

DottedName: String = {
  <n:name> => n,
};

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

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

IfStatement: ast::Statement = {
  "if" <t:Test> ":" <s1:Suite> <s2:("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() {
      last = Some(vec![ast::Statement::If { test: i.1, body: i.3, orelse: last }]);
    }

    ast::Statement::If { test: t, body: s1, orelse: last }
  },
};

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

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

WithStatement: ast::Statement = {
  "with" <t:Test> "as" <_e:Expression> ":" <s:Suite> => ast::Statement::With { items: t, body: s },
};

FuncDef: ast::Statement = {
  "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => ast::Statement::FunctionDef { name: i, args: a, body: s },
};

Parameters: Vec<String> = {
  "(" <a: TypedArgsList> ")" => a,
};

TypedArgsList: Vec<String> = {
  <a: Comma<Identifier>> => a,
};

ClassDef: ast::Statement = {
  "class" <n:Identifier> <_a:("(" ")")?> ":" <_s:Suite> => ast::Statement::ClassDef { name: n },
};

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 = {
  <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?> "]" => {
    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)*> => {
        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<ast::Expression> = {
  <e: Comma<Test>> => e,
};

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

Number: i32 = <s:number> => s;
String: String = {
  <s:string> => s,
};
Identifier: String = <s:name> => s;

// Hook external lexer:
extern {
    type Location = usize;
    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,
        "for" => lexer::Tok::For,
        "if" => lexer::Tok::If,
        "in" => lexer::Tok::In,
        "is" => lexer::Tok::Is,
        "import" => lexer::Tok::Import,
        "not" => lexer::Tok::Not,
        "or" => lexer::Tok::Or,
        "pass" => lexer::Tok::Pass,
        "return" => lexer::Tok::Return,
        "while" => lexer::Tok::While,
        "with" => lexer::Tok::With,
	  "lambda" => lexer::Tok::Lambda,
        "True" => lexer::Tok::True,
        "False" => lexer::Tok::False,
        "None" => lexer::Tok::None,
        number => lexer::Tok::Number { value: <i32> },
        string => lexer::Tok::String { value: <String> },
        name => lexer::Tok::Name { name: <String> },
        "\n" => lexer::Tok::Newline,
        ";" => lexer::Tok::Semi,
    }
}