use crate::ast::*;
use crate::lexer::*;

grammar<'input>(ast: &mut ASTArena);

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

    enum Token<'input> {
        "<" => Token::LAngle,
        ">" => Token::RAngle,
        "(" => Token::LPar,
        ")" => Token::RPar,
        "{" => Token::LCurl,
        "}" => Token::RCurl,
        "|" => Token::Pipe,
        ":" => Token::Colon,
        "," => Token::Comma,
        ";" => Token::Semi,
        "." => Token::Dot,
        "_" => Token::Underscore,
        ".." => Token::DotDot,
        "=>" => Token::Arrow,
        ":=" => Token::Assn,
        "::=" => Token::LitAssn,
        "puts" => Token::Puts,
        "case" => Token::Case,
        "let" => Token::Let,
        "in" => Token::In,
        "fix" => Token::Fix,
        "var" => Token::Var(<&'input str>),
        "atom" => Token::Atom(<&'input str>),
        "num" => Token::Num(<i64>),
        "str" => Token::Str(<String>)
    }
}

pub Stmts: Vec<Stmt> = {
    <mut stmts:(<Stmt> ";")*> <stmt:Stmt?> => match stmt {
        None => stmts,
        Some(stmt) => {
            stmts.push(stmt);
            stmts
        }
    },
};

pub Stmt: Stmt = {
    "puts" <Expr> => Stmt::Puts(<>),
    <fixed:"fix"?> <name:Name> ":=" <expr:Expr> => Stmt::Assn(fixed.is_some(), name, expr),
    <fixed:"fix"?> <name:Name> "::=" <strs:(<AtomLit>)*> =>
        Stmt::LitAssn(
            fixed.is_some(),
            name,
            strs,
        ),
    "fix" <Name> => Stmt::Fix(<>),
};

pub AtomLit: String = {
    "var" => <>.to_string(),
    "atom" => <>.to_string(),
};

pub Name: Name = {
    "var" => ast.add_string(<>),
};

pub Expr: ExprRef = {
    <mut ts:(<Choice> "|")*> <t:Choice> => {
        if ts.len() == 0 {
            t.value
        } else {
            ts.push(t);
            ast.add_expr(Expr::Chc(ts))
        }
    }
};

pub Choice: Choice = {
    <weight:"num"> ":" <value:Term> => Choice {
        weight: Some(weight),
        value
    },
    <value:Term> => Choice {
        weight: None,
        value
    }
};

pub Term: ExprRef = {
    <mut bs:(<Branch>)*> => {
        match bs.len() {
            0 => ast.expr_nil(),
            1 => bs.pop().unwrap(),
            _ => ast.add_expr(Expr::Cat(<>)),
        }
    }
};

pub Branch: ExprRef = {
    <l:Branch> "." <r:Subbranch> => ast.add_expr(Expr::Ap(l, r)),
    <Subbranch> => <>,
};

pub Subbranch: ExprRef = {
    <l:Subbranch> ".." <r:Leaf> => ast.add_expr(Expr::Range(l, r)),
    <Leaf> => <>,
}

pub Leaf: ExprRef = {
    <Literal> => ast.add_expr(Expr::Lit(<>)),
    <Name> => ast.add_expr(Expr::Var(<>)),
    "<" <mut es:(<Expr> ",")*> <e:Expr> ">" => {
        if es.len() == 0 && e == ast.expr_nil() {
            ast.add_expr(Expr::Tup(Vec::new()))
        } else {
            es.push(e);
            ast.add_expr(Expr::Tup(es))
        }
    },
    "let" <name:Name> ":=" <e1:Expr> "in" "{" <e2:Expr> "}" =>
        ast.add_expr(Expr::Let(name, e1, e2)),
    "{" <mut cs:(<Case> ";")*> <c:Case> "}" => {
        cs.push(c);
        ast.add_expr(Expr::Fun(cs))
    },
    "(" <e:Expr> ")" => e,
};

pub Case: Case = {
    <pat:Pat> "=>" <expr:Expr> => Case { pat, expr },
};

pub Pat: Pat = {
    "_" => Pat::Wildcard,
    <Literal> => Pat::Lit(<>),
    <Name> => Pat::Var(<>),
    "<" ">" => Pat::Tup(Vec::new()),
    "<" <mut ps:(<Pat> ",")*> <p:Pat> ">" => {
        ps.push(p);
        Pat::Tup(ps)
    },
};

pub Literal: Literal = {
    "num" => Literal::Num(<>),
    "str" => Literal::Str(<>),
    "atom" => Literal::Atom(ast.add_string(<>)),
};