use crate::{
    any,
    ast::untyped::*,
    combine::{impls::choice::Either, impls::separated::SeparatedItems, parser::Parser},
    lexer::span::Spanned,
    many, map, match_either, opt, recursive, sep, seq, seq_pat, token,
};

// LL(1) grammar (left recursion eliminated via head + tail'):
//
// expr         ::= or
//
// or           ::= and ( "||" and )*
//
// and          ::= equality ( "&&" equality )*
//
// equality     ::= comparison ( ( "==" | "!=" ) comparison )*
//
// comparison   ::= term ( ( "<" | "<=" | ">" | ">=" ) term )*
//
// term         ::= factor ( ( "+" | "-" ) factor )*
//
// factor       ::= unary ( ( "*" | "/" ) unary )*
//
// unary        ::= ( "!" | "-" ) unary       -- right recursive (unchanged)
//               |  primary
//
// primary      ::= atom ( suffix )*
//
// suffix       ::= "(" arglist ")"                      -- call
//               |  "." <ident> "(" arglist ")"          -- method call (before field: same prefix)
//               |  "." <ident>                          -- field access
//
// atom         ::= <literal>
//               |  <ident>
//               |  "(" expr ")"
//               |  <block>
//
// arglist      ::= expr ( "," expr )* ","?
//               |  ε

// Folds a left-associative (head, tail') chain into a binary expression tree.
// Encodes the `A'` production: op B A' | ε
macro_rules! left_assoc {
    ($head:expr, $op:expr, $rhs:expr) => {
        seq!($head, many!(seq!($op, $rhs)); |lhs, tail| {
            let mut node: Spanned<Expr> = lhs;
            for pair in tail.value.into_iter() {
                let (op_sp, rhs_sp) = pair.value;
                let new_expr = Expr::binary(
                    op_sp.value,
                    Node::from(node.clone()),
                    Node::from(rhs_sp.clone()),
                );
                node = Spanned::new(new_expr, node.span.union(&rhs_sp.span));
            }
            node.into_inner()
        })
    };
}

