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::DotDot,
        "=>" => Token::Arrow,
        ":=" => Token::Assn,
        "::=" => Token::LitAssn,
        "puts" => Token::Puts,
        "case" => Token::Case,
        "let" => Token::Let,
        "in" => Token::In,
        "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(<>),
    <Name> ":=" <Expr> => Stmt::Assn(<>),
    <name:Name> "::=" <strs:(<"var">)*> =>
        Stmt::LitAssn(name, strs.iter().map(|x| x.to_string()).collect()),
};

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

pub Expr: Expr = {
    <mut ts:(<Choice> "|")*> <t:Choice> => {
        ts.push(t);
        Expr::Chc(ts)
    }
};

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

pub Term: Expr = {
    (<Branch>)* => Expr::Cat(<>),
};

pub Branch: Expr = {
    <l:Branch> "." <r:Subbranch> => Expr::Ap(Box::new(l), Box::new(r)),
    <Subbranch> => <>,
};

pub Subbranch: Expr = {
    <l:Subbranch> ".." <r:Leaf> => Expr::Range(Box::new(l), Box::new(r)),
    <Leaf> => <>,
}

pub Leaf: Expr = {
    <Literal> => Expr::Lit(<>),
    <Name> => Expr::Var(<>),
    "<" <mut es:(<Expr> ",")*> <e:Expr> ">" => {
        es.push(e);
        Expr::Tup(es)
    },
    "let" <name:Name> ":=" <e1:Expr> "in" "{" <e2:Expr> "}" =>
        Expr::Let(name, Box::new(e1), Box::new(e2)),
    "{" <mut cs:(<Case> ";")*> <c:Case> "}" => {
        cs.push(c);
        Expr::Fun(cs)
    },
    "(" <e:Expr> ")" => e,
};

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

pub Pat: Pat = {
    <Literal> => Pat::Lit(<>),
    <Name> => Pat::Var(<>),
    "<" <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(<>)),
};