mod tests;

use std::collections::VecDeque;

use crate::ast::{
    AllocKind, Block, BlockInner, Call, ComptimeValue, Decl, DeclKind, Expr, ExprKind, IfElse,
    Match, MatchArm, Module, Param, Params, Pat, PatKind, RecordField, Region, Signature,
    SignatureInner, Stmt, StmtKind, Type, TypeAliasDefinition, TypeKind, ValueKind, VariantField,
};
use crate::ctx::{Ctx, Symbol};
use crate::lex::{BinOp, Keyword, Token};
use crate::span::{Diagnostic, Span, SpanExt, Spanned};

#[derive(Debug, PartialEq, Eq)]
pub struct ParseError {
    pub kind: ParseErrorKind,
    pub span: Span,
}

#[derive(Debug, PartialEq, Eq)]
pub enum ParseErrorKind {
    ExpectedRegion,
    CannotAllocateInRegion,
    ExpectedIdentifier,
    MalformedAccess,
    ExpectedType,
    ExpectedPattern,
    ExpectedExpression,
    ExpectedToken(Token),
    MalformedRange,
    UnexpectedEOF,
}

impl ParseError {
    fn new(kind: ParseErrorKind, span: Span) -> Self {
        Self { kind, span }
    }

    fn eof() -> Self {
        Self::new(ParseErrorKind::UnexpectedEOF, 0..0)
    }
}

impl Diagnostic for ParseError {
    fn span(&self) -> &Span {
        &self.span
    }

    fn message(&self) -> String {
        "parse error".to_string()
    }

    fn label(&self) -> Option<String> {
        Some(format!("{:?}", self.kind))
    }
}

type ParseResult<T> = Result<T, ParseError>;

type SpannedToken = Spanned<Token>;

fn expect_next(tokens: &mut VecDeque<SpannedToken>, expected: Token) -> ParseResult<Span> {
    match tokens.pop_front() {
        Some(Spanned { node, span }) if node == expected => Ok(span),
        Some(Spanned { span, .. }) => Err(ParseError::new(
            ParseErrorKind::ExpectedToken(expected),
            span,
        )),
        None => Err(ParseError::eof()),
    }
}