pub fn expression(block: impl Parser<Kw, Item = Block> + 'static) -> impl Parser<Kw, Item = Expr> {
    let eq = token!(Eq);
    let ne = token!(Ne);
    let lt = token!(Lt);
    let le = token!(Le);
    let gt = token!(Gt);
    let ge = token!(Ge);
    let add = token!(Add);
    let sub = token!(Sub);
    let mul = token!(Mul);
    let div = token!(Div);
    let not = token!(Not);
    let lparen = token!(LeftParen);
    let rparen = token!(RightParen);
    let ident = token!(Ident(sym) => sym);
    let lit = token!(Literal(lit) => lit);
    let comma = token!(Comma);
    let dot = token!(Dot);

    let if_kw = token!(Keyword(Kw::If));
    let else_kw = token!(Keyword(Kw::Else));

    let or_op = token!(DoublePipe => BinOp::Or);
    let and_op = token!(And        => BinOp::And);

    let unary_op = any!(not, sub; |op| match_either!(op,
        _ => UnaryOp::Not,
        _ => UnaryOp::Neg,
    ));
    let factor_op = any!(mul, div; |op| match_either!(op,
        _ => BinOp::Mul,
        _ => BinOp::Div,
    ));
    let term_op = any!(add, sub; |op| match_either!(op,
        _ => BinOp::Add,
        _ => BinOp::Sub,
    ));
    let comparison_op = any!(lt, le, gt, ge; |op| match_either!(op,
        _ => BinOp::Lt,
        _ => BinOp::Le,
        _ => BinOp::Gt,
        _ => BinOp::Ge,
    ));
    let equality_op = any!(eq, ne; |op| match_either!(op,
        _ => BinOp::Eq,
        _ => BinOp::Ne,
    ));

    recursive!(|expr| {
        // atom ::= <literal> | <ident> | "(" expr ")"
        // Note: ident must come after grouping to avoid consuming "(" as an ident.
        let grouping = seq!(lparen, expr, rparen; |_, e, _| Expr::grouping(Node::from(e)));
        let atom = any!(lit, grouping, ident, block; |p| match_either!(p,
            lit      => Expr::literal(lit),
            grouping => grouping,
            ident    => Expr::ident(ident),
            block    => Expr::block(block),
        ));

        // Postfix suffixes for primary':
        //   call_suffix   ::= "(" arglist ")"
        //   method_suffix ::= "." <ident> "(" arglist ")"   -- must be tried before field_suffix
        //   field_suffix  ::= "." <ident>
        //
        // method_suffix is tried before field_suffix because both start with "." <ident>;
        // the "(" disambiguates and we need to commit to the longer form first.
        let call_suffix = seq!(lparen, sep!(expr, comma), rparen);
        let method_or_field_suffix = seq!(dot, ident, opt!(lparen, sep!(expr, comma), rparen));

        let suffix = any!(method_or_field_suffix, call_suffix);

        // primary ::= atom primary'
        // primary' is encoded as many!(suffix), then folded left into the base atom.
        let primary = seq!(atom, many!(suffix); |base, suffixes| {
            let mut node = base;
            for suffix in suffixes.value.into_iter() {
                match_either!(suffix.value,
                    seq_pat!(_, method_or_field, opt_params) => match opt_params.value {
                        Some(seq_pat!(_, params, rparen)) => {
                            let new_expr = Expr::MethodCall { receiver: Box::new(Node::from(node.clone())), method: method_or_field, args: Node::from_vec(params.value.0) };
                            node = Spanned::new(new_expr, node.span.union(&rparen.span));
                        }
                        None => {
                            let new_expr = Expr::Field { receiver: Box::new(Node::from(node.clone())), name: method_or_field.clone() };
                            node = Spanned::new(new_expr, node.span.union(&method_or_field.span));
                        }
                    },
                    seq_pat!(_, params, rp) => {
                        let new_expr = Expr::CallExpr { callee: Box::new(Node::from(node.clone())), args: Node::from_vec(params.value.0) };
                        node = Spanned::new(new_expr, node.span.union(&rp.span));
                    },
                )
            }
            node.into_inner()
        });

        // unary    ::= ( "!" | "-" ) unary | primary   (right-recursive — intentional)
        let unary = recursive!(|unary| {
            any!(seq!(unary_op, unary), primary; |u| match_either!(u,
                seq_pat!(op, expr) => Expr::unary(*op, Node::from(expr)),
                primary            => primary,
            ))
        });

        // factor'  ::= ( "*" | "/" ) unary factor' | ε
        let factor = left_assoc!(unary, factor_op, unary);
        // term'    ::= ( "+" | "-" ) factor term' | ε
        let term = left_assoc!(factor, term_op, factor);
        // comparison' ::= ( "<" | … ) term comparison' | ε
        let comparison = left_assoc!(term, comparison_op, term);
        // equality'  ::= ( "==" | "!=" ) comparison equality' | ε
        let equality = left_assoc!(comparison, equality_op, comparison);
        // and'     ::= "&&" equality and' | ε
        let and = left_assoc!(equality, and_op, equality);
        // or'      ::= "||" and or' | ε
        let or = left_assoc!(and, or_op, and);

        let if_expr = seq!(
            if_kw,
            grouping,
            expr,
            opt!(else_kw, expr); |_, condition, body, opt_else| {
                Expr::If {
                    cond: Box::new(condition.into()),
                    then: Box::new(body.into()),
                    else_: opt_else.into_inner().map(|seq_pat!(_, expr)| Box::new(expr.into())),
                }
            }
        );

        any!(if_expr, or; |v| match_either!(v,
            if_expr => if_expr,
            expr => expr,
        ))
    })
}

