use std::collections::BTreeMap;
use std::iter;
use std::iter::FromIterator;

use crate::core;
use crate::core::bx;
use crate::core::Expr::*;
use crate::core::BuiltinType::*;
use crate::grammar_util::*;
use crate::lexer::*;

grammar<'input>;

extern {
    type Location = usize;
    type Error = LexicalError;

    enum Tok<'input> {
        Pi => Tok::Pi,
        Lambda => Tok::Lambda,
        Combine => Tok::Combine,
        "->" => Tok::Arrow,

        Int => Tok::Integer(<isize>),
        Nat => Tok::Natural(<usize>),
        Text => Tok::Text(<String>),
        Bool => Tok::Bool(<bool>),
        Label => Tok::Identifier(<&'input str>),
        Const => Tok::Const(<core::Const>),
        Let => Tok::Keyword(Keyword::Let),
        In => Tok::Keyword(Keyword::In),
        If => Tok::Keyword(Keyword::If),
        Then => Tok::Keyword(Keyword::Then),
        Else => Tok::Keyword(Keyword::Else),
        List => Tok::ListLike(ListLike::List),
        Optional => Tok::ListLike(ListLike::Optional),
        Builtin => Tok::Builtin(<Builtin>),

        "{" => Tok::BraceL,
        "}" => Tok::BraceR,
        "[" => Tok::BracketL,
        "]" => Tok::BracketR,
        "(" => Tok::ParenL,
        ")" => Tok::ParenR,
        "&&" => Tok::BoolAnd,
        "||" => Tok::BoolOr,
        "==" => Tok::CompareEQ,
        "!=" => Tok::CompareNE,
        "++" => Tok::Append,
        "*" => Tok::Times,
        "+" => Tok::Plus,
        "," => Tok::Comma,
        "." => Tok::Dot,
        ":" => Tok::Ascription,
        "=" => Tok::Equals,
    }
}

pub Expr: BoxExpr<'input> = { // exprA
    <ExprB> ":" <Expr> => bx(Annot(<>)),
    ExprB,
};

ExprB: BoxExpr<'input> = {
    Lambda "(" <Label> ":" <Expr> ")" "->" <ExprB> => bx(Lam(<>)),
    Pi "(" <Label> ":" <Expr> ")" "->" <ExprB> => bx(Pi(<>)),
    If <Expr> Then <ExprB> Else <ExprC> => bx(BoolIf(<>)),
    <ExprC> "->" <ExprB> => bx(Pi("_", <>)),
    Let <Label> <(":" <Expr>)?> "=" <Expr> In <ExprB> => bx(Let(<>)),
    "[" <a:Elems> "]" ":" <b:ListLike> <c:ExprE> => bx(b(c, a)),
    ExprC,
};

ListLike: ExprListFn<'input> = {
    List => ListLit,
    Optional => OptionalLit,
};

BoolOr: ExprOpFn<'input> = { "||" => BoolOr };
NaturalPlus: ExprOpFn<'input> = { "+" => NaturalPlus };
TextAppend: ExprOpFn<'input> = { "++" => TextAppend };
BoolAnd: ExprOpFn<'input> = { "&&" => BoolAnd };
CombineOp: ExprOpFn<'input> = { Combine => Combine };
NaturalTimes: ExprOpFn<'input> = { "*" => NaturalTimes };
BoolEQ: ExprOpFn<'input> = { "==" => BoolEQ };
BoolNE: ExprOpFn<'input> = { "!=" => BoolNE };

Tier<NextTier, Op>: BoxExpr<'input> = {
    <a:NextTier> <f:Op> <b:Tier<NextTier, Op>> => bx(f(a, b)),
    // <b:Tier<NextTier, Op>> <f:Op> <a:NextTier> => bx(f(a, b)),
    NextTier,
};

ExprC = Tier<ExprC1, BoolOr>;
ExprC1 = Tier<ExprC2, NaturalPlus>;
ExprC2 = Tier<ExprC3, TextAppend>;
ExprC3 = Tier<ExprC4, BoolAnd>;
ExprC4 = Tier<ExprC5, CombineOp>;
ExprC5 = Tier<ExprC6, NaturalTimes>;
ExprC6 = Tier<ExprC7, BoolEQ>;
ExprC7 = Tier<ExprD, BoolNE>;

ExprD: BoxExpr<'input> = {
    <v:(ExprE)+> => {
        let mut it = v.into_iter();
        let f = it.next().unwrap();
        it.fold(f, |f, x| bx(App(f, x)))
    }
};

ExprE: BoxExpr<'input> = {
    <a:ExprF> <fields:("." <Label>)*> => {
        fields.into_iter().fold(a, |x, f| bx(Field(x, f)))
    },
};

ExprF: BoxExpr<'input> = {
    Nat => bx(NaturalLit(<>)),
    Int => bx(IntegerLit(<>)),
    Text => bx(TextLit(<>)),
    Label => bx(Var(core::V(<>, 0))), // FIXME support var@n syntax
    Const => bx(Const(<>)),
    List => bx(BuiltinType(List)),
    Optional => bx(BuiltinType(Optional)),
    Builtin => bx(builtin_expr(<>)),
    Bool => bx(BoolLit(<>)),
    Record,
    RecordLit,
    "(" <Expr> ")",
};

SepBy<S, T>: iter::Chain<::std::vec::IntoIter<T>, ::std::option::IntoIter<T>> = {
    <v:(<T> S)*> <last:T?> => v.into_iter().chain(last.into_iter()),
};

SepBy1<S, T>: iter::Chain<::std::vec::IntoIter<T>, iter::Once<T>> = {
    <v:(<T> S)*> <last:T> => v.into_iter().chain(iter::once(last)),
};

Elems: Vec<ParsedExpr<'input>> = {
    <v:SepBy<",", Expr>> => {
        v.into_iter()
         .map(|b| *b)
         .collect::<Vec<_>>()
    }
};

RecordLit: BoxExpr<'input> = {
    "{" "=" "}" => bx(RecordLit(BTreeMap::new())),
    "{" <FieldValues> "}" => bx(RecordLit(BTreeMap::from_iter(<>))),
};

Record: BoxExpr<'input> = {
    "{" <FieldTypes> "}" => bx(Record(BTreeMap::from_iter(<>))),
};

FieldValues = SepBy1<",", Field<"=">>;
FieldTypes = SepBy<",", Field<":">>;

Field<Sep>: (&'input str, ParsedExpr<'input>) = {
    <a:Label> Sep <b:Expr> => (a, *b),
};