fn expect_identifier(tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Spanned<Symbol>> {
    match tokens.pop_front() {
        Some(Spanned {
            node: Token::Identifier(name),
            span,
        }) => Ok(Spanned::new(name, span)),
        Some(Spanned { span, .. }) => {
            Err(ParseError::new(ParseErrorKind::ExpectedIdentifier, span))
        }
        None => Err(ParseError::eof()),
    }
}

fn peek_token(tokens: &VecDeque<SpannedToken>) -> Option<&Token> {
    tokens.front().map(|t| &t.node)
}

fn peek_span(tokens: &VecDeque<SpannedToken>) -> Option<Span> {
    tokens.front().map(|t| t.span.clone())
}

fn parse_potential_region_alloc(
    ctx: &mut Ctx,
    mut expr: Expr,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Expr> {
    if let Some(Token::At) = peek_token(tokens) {
        let at_span = tokens.pop_front().unwrap().span;
        let region_expr = parse_expr(ctx, tokens)?;
        let span = at_span.merge(&region_expr.span);

        if let ExprKind::Allocation {
            kind: _,
            elements: _,
            default_elem: _,
            region,
        } = &mut expr.node
        {
            let region_kind = match region_expr.node {
                ExprKind::Value(ValueKind::Ident(name)) => {
                    if ctx.resolve(name) == "local" {
                        Region::Scoped(0)
                    } else {
                        Region::Named(name)
                    }
                }
                ExprKind::Comptime(ComptimeValue::Region(r)) => r,
                _ => {
                    return Err(ParseError::new(
                        ParseErrorKind::ExpectedRegion,
                        region_expr.span,
                    ));
                }
            };

            *region = Some(region_kind);
        } else {
            return Err(ParseError::new(
                ParseErrorKind::CannotAllocateInRegion,
                expr.span,
            ));
        }

        expr.span = span;
    }

    Ok(expr)
}

fn parse_expr(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let expr = parse_logical_or(ctx, tokens)?;
    if !ctx.parsing_region {
        ctx.parsing_region = true;
        let expr = parse_potential_region_alloc(ctx, expr, tokens);
        ctx.parsing_region = false;
        expr
    } else {
        Ok(expr)
    }
}

fn parse_logical_or(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut left = parse_logical_and(ctx, tokens)?;

    while let Some(Token::BinOp(BinOp::Or)) = peek_token(tokens) {
        let _op_token = tokens.pop_front().unwrap();
        let right = parse_logical_and(ctx, tokens)?;
        let span = left.span.merge(&right.span);
        left = Spanned::new(
            ExprKind::BinOp {
                op: BinOp::Or,
                lhs: Box::new(left),
                rhs: Box::new(right),
            },
            span,
        );
    }

    Ok(left)
}

fn parse_logical_and(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut left = parse_comparison(ctx, tokens)?;

    while let Some(Token::BinOp(BinOp::And)) = peek_token(tokens) {
        let _op_token = tokens.pop_front().unwrap();
        let right = parse_comparison(ctx, tokens)?;
        let span = left.span.merge(&right.span);
        left = Spanned::new(
            ExprKind::BinOp {
                op: BinOp::And,
                lhs: Box::new(left),
                rhs: Box::new(right),
            },
            span,
        );
    }

    Ok(left)
}

fn parse_comparison(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut left = parse_additive(ctx, tokens)?;

    while let Some(op) = peek_token(tokens) {
        match op {
            Token::BinOp(
                op @ (BinOp::EqEq | BinOp::NEq | BinOp::Gt | BinOp::GtEq | BinOp::Lt | BinOp::LtEq),
            ) => {
                let op = *op;
                tokens.pop_front();
                let right = parse_additive(ctx, tokens)?;
                let span = left.span.merge(&right.span);
                left = Spanned::new(
                    ExprKind::BinOp {
                        op,
                        lhs: Box::new(left),
                        rhs: Box::new(right),
                    },
                    span,
                );
            }
            _ => break,
        }
    }

    Ok(left)
}

fn parse_additive(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut left = parse_multiplicative(ctx, tokens)?;

    while let Some(op) = peek_token(tokens) {
        match op {
            Token::BinOp(BinOp::Add) | Token::BinOp(BinOp::Sub) => {
                let Spanned {
                    node: Token::BinOp(op),
                    ..
                } = tokens.pop_front().unwrap()
                else {
                    unreachable!()
                };
                let right = parse_multiplicative(ctx, tokens)?;
                let span = left.span.merge(&right.span);
                left = Spanned::new(
                    ExprKind::BinOp {
                        op,
                        lhs: Box::new(left),
                        rhs: Box::new(right),
                    },
                    span,
                );
            }
            _ => break,
        }
    }

    Ok(left)
}

fn parse_postfix(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut expr = parse_primary(ctx, tokens)?;

    loop {
        match peek_token(tokens) {
            Some(Token::Dot) => {
                tokens.pop_front();
                if let Some(Token::Identifier(_)) = peek_token(tokens) {
                    let field = expect_identifier(tokens)?;

                    let span = expr.span.merge(&field.span);
                    expr = Spanned::new(ExprKind::FieldAccess(Box::new(expr), field.node), span);
                } else if let Some(Token::Int(_)) = peek_token(tokens) {
                    let Some(Spanned { node, span }) = tokens.pop_front() else {
                        unreachable!()
                    };

                    let Token::Int(i) = node else { unreachable!() };

                    let index: usize = match (i).try_into() {
                        Ok(i) => i,
                        Err(_) => {
                            return Err(ParseError::new(ParseErrorKind::MalformedAccess, span));
                        }
                    };

                    let span = expr.span.merge(&span);
                    expr = Spanned::new(ExprKind::TupleAccess(Box::new(expr), index), span);
                } else {
                    return Err(ParseError::new(
                        ParseErrorKind::MalformedAccess,
                        peek_span(tokens).unwrap_or(expr.span),
                    ));
                }
            }
            Some(Token::LBracket) => {
                tokens.pop_front();
                let index = parse_expr(ctx, tokens)?;
                expect_next(tokens, Token::RBracket)?;
                let span = expr.span.merge(&index.span);
                expr = Spanned::new(ExprKind::Index(Box::new(expr), Box::new(index)), span);
            }
            Some(Token::LParen) => {
                tokens.pop_front();
                let mut args = Vec::new();

                loop {
                    if let Some(Token::RParen) = peek_token(tokens) {
                        break;
                    }
                    args.push(parse_expr(ctx, tokens)?);
                    if let Some(Token::Comma) = peek_token(tokens) {
                        tokens.pop_front();
                    } else {
                        break;
                    }
                }

                let end_span = expect_next(tokens, Token::RParen)?;
                let span = expr.span.merge(&end_span);

                expr = Spanned::new(
                    ExprKind::Call(Call {
                        callee: Box::new(expr),
                        args,
                        returned_ty: None,
                    }),
                    span,
                );
            }
            _ => break,
        }
    }

    Ok(expr)
}

fn parse_multiplicative(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let mut left = parse_postfix(ctx, tokens)?;

    while let Some(op) = peek_token(tokens) {
        match op {
            Token::BinOp(BinOp::Mul) | Token::BinOp(BinOp::Div) | Token::BinOp(BinOp::Mod) => {
                let Spanned {
                    node: Token::BinOp(op),
                    ..
                } = tokens.pop_front().unwrap()
                else {
                    unreachable!()
                };
                let right = parse_postfix(ctx, tokens)?;
                let span = left.span.merge(&right.span);
                left = Spanned::new(
                    ExprKind::BinOp {
                        op,
                        lhs: Box::new(left),
                        rhs: Box::new(right),
                    },
                    span,
                );
            }
            _ => break,
        }
    }

    Ok(left)
}

fn parse_identifier_expr(
    ctx: &mut Ctx,
    ident: Spanned<Symbol>,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Expr> {
    if let Some(Token::LParen) = peek_token(tokens) {
        let lparen_span = tokens.pop_front().unwrap().span;
        let mut args = Vec::new();

        loop {
            if let Some(Token::RParen) = peek_token(tokens) {
                break;
            }

            args.push(parse_expr(ctx, tokens)?);

            if let Some(Token::Comma) = peek_token(tokens) {
                tokens.pop_front();
            } else {
                break;
            }
        }

        let rparen_span = expect_next(tokens, Token::RParen)?;
        let span = ident.span.merge(&lparen_span).merge(&rparen_span);

        let ident_expr = Spanned::new(ExprKind::Value(ValueKind::Ident(ident.node)), span.clone());

        Ok(Spanned::new(
            ExprKind::Call(Call {
                callee: Box::new(ident_expr),
                args,
                returned_ty: None,
            }),
            span,
        ))
    } else {
        Ok(Spanned::new(
            ExprKind::Value(ValueKind::Ident(ident.node)),
            ident.span,
        ))
    }
}

fn parse_variant_expr(
    ctx: &mut Ctx,
    variant_name: Symbol,
    span: Span,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Expr> {
    if let Some(Token::LParen) = peek_token(tokens) {
        tokens.pop_front();
        let mut args = Vec::new();

        loop {
            if let Some(Token::RParen) = peek_token(tokens) {
                break;
            }

            args.push(parse_expr(ctx, tokens)?);

            if let Some(Token::Comma) = peek_token(tokens) {
                tokens.pop_front();
            } else {
                break;
            }
        }

        let rparen_span = expect_next(tokens, Token::RParen)?;
        let span = span.merge(&rparen_span);

        Ok(Spanned::new(
            ExprKind::Allocation {
                kind: AllocKind::Variant(variant_name),
                elements: args,
                default_elem: None,
                region: None,
            },
            span,
        ))
    } else {
        Ok(Spanned::new(
            ExprKind::Allocation {
                kind: AllocKind::Variant(variant_name),
                elements: vec![],
                default_elem: None,
                region: None,
            },
            span,
        ))
    }
}

fn parse_primary(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Expr> {
    let Some(token) = tokens.pop_front() else {
        return Err(ParseError::eof());
    };

    match token.node {
        Token::Int(i) => {
            if let Some(Token::DotDot) = peek_token(tokens) {
                expect_next(tokens, Token::DotDot)?;

                let end_range = tokens.pop_front().unwrap();
                let Token::Int(end_range_int) = end_range.node else {
                    return Err(ParseError::new(ParseErrorKind::MalformedRange, token.span));
                };

                let total_span = token.span.merge(&end_range.span);
                Ok(Spanned::new(
                    ExprKind::Range {
                        start: Box::new(Spanned::new(
                            ExprKind::Value(ValueKind::Int(i)),
                            token.span,
                        )),
                        end: Box::new(Spanned::new(
                            ExprKind::Value(ValueKind::Int(end_range_int)),
                            end_range.span,
                        )),
                        inclusive: false,
                    },
                    total_span,
                ))
            } else {
                Ok(Spanned::new(ExprKind::Value(ValueKind::Int(i)), token.span))
            }
        }
        Token::Keyword(Keyword::True) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Bool(true)),
            token.span,
        )),
        Token::Keyword(Keyword::False) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Bool(false)),
            token.span,
        )),
        Token::Keyword(Keyword::Int) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Type(TypeKind::Int)),
            token.span,
        )),
        Token::Keyword(Keyword::Bool) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Type(TypeKind::Bool)),
            token.span,
        )),
        Token::Keyword(Keyword::Char) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Type(TypeKind::Char)),
            token.span,
        )),
        Token::Keyword(Keyword::Type) => Ok(Spanned::new(
            ExprKind::Value(ValueKind::Type(TypeKind::Type)),
            token.span,
        )),
        Token::String(s) => Ok(Spanned::new(
            ExprKind::Allocation {
                kind: AllocKind::Str(s),
                elements: vec![],
                default_elem: None,
                region: None,
            },
            token.span,
        )),
        Token::Identifier(name) => {
            if ctx.resolve(name) == "local" {
                return Ok(Spanned::new(
                    ExprKind::Comptime(ComptimeValue::Region(Region::Scoped(0))),
                    token.span,
                ));
            }

            parse_identifier_expr(ctx, Spanned::new(name, token.span), tokens)
        }

        Token::Dot => {
            let dot_span = token.span.clone();
            let Spanned {
                node: Token::Identifier(variant_name),
                span: ident_span,
            } = tokens.pop_front().unwrap()
            else {
                return Err(ParseError::new(
                    ParseErrorKind::ExpectedIdentifier,
                    dot_span,
                ));
            };
            let span = dot_span.merge(&ident_span);
            parse_variant_expr(ctx, variant_name, span, tokens)
        }

        // array expr
        Token::LBracket => {
            let start_span = token.span;

            // Handle empty list []
            if let Some(Token::RBracket) = peek_token(tokens) {
                let end_span = expect_next(tokens, Token::RBracket)?;
                return Ok(Spanned::new(
                    ExprKind::Allocation {
                        kind: AllocKind::Array(TypeKind::Any.into(), ComptimeValue::Int(0).into()),
                        elements: vec![],
                        default_elem: None,
                        region: None,
                    },
                    start_span.merge(&end_span),
                ));
            }

            // Parse first expression
            let first_expr = parse_expr(ctx, tokens)?;

            match peek_token(tokens) {
                // [expr; size] - default array syntax
                Some(Token::SemiColon) => {
                    tokens.pop_front(); // consume ;
                    let size_expr = parse_expr(ctx, tokens)?;
                    let size_cv = match &size_expr.node {
                        ExprKind::Value(ValueKind::Int(i)) => ComptimeValue::Int(*i as i128),
                        ExprKind::Value(ValueKind::Ident(name)) => ComptimeValue::Ident(*name),
                        _ => {
                            return Err(ParseError::new(
                                ParseErrorKind::ExpectedExpression,
                                size_expr.span,
                            ));
                        }
                    };
                    expect_next(tokens, Token::RBracket)?;
                    let region = check_region_type_annotation(ctx, tokens)?;
                    let span = start_span.merge(&size_expr.span);

                    Ok(Spanned::new(
                        ExprKind::Allocation {
                            kind: AllocKind::Array(TypeKind::Any.into(), size_cv.into()),
                            elements: vec![],
                            default_elem: Some(Box::new(first_expr)),
                            region,
                        },
                        span,
                    ))
                }

                // [expr, expr, ...] or [expr] - list literal syntax
                Some(Token::Comma) | Some(Token::RBracket) => {
                    let mut exprs = vec![first_expr];

                    while let Some(Token::Comma) = peek_token(tokens) {
                        tokens.pop_front(); // consume ,
                        if let Some(Token::RBracket) = peek_token(tokens) {
                            break; // trailing comma
                        }
                        exprs.push(parse_expr(ctx, tokens)?);
                    }

                    let end_span = expect_next(tokens, Token::RBracket)?;
                    let span = start_span.merge(&end_span);

                    Ok(Spanned::new(
                        ExprKind::Allocation {
                            kind: AllocKind::Array(
                                TypeKind::Any.into(),
                                ComptimeValue::Int(exprs.len() as i128).into(),
                            ),
                            elements: exprs,
                            default_elem: None,
                            region: None,
                        },
                        span,
                    ))
                }

                _ => Err(ParseError::new(
                    ParseErrorKind::ExpectedExpression,
                    peek_span(tokens).unwrap_or(start_span),
                )),
            }
        }

        // Token::LBracket => {
        //     let start_span = token.span;
        //
        //     let is_default_array = match peek_token(tokens) {
        //         Some(Token::SemiColon) => true,
        //         Some(Token::Int(_) | Token::Identifier(_) | Token::Keyword(_)) => {
        //             let mut temp_tokens = tokens.clone();
        //             let _ = parse_expr(ctx, &mut temp_tokens);
        //             matches!(peek_token(&temp_tokens), Some(Token::SemiColon))
        //         }
        //         _ => false,
        //     };
        //
        //     if is_default_array {
        //         let expr = parse_expr(ctx, tokens)?;
        //         let _ = expect_next(tokens, Token::SemiColon)?;
        //         let default_elem = Box::new(expr);
        //
        //         let size_expr = parse_expr(ctx, tokens)?;
        //         let size_cv = match &size_expr.node {
        //             ExprKind::Value(ValueKind::Int(i)) => ComptimeValue::Int(*i as i128),
        //             ExprKind::Value(ValueKind::Ident(name)) => ComptimeValue::Ident(*name),
        //             _ => {
        //                 return Err(ParseError::new(
        //                     ParseErrorKind::ExpectedExpression,
        //                     size_expr.span,
        //                 ));
        //             }
        //         };
        //         expect_next(tokens, Token::RBracket)?;
        //
        //         let region = check_region_type_annotation(ctx, tokens)?;
        //
        //         let span = start_span.merge(&size_expr.span);
        //
        //         Ok(Spanned::new(
        //             ExprKind::Allocation {
        //                 kind: AllocKind::Array(TypeKind::Any.into(), size_cv.into()),
        //                 elements: vec![],
        //                 default_elem: Some(default_elem),
        //                 region,
        //             },
        //             span,
        //         ))
        //     } else {
        //         let next_tok = tokens.pop_front().ok_or_else(ParseError::eof)?;
        //
        //         let mut alloc_kind = match next_tok.node {
        //             Token::RBracket => AllocKind::DynArray(TypeKind::Int.into()),
        //             Token::Int(i) => {
        //                 expect_next(tokens, Token::RBracket)?;
        //                 AllocKind::Array(TypeKind::Any.into(), ComptimeValue::Int(i as i128).into())
        //             }
        //             Token::Identifier(name) => {
        //                 expect_next(tokens, Token::RBracket)?;
        //                 AllocKind::Array(TypeKind::Any.into(), ComptimeValue::Ident(name).into())
        //             }
        //             _ => {
        //                 return Err(ParseError::new(
        //                     ParseErrorKind::ExpectedExpression,
        //                     next_tok.span,
        //                 ));
        //             }
        //         };
        //
        //         let Some(array_ty) = parse_type(ctx, tokens)? else {
        //             return Err(ParseError::new(
        //                 ParseErrorKind::ExpectedType,
        //                 peek_span(tokens).unwrap_or(start_span),
        //             ));
        //         };
        //
        //         alloc_kind.fill_type(array_ty.node);
        //
        //         expect_next(tokens, Token::LBrace)?;
        //
        //         let mut exprs = Vec::new();
        //
        //         loop {
        //             if let Some(Token::RBrace) = peek_token(tokens) {
        //                 break;
        //             }
        //
        //             exprs.push(parse_expr(ctx, tokens)?);
        //
        //             if let Some(Token::Comma) = peek_token(tokens) {
        //                 tokens.pop_front();
        //             }
        //         }
        //
        //         let end_span = expect_next(tokens, Token::RBrace)?;
        //         let span = start_span.merge(&end_span);
        //
        //         Ok(Spanned::new(
        //             ExprKind::Allocation {
        //                 kind: alloc_kind,
        //                 elements: exprs,
        //                 default_elem: None,
        //                 region: None,
        //             },
        //             span,
        //         ))
        //     }
        // }
        Token::LParen => {
            let start_span = token.span;
            let expr = parse_expr(ctx, tokens)?;

            if let Some(Token::Comma) = peek_token(tokens) {
                tokens.pop_front();
                let mut exprs = vec![expr];

                loop {
                    if let Some(Token::RParen) = peek_token(tokens) {
                        break;
                    }

                    exprs.push(parse_expr(ctx, tokens)?);

                    if let Some(Token::Comma) = peek_token(tokens) {
                        tokens.pop_front();
                    }
                }

                let end_span = expect_next(tokens, Token::RParen)?;
                let span = start_span.merge(&end_span);

                Ok(Spanned::new(
                    ExprKind::Allocation {
                        // filled in during typechecking
                        kind: AllocKind::Tuple(vec![]),
                        elements: exprs,
                        default_elem: None,
                        region: None,
                    },
                    span,
                ))
            } else {
                let _ = expect_next(tokens, Token::RParen)?;
                Ok(expr)
            }
        }

        Token::LBrace => {
            // parsing a record
            let start_span = token.span;
            let mut fields = Vec::new();
            let mut exprs = Vec::new();

            loop {
                if let Some(Token::RBrace) = peek_token(tokens) {
                    break;
                }

                let field_name = expect_identifier(tokens)?;
                let ty = parse_type_annot(ctx, tokens)?;
                expect_next(tokens, Token::Eq)?;
                let expr = parse_expr(ctx, tokens)?;

                fields.push(RecordField {
                    name: field_name.node,
                    ty,
                });
                exprs.push(expr);

                if let Some(Token::Comma) = peek_token(tokens) {
                    tokens.pop_front();
                }
            }

            let end_span = expect_next(tokens, Token::RBrace)?;
            let span = start_span.merge(&end_span);
            Ok(Spanned::new(
                ExprKind::Allocation {
                    kind: AllocKind::Record(fields),
                    elements: exprs,
                    default_elem: None,
                    region: None,
                },
                span,
            ))
        }

        _ => Err(ParseError::new(
            ParseErrorKind::ExpectedExpression,
            token.span,
        )),
    }
}