pub fn annotations(expr: impl Parser<Kw, Item = Expr>) -> impl Parser<Kw, Item = Annotations> {
    let lparen = token!(LeftParen);
    let rparen = token!(RightParen);
    let ident = token!(Ident(sym) => sym);
    let hashtag = token!(HashTag);
    let lbrace = token!(LeftBrace);
    let rbrace = token!(RightBrace);
    let colon = token!(Colon);
    let comma = token!(Comma);
    let double_colon = seq!(colon, colon);

    let annotation_call = map!(opt!(lparen, sep!(expr, comma), rparen), |opt| match opt {
        Some(seq_pat!(_, args, _)) => Some(args),
        _ => None,
    });
    let annotation_item = map!(
        opt!(double_colon, ident, annotation_call),
        |field| match field {
            Some(seq_pat!(_, field, call)) => Some((field, call)),
            _ => None,
        }
    );
    let annotation = seq!(hashtag,lbrace,ident, annotation_item,rbrace; |_, _, namespace, item, _| match item.into_inner() {
        None => Annotation::Namespace { namespace },
        Some((field, call)) => match call.into_inner() {
            Some(args) => Annotation::Call { namespace, method: field, args: Node::from_vec(args.into_inner().0) },
            None => Annotation::Field { namespace, name: field },
        }
    });
    many!(annotation; |values| Annotations(values))
}

pub fn statement(expr: impl Parser<Kw, Item = Expr>) -> impl Parser<Kw, Item = Stmt> {
    let let_kw = token!(Keyword(Kw::Let));
    let colon = token!(Colon);
    let assign = token!(Assign);
    let ident = token!(Ident(sym) => sym);

    let let_stmt = seq!(let_kw, ident, opt!(colon, ident), assign, expr; |_, name, ty, _, val| {
        LetStmt {
            name: *name,
            ty: ty.into_inner().map(|v| *v.1),
            value: val.into()
        }
    });

    // let if_stmt = seq!(if_kw, expr, block; |_, cond, then| IfStmt { cond: cond.into(), then: then.into() });

    seq!(any!(let_stmt, expr; |stmt| match_either!(stmt,
        let_stmt => Stmt::Let(let_stmt),
        expr_stmt => Stmt::Expr(expr_stmt)
    )))
}

pub fn params() -> impl Parser<Kw, Item = Vec<Node<Param>>> {
    let colon = token!(Colon);
    let ident = token!(Ident(sym) => sym);
    let rparen = token!(RightParen);
    let comma = token!(Comma);
    let lparen = token!(LeftParen);
    let param = seq!(ident, colon, ident; |name, _, ty| Param { name, ty });
    seq!(lparen, sep!(param, comma), rparen; |_, params, _| Node::from_vec(params.into_inner().0))
}

pub fn ctor_statement(expr: impl Parser<Kw, Item = Expr>) -> impl Parser<Kw, Item = CtorStmt> {
    let ident = token!(Ident(sym) => sym);
    let lparen = token!(LeftParen);
    let rparen = token!(RightParen);
    let comma = token!(Comma);
    seq!(ident, lparen, sep!(expr, comma), rparen; |ty, _, args, _| CtorStmt { ty, args: Node::from_vec(args.value.0) })
}

pub fn layout_declaration(
    ctor_stmt: impl Parser<Kw, Item = CtorStmt>,
) -> impl Parser<Kw, Item = LayoutDecl> {
    let layout_kw = token!(Keyword(Kw::Layout));
    let ident = token!(Ident(sym) => sym);
    let colon = token!(Colon);
    let lbrace = token!(LeftBrace);
    let rbrace = token!(RightBrace);

    let layout_item =
        seq!(ident, colon, ctor_stmt; |name, _, value| Field { name, value: value.into() });
    seq!(layout_kw, lbrace, many!(layout_item), rbrace; |_, _, items, _| LayoutDecl { fields: Node::from_vec(items.value) })
}

pub fn when_declaration(
    expr: impl Parser<Kw, Item = Expr>,
    block: impl Parser<Kw, Item = Block>,
) -> impl Parser<Kw, Item = WhenDecl> {
    let when_kw = token!(Keyword(Kw::When));
    seq!(when_kw, expr, block; |_, cond, body| WhenDecl { cond: cond.into(), body: body.into() })
}

pub fn cond_declaration(expr: impl Parser<Kw, Item = Expr>) -> impl Parser<Kw, Item = CondDecl> {
    let cond_kw = token!(Keyword(Kw::Condition));
    let assign = token!(Assign);
    let ident = token!(Ident(sym) => sym);
    let annotation = annotations(expr.clone());
    seq!(opt!(annotation), cond_kw, ident, assign, expr; |annotation, _, name, _, value| CondDecl { annotations: annotation.into_inner(), name, value: value.into() })
}