fn parse_params(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<(Params, Span)> {
    let start_span = expect_next(tokens, Token::LParen)?;
    let mut params = Params { params: Vec::new() };

    while !tokens.is_empty() {
        if let Some(Token::RParen) = peek_token(tokens) {
            break;
        }

        // Check for variadic parameter `...`
        if let Some(Token::DotDotDot) = peek_token(tokens) {
            tokens.pop_front();
            // Variadic params don't require a type annotation
            // They default to a generic variadic type
            let ty = Type::synthetic(TypeKind::Variadic);

            let pattern = Spanned::new(PatKind::Symbol(ctx.intern("__variadic_arg")), 0..0);
            params.params.push(Param {
                pattern,
                ty,
                is_comptime: false,
            });
            break;
        }

        let is_comptime = if let Some(Token::Keyword(Keyword::Comptime)) = peek_token(tokens) {
            tokens.pop_front();
            true
        } else {
            false
        };

        let name = expect_identifier(tokens)?;
        let ty = parse_type_annot(ctx, tokens)?
            .ok_or_else(|| ParseError::new(ParseErrorKind::ExpectedType, name.span.clone()))?;

        let pattern = Spanned::new(PatKind::Symbol(name.node), name.span);
        params.params.push(Param {
            pattern,
            ty,
            is_comptime,
        });

        if let Some(Token::Comma) = peek_token(tokens) {
            tokens.pop_front();
        } else {
            break;
        }
    }

    let end_span = expect_next(tokens, Token::RParen)?;
    Ok((params, start_span.merge(&end_span)))
}

fn parse_proc_sig(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Signature> {
    let (params, params_span) = parse_params(ctx, tokens)?;

    let (return_ty, end_span) = if let Some(Token::Colon) = peek_token(tokens) {
        let ty = parse_type_annot(ctx, tokens)?;
        let end = ty
            .as_ref()
            .map(|t| t.span.clone())
            .unwrap_or(params_span.clone());
        (ty, end)
    } else {
        (None, params_span.clone())
    };

    let span = params_span.merge(&end_span);
    Ok(Spanned::new(SignatureInner { params, return_ty }, span))
}

fn check_region_type_annotation(
    ctx: &mut Ctx,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Option<Region>> {
    let Some(Token::At) = peek_token(tokens) else {
        return Ok(None);
    };

    let _ = tokens.pop_front().unwrap().span;
    let region_ident = expect_identifier(tokens)?.node;

    let region = if ctx.resolve(region_ident) == "local" {
        Region::Scoped(0)
    } else {
        Region::Named(region_ident)
    };

    Ok(Some(region))
}

fn parse_type(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Option<Type>> {
    let ty = match peek_token(tokens) {
        Some(Token::Keyword(Keyword::Int)) => {
            let span = tokens.pop_front().unwrap().span;

            Spanned::new(TypeKind::Int, span)
        }
        Some(Token::Keyword(Keyword::Bool)) => {
            let span = tokens.pop_front().unwrap().span;

            Spanned::new(TypeKind::Bool, span)
        }
        Some(Token::Keyword(Keyword::Char)) => {
            let span = tokens.pop_front().unwrap().span;

            Spanned::new(TypeKind::Char, span)
        }
        Some(Token::Keyword(Keyword::Type)) => {
            let span = tokens.pop_front().unwrap().span;
            Spanned::new(TypeKind::Type, span)
        }
        Some(Token::Keyword(Keyword::Region)) => {
            let span = tokens.pop_front().unwrap().span;
            Spanned::new(TypeKind::Region, span)
        }
        Some(Token::Dot) => {
            // variant type starting with dot
            let dot_span = tokens.pop_front().unwrap().span;
            let Spanned {
                node: Token::Identifier(variant_name),
                span: ident_span,
            } = tokens.pop_front().unwrap()
            else {
                return Err(ParseError::new(
                    ParseErrorKind::ExpectedIdentifier,
                    dot_span,
                ));
            };
            let span = dot_span.merge(&ident_span);
            let variant_name = Spanned::new(variant_name, span);

            let mut fields = Vec::new();

            if let Some(Token::LParen) = peek_token(tokens) {
                tokens.pop_front();
                let mut adts = Vec::new();
                loop {
                    if let Some(Token::RParen) = peek_token(tokens) {
                        break;
                    }
                    adts.push(parse_type(ctx, tokens)?.unwrap());
                    if let Some(Token::Comma) = peek_token(tokens) {
                        tokens.pop_front();
                    }
                }
                expect_next(tokens, Token::RParen)?;
                fields.push(VariantField {
                    name: variant_name,
                    adts,
                });
            } else {
                fields.push(VariantField {
                    name: variant_name,
                    adts: Vec::new(),
                });
            }

            while let Some(Token::Bar) = peek_token(tokens) {
                tokens.pop_front();

                expect_next(tokens, Token::Dot)?;
                let Spanned {
                    node: Token::Identifier(variant_name),
                    span: ident_span,
                } = tokens.pop_front().unwrap()
                else {
                    return Err(ParseError::new(
                        ParseErrorKind::ExpectedIdentifier,
                        dot_span,
                    ));
                };
                let variant_span = dot_span.merge(&ident_span);
                let variant_name = Spanned::new(variant_name, variant_span);

                if let Some(Token::LParen) = peek_token(tokens) {
                    tokens.pop_front();
                    let mut adts = Vec::new();
                    loop {
                        if let Some(Token::RParen) = peek_token(tokens) {
                            break;
                        }
                        adts.push(parse_type(ctx, tokens)?.unwrap());
                        if let Some(Token::Comma) = peek_token(tokens) {
                            tokens.pop_front();
                        }
                    }
                    expect_next(tokens, Token::RParen)?;
                    fields.push(VariantField {
                        name: variant_name,
                        adts,
                    });
                } else {
                    fields.push(VariantField {
                        name: variant_name,
                        adts: Vec::new(),
                    });
                }
            }

            let first_span = fields.first().unwrap().name.span.clone();
            let span = fields.last().unwrap().name.span.merge(&first_span);
            Spanned::new(TypeKind::Variant(fields), span)
        }
        Some(Token::Identifier(_name)) => {
            // Type alias
            let Spanned {
                node: Token::Identifier(name),
                span,
            } = tokens.pop_front().unwrap()
            else {
                unreachable!()
            };
            Spanned::new(TypeKind::TypeAlias(name), span)
        }
        Some(Token::Caret) => {
            let caret_span = tokens.pop_front().unwrap().span;
            let ty = match parse_type(ctx, tokens)? {
                Some(ty) => ty,
                None => return Err(ParseError::new(ParseErrorKind::ExpectedType, caret_span)),
            };

            let region = check_region_type_annotation(ctx, tokens)?;

            let span = caret_span.merge(&ty.span);
            Spanned::new(TypeKind::Ptr(Box::new(ty.node), region), span)
        }
        Some(Token::LBrace) => {
            let fields = extract_typedef_fields(tokens)?;
            let record_fields = parse_record_fields(ctx, fields);

            let start_span = record_fields
                .first()
                .map(|f| f.ty.as_ref().unwrap().span.clone())
                .unwrap();

            let end_span = record_fields
                .last()
                .map(|f| f.ty.as_ref().unwrap().span.clone())
                .unwrap();

            let span = start_span.merge(&end_span);

            Spanned::new(TypeKind::Record(record_fields), span)
        }
        Some(Token::LBracket) => {
            let start_span = tokens.pop_front().unwrap().span;

            let arr_ty = match parse_type(ctx, tokens)? {
                Some(ty) => ty,
                None => {
                    return Err(ParseError::new(
                        ParseErrorKind::ExpectedType,
                        peek_span(tokens).unwrap_or(start_span),
                    ));
                }
            };

            expect_next(tokens, Token::SemiColon)?;

            let next_tok = tokens.pop_front().ok_or_else(ParseError::eof)?;
            let alloc_kind = match next_tok.node {
                Token::Int(i) => {
                    AllocKind::Array(arr_ty.node.into(), ComptimeValue::Int(i as i128).into())
                }
                Token::Identifier(name) => {
                    // assuming comptime identifier that will be resolved during typechecking
                    AllocKind::Array(arr_ty.node.into(), ComptimeValue::Ident(name).into())
                }
                _ => {
                    return Err(ParseError::new(
                        ParseErrorKind::ExpectedExpression,
                        next_tok.span,
                    ));
                }
            };

            expect_next(tokens, Token::RBracket)?;

            let span = start_span.merge(&arr_ty.span);

            let region = check_region_type_annotation(ctx, tokens)?;
            Spanned::new(
                TypeKind::Alloc(alloc_kind, region.unwrap_or(Region::Scoped(0))),
                span,
            )
        }
        _ => {
            return Ok(None);
        }
    };

    Ok(Some(ty))
}

fn parse_type_annot(
    ctx: &mut Ctx,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Option<Type>> {
    let Some(Spanned {
        node: Token::Colon,
        span: _,
    }) = tokens.pop_front()
    else {
        return Err(ParseError::new(ParseErrorKind::ExpectedType, 0..0));
    };

    parse_type(ctx, tokens)
}

fn parse_ifelse(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Stmt> {
    let if_span = expect_next(tokens, Token::Keyword(Keyword::If))?;
    let cond = parse_expr(ctx, tokens)?;
    let then = parse_block(ctx, tokens)?;

    let (else_, end_span) = if let Some(Token::Keyword(Keyword::Else)) = peek_token(tokens) {
        expect_next(tokens, Token::Keyword(Keyword::Else))?;
        let else_block = parse_block(ctx, tokens)?;
        let span = else_block.span.clone();
        (Some(else_block), span)
    } else {
        (None, then.span.clone())
    };

    let span = if_span.merge(&end_span);
    Ok(Spanned::new(
        StmtKind::IfElse(Box::new(IfElse { cond, then, else_ })),
        span,
    ))
}

fn parse_pat(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Pat> {
    let Some(token) = tokens.front() else {
        return Err(ParseError::eof());
    };
    let token_span = token.span.clone();
    let token_node = token.node.clone();

    match token_node {
        Token::Dot => {
            let dot_span = token_span.clone();
            tokens.pop_front();
            let Some(ident_token) = tokens.pop_front() else {
                return Err(ParseError::new(
                    ParseErrorKind::ExpectedIdentifier,
                    dot_span,
                ));
            };
            let Spanned {
                node: Token::Identifier(variant_name),
                span: ident_span,
            } = ident_token
            else {
                return Err(ParseError::new(
                    ParseErrorKind::ExpectedIdentifier,
                    dot_span,
                ));
            };
            let span = dot_span.merge(&ident_span);

            if let Some(Token::LParen) = peek_token(tokens) {
                tokens.pop_front();
                let mut bindings = Vec::new();

                loop {
                    if let Some(Token::RParen) = peek_token(tokens) {
                        break;
                    }

                    bindings.push(parse_pat(ctx, tokens)?);

                    if let Some(Token::Comma) = peek_token(tokens) {
                        tokens.pop_front();
                    } else {
                        break;
                    }
                }

                let end_span = expect_next(tokens, Token::RParen)?;
                let span = span.merge(&end_span);

                Ok(Spanned::new(
                    PatKind::Variant {
                        name: variant_name,
                        bindings,
                    },
                    span,
                ))
            } else {
                Ok(Spanned::new(
                    PatKind::Variant {
                        name: variant_name,
                        bindings: vec![],
                    },
                    span,
                ))
            }
        }
        Token::Identifier(name) => {
            let name_str = ctx.resolve(name);

            if name_str == "_" {
                let span = tokens.pop_front().unwrap().span;
                return Ok(Spanned::new(PatKind::Wildcard, span));
            }

            let ident = expect_identifier(tokens)?;
            Ok(Spanned::new(PatKind::Symbol(ident.node), ident.span))
        }
        // TOOD: clean this up to a parse literal function
        Token::Int(i) => {
            let span = tokens.pop_front().unwrap().span;
            Ok(Spanned::new(PatKind::Literal(ValueKind::Int(i)), span))
        }
        Token::Keyword(Keyword::True) => {
            let span = tokens.pop_front().unwrap().span;
            Ok(Spanned::new(PatKind::Literal(ValueKind::Bool(true)), span))
        }
        Token::Keyword(Keyword::False) => {
            let span = tokens.pop_front().unwrap().span;
            Ok(Spanned::new(PatKind::Literal(ValueKind::Bool(false)), span))
        }
        Token::LParen => {
            let start_span = tokens.pop_front().unwrap().span;
            let mut patterns = Vec::new();

            loop {
                if let Some(Token::RParen) = peek_token(tokens) {
                    break;
                }

                patterns.push(parse_pat(ctx, tokens)?);

                if let Some(Token::Comma) = peek_token(tokens) {
                    tokens.pop_front();
                } else {
                    break;
                }
            }

            let end_span = expect_next(tokens, Token::RParen)?;
            let span = start_span.merge(&end_span);

            Ok(Spanned::new(PatKind::Tuple(patterns), span))
        }
        // Token::LBrace => {
        //     let start_span = expect_next(tokens, Token::LBrace)?;
        //     let mut fields = Vec::new();
        //
        //     loop {
        //         if let Some(Token::RBrace) = peek_token(tokens) {
        //             break;
        //         }
        //
        //         let field_name = expect_identifier(tokens)?;
        //         let ty = parse_type_annot(ctx, tokens)?;
        //
        //         fields.push(RecordField {
        //             name: field_name.node,
        //             ty,
        //         });
        //
        //         if let Some(Token::Comma) = peek_token(tokens) {
        //             tokens.pop_front();
        //
        //             if let Some(Token::DotDot) = peek_token(tokens) {
        //                 tokens.pop_front();
        //                 break;
        //             }
        //         } else {
        //             break;
        //         }
        //     }
        //
        //     let end_span = expect_next(tokens, Token::RBrace)?;
        //
        //     let span = start_span.merge(&end_span);
        //     Ok(Spanned::new(PatKind::Record(fields), span))
        // }
        _ => Err(ParseError::new(
            ParseErrorKind::ExpectedPattern,
            token.span.clone(),
        )),
    }
}

fn parse_match_arm(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<MatchArm> {
    let pat = parse_pat(ctx, tokens)?;
    expect_next(tokens, Token::FatArrow)?;
    let body = parse_block(ctx, tokens)?;

    Ok(MatchArm { pat, body })
}

fn parse_match(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Stmt> {
    let match_span = expect_next(tokens, Token::Keyword(Keyword::Match))?;
    let scrutinee = parse_expr(ctx, tokens)?;
    expect_next(tokens, Token::Keyword(Keyword::With))?;

    let mut arms = Vec::new();

    loop {
        let arm = parse_match_arm(ctx, tokens)?;
        arms.push(arm);

        if let Some(Token::Bar) = peek_token(tokens) {
            tokens.pop_front();
        } else {
            break;
        }
    }

    let span = match_span.merge(&arms.last().unwrap().body.span);
    Ok(Spanned::new(
        StmtKind::Match(Match { scrutinee, arms }),
        span,
    ))
}

fn parse_for(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Stmt> {
    let for_span = expect_next(tokens, Token::Keyword(Keyword::For))?;
    let in_pat = parse_pat(ctx, tokens)?;
    expect_next(tokens, Token::Keyword(Keyword::In))?;
    let in_expr = parse_expr(ctx, tokens)?;
    let in_block = parse_block(ctx, tokens)?;

    let span = for_span.merge(&in_block.span);
    Ok(Spanned::new(
        StmtKind::For {
            binding: in_pat,
            iter: in_expr,
            body: in_block,
        },
        span,
    ))
}

fn parse_stmt(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Stmt> {
    match peek_token(tokens) {
        Some(Token::Keyword(Keyword::For)) => parse_for(ctx, tokens),
        Some(Token::Keyword(Keyword::If)) => parse_ifelse(ctx, tokens),
        Some(Token::Keyword(Keyword::Match)) => parse_match(ctx, tokens),
        Some(Token::Keyword(Keyword::Return)) => {
            let ret_span = expect_next(tokens, Token::Keyword(Keyword::Return))?;
            // TODO: is this the right way to scan for empty return statements?
            if let Some(Token::RBrace) = peek_token(tokens) {
                return Ok(Spanned::new(StmtKind::Return(None), ret_span));
            }

            let expr = parse_expr(ctx, tokens)?;
            let span = ret_span.merge(&expr.span);
            Ok(Spanned::new(StmtKind::Return(Some(expr)), span))
        }
        _ => {
            // let name = expect_identifier(tokens)?;
            let left_expr = parse_expr(ctx, tokens)?;
            let start_span = left_expr.span.clone();

            if let ExprKind::Call(call) = &left_expr.node {
                let span = left_expr.span.clone();
                return Ok(Spanned::new(StmtKind::Call(call.clone()), span));
            }

            match peek_token(tokens) {
                Some(Token::Colon) => {
                    let ExprKind::Value(ValueKind::Ident(name)) = left_expr.node else {
                        return Err(ParseError::new(
                            ParseErrorKind::ExpectedIdentifier,
                            left_expr.span,
                        ));
                    };

                    let ty = parse_type_annot(ctx, tokens)?;
                    expect_next(tokens, Token::Eq)?;
                    let expr = parse_expr(ctx, tokens)?;
                    let span = start_span.merge(&expr.span);
                    Ok(Spanned::new(
                        StmtKind::ValDec {
                            name: Spanned::new(name, left_expr.span),
                            ty,
                            expr,
                            is_comptime: false,
                        },
                        span,
                    ))
                }
                Some(Token::Eq) => {
                    expect_next(tokens, Token::Eq)?;
                    let expr = parse_expr(ctx, tokens)?;
                    let span = start_span.merge(&expr.span);
                    Ok(Spanned::new(
                        StmtKind::Assign {
                            location: left_expr,
                            expr,
                        },
                        span,
                    ))
                }
                // Some(Token::LParen) => {
                //     let ExprKind::Value(ValueKind::Ident(name)) = left_expr.node else {
                //         return Err(ParseError::new(
                //             ParseErrorKind::ExpectedIdentifier,
                //             left_expr.span,
                //         ));
                //     };
                //     let name = Spanned::new(name, left_expr.span);
                //
                //     let expr = parse_identifier_expr(ctx, name, tokens)?;
                //     let ExprKind::Call(call) = expr.node else {
                //         unreachable!()
                //     };
                //     let span = expr.span;
                //     Ok(Spanned::new(StmtKind::Call(call), span))
                // }
                t => {
                    println!("here {t:?} {left_expr:?}");
                    Err(ParseError::new(
                        ParseErrorKind::ExpectedExpression,
                        peek_span(tokens).unwrap_or(start_span),
                    ))
                }
            }
        }
    }
}

fn parse_block(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Block> {
    let start_span = expect_next(tokens, Token::LBrace)?;
    let mut stmts = Vec::new();

    while !tokens.is_empty() {
        if let Some(Token::RBrace) = peek_token(tokens) {
            break;
        }

        stmts.push(parse_stmt(ctx, tokens)?);
    }

    let end_span = expect_next(tokens, Token::RBrace)?;
    let span = start_span.merge(&end_span);

    let mut expr = None;
    if !stmts.is_empty() {
        let last_stmt = &stmts[stmts.len() - 1];
        if let StmtKind::Return(expr_) = &last_stmt.node {
            expr = expr_.clone();
        }
    }

    Ok(Spanned::new(BlockInner { stmts, expr }, span))
}

// fn parse_constraints(
//     ctx: &mut Ctx,
//     tokens: &mut VecDeque<SpannedToken>,
// ) -> ParseResult<Vec<Constraint>> {
//     let mut constraints = Vec::new();
//
//     if let Some(Token::Keyword(Keyword::With)) = peek_token(tokens) {
//         let _ = tokens.pop_front().unwrap();
//
//         while !tokens.is_empty() {
//             match peek_token(tokens) {
//                 Some(Token::Keyword(Keyword::Allocates)) => {
//                     let t = tokens.pop_front().unwrap();
//                     let Some(region) = check_region_type_annotation(ctx, tokens)? else {
//                         return Err(ParseError::new(ParseErrorKind::ExpectedRegion, t.span));
//                     };
//                     constraints.push(Constraint::Allocates(region));
//                 }
//                 _ => break,
//             }
//         }
//     }
//
//     Ok(constraints)
// }

fn parse_procedure(
    ctx: &mut Ctx,
    name: Spanned<Symbol>,
    fn_ty: Option<Type>,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Decl> {
    let sig = parse_proc_sig(ctx, tokens)?;
    // let constraints = parse_constraints(ctx, tokens)?;
    let block = parse_block(ctx, tokens)?;
    let span = name.span.merge(&block.span);

    Ok(Spanned::new(
        DeclKind::Procedure {
            name,
            fn_ty,
            sig,
            // constraints,
            block,
            monomorph_of: None,
            is_comptime: false,
        },
        span,
    ))
}

fn extract_typedef_fields(
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Vec<VecDeque<SpannedToken>>> {
    let mut fields = Vec::new();
    expect_next(tokens, Token::LBrace)?;

    while !tokens.is_empty() {
        let mut field = VecDeque::new();
        let mut in_parens = false;

        loop {
            if let Some(Token::Comma) = peek_token(tokens) {
                tokens.pop_front();
                if !in_parens {
                    break;
                } else {
                    // Keep the comma token with a dummy span for now
                    field.push_back(Spanned::new(Token::Comma, 0..0));
                }
            } else if let Some(Token::RBrace) = peek_token(tokens) {
                break;
            } else {
                let token = tokens.pop_front().ok_or_else(ParseError::eof)?;
                if token.node == Token::LParen {
                    in_parens = true;
                } else if token.node == Token::RParen {
                    in_parens = false;
                }
                field.push_back(token);
            }
        }

        fields.push(field);

        if let Some(Token::RBrace) = peek_token(tokens) {
            break;
        }
    }

    expect_next(tokens, Token::RBrace)?;
    Ok(fields)
}

#[inline]
fn is_record(fields: &[VecDeque<SpannedToken>]) -> bool {
    fields
        .iter()
        .all(|f| f.iter().any(|t| t.node == Token::Colon))
}

fn parse_record_fields(ctx: &mut Ctx, fields: Vec<VecDeque<SpannedToken>>) -> Vec<RecordField> {
    fields
        .into_iter()
        .map(|mut field_tokens| {
            let name = expect_identifier(&mut field_tokens).unwrap();
            let ty = parse_type_annot(ctx, &mut field_tokens)
                .expect("to parse")
                .expect("expected type annot on struct field");

            RecordField {
                name: name.node,
                ty: Some(ty),
            }
        })
        .collect()
}

fn parse_enum_fields(fields: Vec<VecDeque<SpannedToken>>) -> Vec<VariantField> {
    fields
        .into_iter()
        .map(|mut field_tokens| {
            let name = expect_identifier(&mut field_tokens).unwrap();

            let mut adts = Vec::new();
            if let Some(Token::LParen) = peek_token(&field_tokens) {
                field_tokens.pop_front();
                loop {
                    if let Some(Token::RParen) = peek_token(&field_tokens) {
                        break;
                    }

                    if let Some(Token::Keyword(Keyword::Int)) = peek_token(&field_tokens) {
                        let span = field_tokens.pop_front().unwrap().span;
                        adts.push(Spanned::new(TypeKind::Int, span));
                    } else if let Some(Token::Identifier(_)) = peek_token(&field_tokens) {
                        let Spanned {
                            node: Token::Identifier(type_name),
                            span,
                        } = field_tokens.pop_front().unwrap()
                        else {
                            unreachable!()
                        };

                        adts.push(Spanned::new(TypeKind::TypeAlias(type_name), span));
                    } else {
                        panic!("expected type name or int");
                    }

                    if let Some(Token::Comma) = peek_token(&field_tokens) {
                        field_tokens.pop_front();
                    }
                }

                field_tokens.pop_front();
            }

            VariantField { name, adts }
        })
        .collect()
}

fn parse_typedef(
    ctx: &mut Ctx,
    name: Spanned<Symbol>,
    _ty: Option<Type>,
    tokens: &mut VecDeque<SpannedToken>,
) -> ParseResult<Decl> {
    let start_span = name.span.clone();
    let fields = extract_typedef_fields(tokens)?;

    if is_record(&fields) {
        let record_fields = parse_record_fields(ctx, fields);

        let end_span = record_fields
            .last()
            .map(|f| f.ty.as_ref().unwrap().span.clone())
            .unwrap_or(start_span.clone());
        let span = start_span.merge(&end_span);

        Ok(Spanned::new(
            DeclKind::TypeDef {
                name,
                def: TypeAliasDefinition::Record(record_fields),
            },
            span,
        ))
    } else {
        let enum_fields = parse_enum_fields(fields);
        let end_span = enum_fields
            .last()
            .map(|f| f.name.span.clone())
            .unwrap_or(start_span.clone());
        let span = start_span.merge(&end_span);

        Ok(Spanned::new(
            DeclKind::TypeDef {
                name,
                def: TypeAliasDefinition::Variant(enum_fields),
            },
            span,
        ))
    }
}

fn parse_decls(ctx: &mut Ctx, tokens: &mut VecDeque<SpannedToken>) -> ParseResult<Vec<Decl>> {
    let mut decs = Vec::new();

    while !tokens.is_empty() {
        let name = expect_identifier(tokens)?;
        let start_span = name.span.clone();

        let ty = parse_type_annot(ctx, tokens)?;

        let Some(Spanned {
            node: Token::Colon, ..
        }) = tokens.pop_front()
        else {
            return Err(ParseError::new(
                ParseErrorKind::ExpectedExpression,
                peek_span(tokens).unwrap_or(start_span),
            ));
        };

        match peek_token(tokens) {
            Some(Token::LParen) => {
                decs.push(parse_procedure(ctx, name, ty, tokens)?);
            }
            Some(Token::LBrace) => {
                decs.push(parse_typedef(ctx, name, ty, tokens)?);
            }
            Some(Token::Keyword(Keyword::Extern)) => {
                expect_next(tokens, Token::Keyword(Keyword::Extern))?;
                let sig = parse_proc_sig(ctx, tokens)?;
                let span = name.span.merge(&sig.span);
                decs.push(Spanned::new(
                    DeclKind::Extern {
                        name,
                        sig,
                        generic_params: None,
                    },
                    span,
                ));
            }
            Some(_) => {
                let expr = parse_expr(ctx, tokens)?;
                let span = start_span.merge(&expr.span);
                decs.push(Spanned::new(
                    DeclKind::Constant {
                        name,
                        ty,
                        expr,
                        is_comptime: false,
                    },
                    span,
                ));
            }
            None => return Err(ParseError::eof()),
        }
    }

    Ok(decs)
}

pub fn parse(ctx: &mut Ctx, tokens: VecDeque<SpannedToken>) -> ParseResult<Module> {
    let mut tokens = tokens;

    let decls = parse_decls(ctx, &mut tokens)?;

    let module = Module {
        declarations: decls,
    };

    Ok(module)
}