pub fn proc_declaration(
    params: impl Parser<Kw, Item = Vec<Node<Param>>>,
    block: impl Parser<Kw, Item = Block>,
    annotation: impl Parser<Kw, Item = Annotations>,
) -> impl Parser<Kw, Item = ProcDecl> {
    let proc_kw = token!(Keyword(Kw::Proc));
    let ident = token!(Ident(sym) => sym);
    let arrow = token!(Arrow);

    seq!(opt!(annotation),proc_kw, ident, params, opt!(arrow, ident), block; |annotation, _, name, params, ret, body| ProcDecl { annotations: annotation.into_inner(), name, params: params.into_inner(), ret: ret.into_inner().map(|seq_pat!(_, ty)| ty), body: body.into() })
}

pub fn group_item(
    expr: impl Parser<Kw, Item = Expr>,
    block: impl Parser<Kw, Item = Block>,
    annotation: impl Parser<Kw, Item = Annotations>,
) -> impl Parser<Kw, Item = GroupItem> {
    let ctor_stmt = ctor_statement(expr.clone());
    let layout_decl = layout_declaration(ctor_stmt);
    let when_decl = when_declaration(expr.clone(), block.clone());
    let params = params();
    let proc_decl = proc_declaration(params, block, annotation);
    let cond_decl = cond_declaration(expr.clone());
    any!(layout_decl, when_decl, proc_decl, cond_decl; |v| match_either!(v,
        layout => GroupItem::Layout(layout),
        when => GroupItem::When(when),
        proc => GroupItem::Proc(proc),
        cond => GroupItem::Cond(cond)
    ))
}

pub fn group(
    expr: impl Parser<Kw, Item = Expr>,
    block: impl Parser<Kw, Item = Block>,
) -> impl Parser<Kw, Item = GroupDecl> {
    let ident = token!(Ident(sym) => sym);
    let group_kw = token!(Keyword(Kw::Group));
    let lbrace = token!(LeftBrace);
    let rbrace = token!(RightBrace);
    let annotation = annotations(expr.clone());
    let group_item = group_item(expr.clone(), block.clone(), annotation.clone());
    let params = params();

    seq!(group_kw, ident, params, lbrace, many!(group_item), rbrace; |_, name, params, _, items, _| GroupDecl { name, params: params.into_inner(), items: Node::from_vec(items.into_inner()) })
}

pub fn block_parser(stmt: impl Parser<Kw, Item = Stmt>) -> impl Parser<Kw, Item = Block> {
    let lbrace = token!(LeftBrace);
    let rbrace = token!(RightBrace);
    let semicolon = token!(Semicolon);
    seq!(lbrace, sep!(stmt, semicolon), rbrace; |_, stmts, _| {
        let SeparatedItems(mut stmts, trailing) = stmts.into_inner();
        let tail = if !trailing { stmts.pop().map(|s| Box::new(s.into())) } else { None };
        Block { stmts: Node::from_vec(stmts), tail }
    })
}

pub fn toplevel_item(
    block: impl Parser<Kw, Item = Block>,
    expr: impl Parser<Kw, Item = Expr>,
) -> impl Parser<Kw, Item = Item> {
    let start_kw = token!(Keyword(Kw::Start));
    let start = seq!(start_kw, block);
    let layout_decl = layout_declaration(ctor_statement(expr.clone()));
    let group_decl = group(expr.clone(), block.clone());
    let proc_decl = proc_declaration(params(), block, annotations(expr.clone()));

    any!(start, layout_decl, group_decl, proc_decl; |v| match_either!(v,
        seq_pat!(_, block) => Item::Start(block.into_inner()),
        layout => Item::Layout(layout),
        group => Item::Group(group),
        proc => Item::Proc(proc)
    ))
}

pub fn parser() -> impl Parser<Kw, Item = crate::ast::untyped::Prog> {
    let block = recursive!(|block| {
        let expr = expression(block.clone());
        let stmt = statement(expr.clone());
        block_parser(stmt)
    });
    let expr = expression(block.clone());
    let item = toplevel_item(block, expr);
    many!(item; |items| Prog { items: Node::from_vec(items) })
}
