const std = @import("std");
const Token = @import("token.zig").Token;
const Tag = @import("token.zig").Tag;
const Lexer = @import("lexer.zig").Lexer;
const ast = @import("ast.zig");
const Node = ast.Node;
const Type = @import("types.zig").Type;
const errors = @import("errors.zig");

pub const Parser = struct {
    lexer: Lexer,
    current: Token,
    source: [:0]const u8,
    allocator: std.mem.Allocator,
    err_msg: ?[]const u8,
    err_offset: ?u32 = null,
    prev_end: u32 = 0,
    diagnostics: ?*errors.DiagnosticList = null,

    pub fn init(allocator: std.mem.Allocator, source: [:0]const u8) Parser {
        var lexer = Lexer.init(source);
        const first = lexer.next();
        return .{
            .lexer = lexer,
            .current = first,
            .source = source,
            .allocator = allocator,
            .err_msg = null,
            .err_offset = null,
        };
    }

    fn createNode(self: *Parser, start: u32, data: Node.Data) !*Node {
        const node = try self.allocator.create(Node);
        node.* = .{ .span = .{ .start = start, .end = self.prev_end }, .data = data };
        return node;
    }

    pub fn parse(self: *Parser) anyerror!*Node {
        var decls = std.ArrayList(*Node).empty;
        while (self.current.tag != .eof) {
            const decl = try self.parseTopLevel();
            try decls.append(self.allocator, decl);
        }
        const node = try self.createNode(0, .{ .root = .{ .decls = try decls.toOwnedSlice(self.allocator) } });
        return node;
    }

    fn parseTopLevel(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;

        // Top-level flat import: #import "path";
        if (self.current.tag == .hash_import) {
            self.advance();
            if (self.current.tag != .string_literal) {
                return self.fail("expected string path after '#import'");
            }
            const raw = self.tokenSlice(self.current);
            const path = raw[1 .. raw.len - 1];
            self.advance();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .import_decl = .{ .path = path, .name = null } });
        }

        // Top-level #library directive: #library "libname";
        if (self.current.tag == .hash_library) {
            self.advance();
            if (self.current.tag != .string_literal) {
                return self.fail("expected string after '#library'");
            }
            const raw = self.tokenSlice(self.current);
            const lib_name = raw[1 .. raw.len - 1];
            self.advance();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .library_decl = .{ .lib_name = lib_name } });
        }

        // Top-level #run directive
        if (self.current.tag == .hash_run) {
            self.advance();
            const expr = try self.parseExpr();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .comptime_expr = .{ .expr = expr } });
        }

        // All top-level declarations start with an identifier
        if (self.current.tag != .identifier) {
            return self.fail("expected identifier at top level");
        }
        const name = self.tokenSlice(self.current);
        self.advance();

        // IDENT :: ...
        if (self.current.tag == .colon_colon) {
            self.advance();
            return self.parseConstBinding(name, start);
        }

        // IDENT : type : value; (typed constant)
        // IDENT : type = value; (typed variable)
        if (self.current.tag == .colon) {
            self.advance();
            return self.parseTypedBinding(name, start);
        }

        // IDENT := value; (variable)
        if (self.current.tag == .colon_equal) {
            self.advance();
            const value = try self.parseExpr();
            try self.expectSemicolonAfter(value);
            return try self.createNode(start, .{ .var_decl = .{ .name = name, .type_annotation = null, .value = value } });
        }

        return self.fail("expected '::', ':=', or ':' after identifier");
    }

    fn parseConstBinding(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        // After `::`
        // Could be: #run expr, enum { ... }, (params) -> type { body }, or expr;

        // Namespaced import: name :: #import "path";
        if (self.current.tag == .hash_import) {
            self.advance();
            if (self.current.tag != .string_literal) {
                return self.fail("expected string path after '#import'");
            }
            const raw = self.tokenSlice(self.current);
            const path = raw[1 .. raw.len - 1];
            self.advance();
            try self.expect(.semicolon);
            return try self.createNode(start_pos, .{ .import_decl = .{ .path = path, .name = name } });
        }

        // Compile-time evaluation: name :: #run expr;
        if (self.current.tag == .hash_run) {
            const run_start = self.current.loc.start;
            self.advance();
            const inner = try self.parseExpr();
            try self.expect(.semicolon);
            const ct = try self.createNode(run_start, .{ .comptime_expr = .{ .expr = inner } });
            return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = null, .value = ct } });
        }

        // Built-in declaration: name :: #builtin;
        if (self.current.tag == .hash_builtin) {
            const bi_start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            const bi = try self.createNode(bi_start, .{ .builtin_expr = {} });
            return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = null, .value = bi } });
        }

        // Enum declaration
        if (self.current.tag == .kw_enum) {
            return self.parseEnumDecl(name, start_pos);
        }

        // Struct declaration
        if (self.current.tag == .kw_struct) {
            return self.parseStructDecl(name, start_pos);
        }

        // C-style union declaration
        if (self.current.tag == .kw_union) {
            return self.parseUnionDecl(name, start_pos);
        }

        // Function declaration: (params) -> type { body } or () { body }
        if (self.current.tag == .l_paren) {
            // Look ahead: is this a function or an expression starting with `(`?
            // Heuristic: if after matching parens we see `{` or `->`, it's a function.
            if (self.isFunctionDef()) {
                return self.parseFnDecl(name, start_pos);
            }
        }

        // Bare block shorthand: name :: { body } is equivalent to name :: () { body }
        if (self.current.tag == .l_brace) {
            const body = try self.parseBlock();
            return try self.createNode(start_pos, .{ .fn_decl = .{ .name = name, .params = &.{}, .return_type = null, .body = body } });
        }

        // Otherwise it's a constant expression
        const value = try self.parseExpr();

        // name :: type_expr #builtin; — builtin with type annotation
        if (self.current.tag == .hash_builtin) {
            const bi_start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            const bi = try self.createNode(bi_start, .{ .builtin_expr = {} });
            return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = value, .value = bi } });
        }

        // name :: type_expr #foreign; — foreign with type annotation
        if (self.current.tag == .hash_foreign) {
            const fi_start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            const fi = try self.createNode(fi_start, .{ .foreign_expr = {} });
            return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = value, .value = fi } });
        }

        try self.expect(.semicolon);
        return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = null, .value = value } });
    }

    fn parseTypedBinding(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        // After `name :`
        // Parse type
        const type_node = try self.parseTypeExpr();

        if (self.current.tag == .colon) {
            // name : type : value; (typed constant)
            self.advance();
            const value = try self.parseExpr();
            try self.expectSemicolonAfter(value);
            return try self.createNode(start_pos, .{ .const_decl = .{ .name = name, .type_annotation = type_node, .value = value } });
        }

        if (self.current.tag == .equal) {
            // name : type = value; (typed variable)
            self.advance();
            const value = try self.parseExpr();
            try self.expectSemicolonAfter(value);
            return try self.createNode(start_pos, .{ .var_decl = .{ .name = name, .type_annotation = type_node, .value = value } });
        }

        if (self.current.tag == .semicolon) {
            // name : type; (default-initialized variable)
            self.advance();
            return try self.createNode(start_pos, .{ .var_decl = .{ .name = name, .type_annotation = type_node, .value = null } });
        }

        return self.fail("expected ':', '=' or ';' after type annotation");
    }

    fn parseTypeExpr(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;

        // Pointer type: *T
        if (self.current.tag == .star) {
            self.advance(); // skip '*'
            const pointee_type = try self.parseTypeExpr();
            return try self.createNode(start, .{ .pointer_type_expr = .{ .pointee_type = pointee_type } });
        }

        // Array type: [N]T, Slice type: []T, Many-pointer type: [*]T, Sentinel slice: [:0]T
        if (self.current.tag == .l_bracket) {
            self.advance(); // skip '['
            if (self.current.tag == .colon) {
                // Sentinel-terminated slice: [:0]T
                self.advance(); // skip ':'
                if (self.current.tag != .int_literal) {
                    return self.fail("expected sentinel value after ':'");
                }
                const sentinel_str = self.tokenSlice(self.current);
                self.advance(); // skip sentinel value
                try self.expect(.r_bracket); // expect ']'
                const elem_type = try self.parseTypeExpr();
                // Build name like "[:0]u8" for type resolution
                const elem_name = if (elem_type.data == .type_expr) elem_type.data.type_expr.name else "?";
                const name = try std.fmt.allocPrint(self.allocator, "[:{s}]{s}", .{ sentinel_str, elem_name });
                return try self.createNode(start, .{ .type_expr = .{ .name = name } });
            }
            if (self.current.tag == .r_bracket) {
                // Slice type: []T
                self.advance(); // skip ']'
                const elem_type = try self.parseTypeExpr();
                return try self.createNode(start, .{ .slice_type_expr = .{ .element_type = elem_type } });
            }
            if (self.current.tag == .star) {
                // Many-pointer type: [*]T
                self.advance(); // skip '*'
                try self.expect(.r_bracket); // expect ']'
                const elem_type = try self.parseTypeExpr();
                return try self.createNode(start, .{ .many_pointer_type_expr = .{ .element_type = elem_type } });
            }
            const len_node = try self.parseExpr();
            try self.expect(.r_bracket);
            const elem_type = try self.parseTypeExpr();
            return try self.createNode(start, .{ .array_type_expr = .{ .length = len_node, .element_type = elem_type } });
        }

        // Generic type parameter introduction: $T
        if (self.current.tag == .dollar) {
            self.advance();
            if (self.current.tag != .identifier) {
                return self.fail("expected type parameter name after '$'");
            }
            const name = self.tokenSlice(self.current);
            self.advance();
            return try self.createNode(start, .{ .type_expr = .{ .name = name, .is_generic = true } });
        }
        // Function pointer type: (ParamTypes) -> ReturnType
        if (self.current.tag == .l_paren) {
            self.advance(); // skip '('
            var param_types = std.ArrayList(*Node).empty;
            while (self.current.tag != .r_paren and self.current.tag != .eof) {
                if (param_types.items.len > 0) {
                    try self.expect(.comma);
                }
                try param_types.append(self.allocator, try self.parseTypeExpr());
            }
            try self.expect(.r_paren);
            var return_type: ?*Node = null;
            if (self.current.tag == .arrow) {
                self.advance(); // skip '->'
                return_type = try self.parseTypeExpr();
            }
            return try self.createNode(start, .{ .function_type_expr = .{
                .param_types = try param_types.toOwnedSlice(self.allocator),
                .return_type = return_type,
            } });
        }

        if (self.current.tag.isTypeKeyword() or self.current.tag == .identifier) {
            var name = self.tokenSlice(self.current);
            self.advance();

            // Qualified name: ns.Type or ns.Type(args)
            while (self.current.tag == .dot) {
                const dot_lexer = self.lexer;
                const dot_current = self.current;
                const dot_prev_end = self.prev_end;
                self.advance();
                if (self.current.tag == .identifier or self.current.tag.isTypeKeyword()) {
                    name = try std.fmt.allocPrint(self.allocator, "{s}.{s}", .{ name, self.tokenSlice(self.current) });
                    self.advance();
                } else {
                    // Not a qualified name continuation — restore the dot
                    self.lexer = dot_lexer;
                    self.current = dot_current;
                    self.prev_end = dot_prev_end;
                    break;
                }
            }

            // Parameterized type: Vector(N, T) or later generic struct instantiation
            if (self.current.tag == .l_paren) {
                self.advance(); // skip '('
                var args = std.ArrayList(*Node).empty;
                while (self.current.tag != .r_paren and self.current.tag != .eof) {
                    if (args.items.len > 0) {
                        try self.expect(.comma);
                    }
                    // Args can be int literals (for lengths) or type expressions
                    if (self.current.tag == .int_literal) {
                        const arg_start = self.current.loc.start;
                        const text = self.tokenSlice(self.current);
                        const value = std.fmt.parseInt(i64, text, 10) catch 0;
                        self.advance();
                        try args.append(self.allocator, try self.createNode(arg_start, .{ .int_literal = .{ .value = value } }));
                    } else {
                        try args.append(self.allocator, try self.parseTypeExpr());
                    }
                }
                try self.expect(.r_paren);
                return try self.createNode(start, .{ .parameterized_type_expr = .{
                    .name = name,
                    .args = try args.toOwnedSlice(self.allocator),
                } });
            }

            return try self.createNode(start, .{ .type_expr = .{ .name = name } });
        }
        // Inline struct type in type position: struct { ... }
        if (self.current.tag == .kw_struct) {
            return try self.parseStructDecl("__anon", start);
        }
        // Inline C-style union in type position: union { ... }
        if (self.current.tag == .kw_union) {
            return try self.parseUnionDecl("__anon", start);
        }
        // Inline enum type in type position: enum { ... }
        if (self.current.tag == .kw_enum) {
            return try self.parseEnumDecl("__anon", start);
        }
        return self.fail("expected type name");
    }

    fn parseEnumDecl(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        self.advance(); // skip 'enum'

        // Check for 'flags' modifier: enum flags { ... }
        var is_flags = false;
        if (self.current.tag == .identifier and std.mem.eql(u8, self.tokenSlice(self.current), "flags")) {
            is_flags = true;
            self.advance();
        }

        // Check for optional backing type: enum u8 { ... } or enum flags u32 { ... }
        var backing_type: ?*Node = null;
        if (self.current.tag != .l_brace) {
            backing_type = try self.parseTypeExpr();
        }

        try self.expect(.l_brace);
        var variant_names = std.ArrayList([]const u8).empty;
        var variant_types = std.ArrayList(?*Node).empty;
        var variant_values = std.ArrayList(?*Node).empty;
        var has_any_type = false;
        var has_any_value = false;
        while (self.current.tag != .r_brace and self.current.tag != .eof) {
            if (self.current.tag != .identifier) {
                return self.fail("expected variant name");
            }
            try variant_names.append(self.allocator, self.tokenSlice(self.current));
            self.advance();
            if (self.current.tag == .colon_colon) {
                // Explicit value: name :: expr;  or  name :: expr: type;
                self.advance();
                const val_expr = try self.parseExpr();
                try variant_values.append(self.allocator, val_expr);
                has_any_value = true;
                // Check for payload type after value: name :: 0x300: KeyData
                if (self.current.tag == .colon) {
                    if (is_flags) {
                        return self.fail("flags enum variants cannot have payloads");
                    }
                    self.advance();
                    const vtype = try self.parseTypeExpr();
                    try variant_types.append(self.allocator, vtype);
                    has_any_type = true;
                } else {
                    try variant_types.append(self.allocator, null);
                }
            } else if (self.current.tag == .colon) {
                // Typed variant: name: type;
                if (is_flags) {
                    return self.fail("flags enum variants cannot have payloads");
                }
                self.advance();
                const vtype = try self.parseTypeExpr();
                try variant_types.append(self.allocator, vtype);
                try variant_values.append(self.allocator, null);
                has_any_type = true;
            } else {
                // Void variant: name;
                try variant_types.append(self.allocator, null);
                try variant_values.append(self.allocator, null);
            }
            if (self.current.tag == .semicolon) {
                self.advance();
            }
        }
        try self.expect(.r_brace);
        // Always produce enum_decl; variant_types distinguishes payload-less from tagged
        return try self.createNode(start_pos, .{ .enum_decl = .{
            .name = name,
            .variant_names = try variant_names.toOwnedSlice(self.allocator),
            .variant_types = if (has_any_type) try variant_types.toOwnedSlice(self.allocator) else &.{},
            .is_flags = is_flags,
            .variant_values = if (has_any_value) try variant_values.toOwnedSlice(self.allocator) else &.{},
            .backing_type = backing_type,
        } });
    }

    fn parseUnionDecl(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        self.advance(); // skip 'union'
        try self.expect(.l_brace);
        var field_names = std.ArrayList([]const u8).empty;
        var field_types = std.ArrayList(*Node).empty;
        var anon_idx: u32 = 0;
        while (self.current.tag != .r_brace and self.current.tag != .eof) {
            // Anonymous struct field: struct { x, y: f32; };
            if (self.current.tag == .kw_struct) {
                const anon_field = try std.fmt.allocPrint(self.allocator, "__anon_{d}", .{anon_idx});
                anon_idx += 1;
                const anon_struct_name = try std.fmt.allocPrint(self.allocator, "{s}.{s}", .{ name, anon_field });
                const struct_node = try self.parseStructDecl(anon_struct_name, self.current.loc.start);
                try field_names.append(self.allocator, anon_field);
                try field_types.append(self.allocator, struct_node);
                if (self.current.tag == .semicolon) {
                    self.advance();
                }
                continue;
            }
            if (self.current.tag != .identifier) {
                return self.fail("expected field name or 'struct'");
            }
            try field_names.append(self.allocator, self.tokenSlice(self.current));
            self.advance();
            if (self.current.tag != .colon) {
                return self.fail("union fields must have a type");
            }
            self.advance();
            const ftype = try self.parseTypeExpr();
            try field_types.append(self.allocator, ftype);
            if (self.current.tag == .semicolon) {
                self.advance();
            }
        }
        try self.expect(.r_brace);
        return try self.createNode(start_pos, .{ .union_decl = .{
            .name = name,
            .field_names = try field_names.toOwnedSlice(self.allocator),
            .field_types = try field_types.toOwnedSlice(self.allocator),
        } });
    }

    fn parseStructDecl(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        self.advance(); // skip 'struct'

        // Optional type params: struct($N: u32, $T: Type) { ... }
        var type_params = std.ArrayList(ast.StructTypeParam).empty;
        if (self.current.tag == .l_paren) {
            self.advance(); // skip '('
            while (self.current.tag != .r_paren and self.current.tag != .eof) {
                if (type_params.items.len > 0) {
                    try self.expect(.comma);
                }
                // Expect $name : constraint
                try self.expect(.dollar);
                if (self.current.tag != .identifier) {
                    return self.fail("expected type parameter name after '$'");
                }
                const param_name = self.tokenSlice(self.current);
                self.advance();
                try self.expect(.colon);
                const constraint = try self.parseTypeExpr();
                try type_params.append(self.allocator, .{ .name = param_name, .constraint = constraint });
            }
            try self.expect(.r_paren);
        }

        try self.expect(.l_brace);

        var field_names = std.ArrayList([]const u8).empty;
        var field_types = std.ArrayList(*Node).empty;
        var field_defaults = std.ArrayList(?*Node).empty;

        while (self.current.tag != .r_brace and self.current.tag != .eof) {
            // Parse field group: name1, name2, ...: type (= default)?;
            var group_names = std.ArrayList([]const u8).empty;

            if (self.current.tag != .identifier) {
                return self.fail("expected field name in struct");
            }
            try group_names.append(self.allocator, self.tokenSlice(self.current));
            self.advance();

            while (self.current.tag == .comma) {
                self.advance(); // skip ','
                if (self.current.tag != .identifier) {
                    return self.fail("expected field name after ','");
                }
                try group_names.append(self.allocator, self.tokenSlice(self.current));
                self.advance();
            }

            try self.expect(.colon);
            const field_type = try self.parseTypeExpr();

            // Check for default value: = expr
            var default_val: ?*Node = null;
            if (self.current.tag == .equal) {
                self.advance();
                default_val = try self.parseExpr();
            }

            // All names in the group share the same type and default
            for (group_names.items) |fname| {
                // `_` is an ignore identifier — auto-rename to unique internal name
                const actual_name = if (std.mem.eql(u8, fname, "_"))
                    try std.fmt.allocPrint(self.allocator, "_{d}", .{field_names.items.len})
                else
                    fname;
                try field_names.append(self.allocator, actual_name);
                try field_types.append(self.allocator, field_type);
                try field_defaults.append(self.allocator, default_val);
            }

            if (self.current.tag == .semicolon) {
                self.advance();
            }
        }
        try self.expect(.r_brace);

        return try self.createNode(start_pos, .{ .struct_decl = .{
            .name = name,
            .field_names = try field_names.toOwnedSlice(self.allocator),
            .field_types = try field_types.toOwnedSlice(self.allocator),
            .field_defaults = try field_defaults.toOwnedSlice(self.allocator),
            .type_params = try type_params.toOwnedSlice(self.allocator),
        } });
    }

    fn parseStructLiteral(self: *Parser, struct_name: ?[]const u8, type_expr: ?*Node, start_pos: u32) anyerror!*Node {
        try self.expect(.l_brace);

        var field_inits = std.ArrayList(ast.StructFieldInit).empty;

        while (self.current.tag != .r_brace and self.current.tag != .eof) {
            if (field_inits.items.len > 0) {
                try self.expect(.comma);
            }

            // Check if this is a named field: identifier followed by '='
            if (self.current.tag == .identifier) {
                const saved_lexer = self.lexer;
                const saved_current = self.current;
                const saved_prev_end = self.prev_end;
                const fname = self.tokenSlice(self.current);
                const ident_start = self.current.loc.start;
                self.advance();

                if (self.current.tag == .equal) {
                    // Named field: name = expr
                    self.advance(); // skip '='
                    const value = try self.parseExpr();
                    try field_inits.append(self.allocator, .{ .name = fname, .value = value });
                    continue;
                } else if (self.current.tag == .comma or self.current.tag == .r_brace) {
                    // Shorthand: just an identifier (name = identifier with same name)
                    const ident_node = try self.createNode(ident_start, .{ .identifier = .{ .name = fname } });
                    try field_inits.append(self.allocator, .{ .name = fname, .value = ident_node });
                    continue;
                }

                // Not named — backtrack and parse as positional expression
                self.lexer = saved_lexer;
                self.current = saved_current;
                self.prev_end = saved_prev_end;
            }

            // Positional field: just an expression
            const value = try self.parseExpr();
            try field_inits.append(self.allocator, .{ .name = null, .value = value });
        }
        try self.expect(.r_brace);

        return try self.createNode(start_pos, .{ .struct_literal = .{
            .struct_name = struct_name,
            .type_expr = type_expr,
            .field_inits = try field_inits.toOwnedSlice(self.allocator),
        } });
    }

    fn reconstructQualifiedName(self: *Parser, node: *Node) ![]const u8 {
        if (node.data == .identifier) return node.data.identifier.name;
        if (node.data == .field_access) {
            const obj_name = try self.reconstructQualifiedName(node.data.field_access.object);
            return std.fmt.allocPrint(self.allocator, "{s}.{s}", .{ obj_name, node.data.field_access.field });
        }
        return error.ParseError;
    }

    /// Parse a parenthesized parameter list: `(name: type, $T: Type, args: ..Any)`
    /// Handles `$` generic params, `..` variadic marker, and comptime detection.
    /// Expects opening `(` already NOT consumed — this function consumes `(` through `)`.
    fn parseParams(self: *Parser) anyerror![]const ast.Param {
        try self.expect(.l_paren);
        var params = std.ArrayList(ast.Param).empty;
        while (self.current.tag != .r_paren and self.current.tag != .eof) {
            if (params.items.len > 0) {
                try self.expect(.comma);
            }
            var is_ct_param = false;
            if (self.current.tag == .dollar) {
                is_ct_param = true;
                self.advance();
            }
            if (self.current.tag != .identifier) {
                return self.fail("expected parameter name");
            }
            const param_name = self.tokenSlice(self.current);
            const param_name_span = ast.Span{ .start = self.current.loc.start, .end = self.current.loc.end };
            self.advance();
            try self.expect(.colon);
            const is_variadic = self.current.tag == .dot_dot;
            if (is_variadic) self.advance();
            const param_type = try self.parseTypeExpr();
            var is_comptime_param = false;
            if (is_ct_param and param_type.data == .type_expr) {
                const constraint_name = param_type.data.type_expr.name;
                if (std.mem.eql(u8, constraint_name, "Type")) {
                    param_type.data = .{ .type_expr = .{ .name = param_name, .is_generic = true } };
                } else {
                    is_comptime_param = true;
                }
            }
            try params.append(self.allocator, .{ .name = param_name, .name_span = param_name_span, .type_expr = param_type, .is_variadic = is_variadic, .is_comptime = is_comptime_param });
        }
        for (params.items, 0..) |param, i| {
            if (param.is_variadic and i != params.items.len - 1) {
                return self.fail("variadic parameter must be the last parameter");
            }
        }
        try self.expect(.r_paren);
        return try params.toOwnedSlice(self.allocator);
    }

    /// Recursively find all generic type names ($T) in a type expression tree.
    fn collectGenericNames(node: *Node, list: *std.ArrayList([]const u8), allocator: std.mem.Allocator) void {
        switch (node.data) {
            .type_expr => |te| {
                if (te.is_generic) list.append(allocator, te.name) catch {};
            },
            .pointer_type_expr => |pte| collectGenericNames(pte.pointee_type, list, allocator),
            .many_pointer_type_expr => |mpte| collectGenericNames(mpte.element_type, list, allocator),
            .slice_type_expr => |ste| collectGenericNames(ste.element_type, list, allocator),
            .array_type_expr => |ate| collectGenericNames(ate.element_type, list, allocator),
            .parameterized_type_expr => |pte| {
                for (pte.args) |arg| collectGenericNames(arg, list, allocator);
            },
            else => {},
        }
    }

    /// Collect generic type params and comptime value params from parameter annotations.
    fn collectTypeParams(self: *Parser, params: []const ast.Param) ![]const ast.StructTypeParam {
        var type_params = std.ArrayList(ast.StructTypeParam).empty;
        for (params) |param| {
            if (param.is_comptime) {
                var found = false;
                for (type_params.items) |existing| {
                    if (std.mem.eql(u8, existing.name, param.name)) {
                        found = true;
                        break;
                    }
                }
                if (!found) {
                    try type_params.append(self.allocator, .{ .name = param.name, .constraint = param.type_expr });
                }
            } else {
                // Collect all generic type params found anywhere in the type expression
                var generic_names = std.ArrayList([]const u8).empty;
                collectGenericNames(param.type_expr, &generic_names, self.allocator);
                for (generic_names.items) |gen_name| {
                    var found = false;
                    for (type_params.items) |existing| {
                        if (std.mem.eql(u8, existing.name, gen_name)) {
                            found = true;
                            break;
                        }
                    }
                    if (!found) {
                        const type_constraint = self.createNode(param.type_expr.span.start, .{ .type_expr = .{ .name = "Type" } }) catch continue;
                        type_params.append(self.allocator, .{ .name = gen_name, .constraint = type_constraint }) catch {};
                    }
                }
            }
        }
        return try type_params.toOwnedSlice(self.allocator);
    }

    fn parseFnDecl(self: *Parser, name: []const u8, start_pos: u32) anyerror!*Node {
        const params = try self.parseParams();

        // Optional return type
        var return_type: ?*Node = null;
        if (self.current.tag == .arrow) {
            self.advance();
            return_type = try self.parseTypeExpr();
        }

        // Body: block `{ ... }`, arrow `=> expr;`, #builtin, or #foreign marker
        const body = if (self.current.tag == .hash_builtin) blk: {
            const bi_start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            break :blk try self.createNode(bi_start, .{ .builtin_expr = {} });
        } else if (self.current.tag == .hash_foreign) blk: {
            const fi_start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            break :blk try self.createNode(fi_start, .{ .foreign_expr = {} });
        } else if (self.current.tag == .fat_arrow) blk: {
            self.advance();
            const expr = try self.parseExpr();
            try self.expect(.semicolon);
            const stmts = try self.allocator.alloc(*Node, 1);
            stmts[0] = expr;
            const block_start = expr.span.start;
            const block = try self.createNode(block_start, .{ .block = .{ .stmts = stmts } });
            break :blk block;
        } else try self.parseBlock();

        const type_params = try self.collectTypeParams(params);

        return try self.createNode(start_pos, .{ .fn_decl = .{
            .name = name,
            .params = params,
            .return_type = return_type,
            .body = body,
            .type_params = type_params,
        } });
    }

    fn parseBlock(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        try self.expect(.l_brace);
        var stmts = std.ArrayList(*Node).empty;
        while (self.current.tag != .r_brace and self.current.tag != .eof) {
            const stmt = try self.parseStmt();
            try stmts.append(self.allocator, stmt);
        }
        try self.expect(.r_brace);
        return try self.createNode(start, .{ .block = .{ .stmts = try stmts.toOwnedSlice(self.allocator) } });
    }

    /// Block-form if/match/while/bare blocks don't require trailing semicolon.
    fn expectSemicolonAfter(self: *Parser, expr: *Node) anyerror!void {
        const needs_semi = switch (expr.data) {
            .if_expr => |ie| ie.is_inline,
            .match_expr => false,
            .while_expr => false,
            .for_expr => false,
            .block => false,
            else => true,
        };
        if (needs_semi) {
            try self.expect(.semicolon);
        } else if (self.current.tag == .semicolon) {
            self.advance(); // consume optional ;
        }
    }

    pub fn parseStmt(self: *Parser) anyerror!*Node {
        // Check if this is a declaration (IDENT followed by ::, :=, or : type)
        if (self.current.tag == .identifier) {
            const saved_lexer = self.lexer;
            const saved_current = self.current;
            const saved_prev_end = self.prev_end;
            const start = self.current.loc.start;
            const name = self.tokenSlice(self.current);
            self.advance();

            if (self.current.tag == .colon_colon) {
                self.advance();
                return self.parseConstBinding(name, start);
            }
            if (self.current.tag == .colon_equal) {
                self.advance();
                const value = try self.parseExpr();
                try self.expectSemicolonAfter(value);
                return try self.createNode(start, .{ .var_decl = .{ .name = name, .type_annotation = null, .value = value } });
            }
            if (self.current.tag == .colon) {
                self.advance();
                return self.parseTypedBinding(name, start);
            }

            // Check for assignment operators
            if (self.isAssignOp()) {
                const op = self.assignOp();
                self.advance();
                const value = try self.parseExpr();
                try self.expect(.semicolon);
                const target = try self.createNode(start, .{ .identifier = .{ .name = name } });
                return try self.createNode(start, .{ .assignment = .{ .target = target, .op = op, .value = value } });
            }

            // Not a declaration or assignment — backtrack and parse as expression
            self.lexer = saved_lexer;
            self.current = saved_current;
            self.prev_end = saved_prev_end;
        }

        // Return statement: return expr; or return;
        if (self.current.tag == .kw_return) {
            const start = self.current.loc.start;
            self.advance();
            if (self.current.tag == .semicolon) {
                self.advance();
                return try self.createNode(start, .{ .return_stmt = .{ .value = null } });
            }
            const value = try self.parseExpr();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .return_stmt = .{ .value = value } });
        }

        // Defer statement: defer <expr>;
        if (self.current.tag == .kw_defer) {
            const start = self.current.loc.start;
            self.advance();
            const deferred = try self.parseExpr();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .defer_stmt = .{ .expr = deferred } });
        }

        // Break statement: break;
        if (self.current.tag == .kw_break) {
            const start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .break_expr = {} });
        }

        // Continue statement: continue;
        if (self.current.tag == .kw_continue) {
            const start = self.current.loc.start;
            self.advance();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .continue_expr = {} });
        }

        // Insert directive: #insert <expr>;
        if (self.current.tag == .hash_insert) {
            const start = self.current.loc.start;
            self.advance();
            const inner = try self.parseExpr();
            try self.expect(.semicolon);
            return try self.createNode(start, .{ .insert_expr = .{ .expr = inner } });
        }

        // Block-form if/while/for as statements — parse directly to prevent
        // postfix chaining (e.g. `if cond { ... }.field` being misparsed)
        if (self.current.tag == .kw_if) {
            const expr = try self.parseIfExpr();
            try self.expectSemicolonAfter(expr);
            return expr;
        }
        if (self.current.tag == .kw_while) {
            const expr = try self.parsePrimary();
            try self.expectSemicolonAfter(expr);
            return expr;
        }
        if (self.current.tag == .kw_for) {
            const expr = try self.parsePrimary();
            try self.expectSemicolonAfter(expr);
            return expr;
        }

        // Expression statement
        const expr = try self.parseExpr();

        // Check for field assignment: expr = value; (e.g. a.b = 1;)
        if (self.isAssignOp()) {
            const op = self.assignOp();
            self.advance();
            const value = try self.parseExpr();
            try self.expect(.semicolon);
            return try self.createNode(expr.span.start, .{ .assignment = .{ .target = expr, .op = op, .value = value } });
        }

        // Block-form if/match/while/bare blocks don't require trailing semicolon
        try self.expectSemicolonAfter(expr);
        return expr;
    }

    // ---- Expression parsing (Pratt / precedence climbing) ----

    pub fn parseExpr(self: *Parser) anyerror!*Node {
        return self.parseBinary(0);
    }

    fn parseBinary(self: *Parser, min_prec: u8) anyerror!*Node {
        const lhs = try self.parseUnary();
        return self.parseBinaryRhs(lhs, min_prec);
    }

    fn parseBinaryRhs(self: *Parser, initial_lhs: *Node, min_prec: u8) anyerror!*Node {
        var lhs = initial_lhs;

        while (true) {
            const prec = self.binaryPrec();
            if (prec == 0 or prec < min_prec) break;

            const op = self.binaryOp() orelse break;
            self.advance();

            const rhs = try self.parseBinary(prec + 1);

            // Chained comparison detection: if op is a comparison and the next
            // token is also a comparison at the same precedence, accumulate
            // into a ChainedComparison node.
            if (isComparisonOp(op) and self.binaryPrec() == prec and self.isComparisonToken()) {
                var operands = std.ArrayList(*Node).empty;
                var ops = std.ArrayList(ast.BinaryOp.Op).empty;
                try operands.append(self.allocator, lhs);
                try operands.append(self.allocator, rhs);
                try ops.append(self.allocator, op);

                while (self.binaryPrec() == prec and self.isComparisonToken()) {
                    const chain_op = self.binaryOp() orelse break;
                    self.advance();
                    const chain_rhs = try self.parseBinary(prec + 1);
                    try operands.append(self.allocator, chain_rhs);
                    try ops.append(self.allocator, chain_op);
                }

                lhs = try self.createNode(lhs.span.start, .{ .chained_comparison = .{
                    .operands = try operands.toOwnedSlice(self.allocator),
                    .ops = try ops.toOwnedSlice(self.allocator),
                } });
            } else {
                lhs = try self.createNode(lhs.span.start, .{ .binary_op = .{ .op = op, .lhs = lhs, .rhs = rhs } });
            }
        }

        return lhs;
    }

    fn parseUnary(self: *Parser) anyerror!*Node {
        if (self.current.tag == .minus) {
            const start = self.current.loc.start;
            self.advance();
            const operand = try self.parseUnary();
            return try self.createNode(start, .{ .unary_op = .{ .op = .negate, .operand = operand } });
        }
        if (self.current.tag == .bang) {
            const start = self.current.loc.start;
            self.advance();
            const operand = try self.parseUnary();
            return try self.createNode(start, .{ .unary_op = .{ .op = .not, .operand = operand } });
        }
        if (self.current.tag == .kw_xx) {
            const start = self.current.loc.start;
            self.advance();
            const operand = try self.parseUnary();
            return try self.createNode(start, .{ .unary_op = .{ .op = .xx, .operand = operand } });
        }
        if (self.current.tag == .at) {
            const start = self.current.loc.start;
            self.advance();
            const operand = try self.parseUnary();
            return try self.createNode(start, .{ .unary_op = .{ .op = .address_of, .operand = operand } });
        }
        // cast(Type) expr — prefix operator with type parameter
        if (self.current.tag == .identifier and std.mem.eql(u8, self.tokenSlice(self.current), "cast")) {
            const saved_lexer = self.lexer;
            const next_tok = self.lexer.next();
            self.lexer = saved_lexer;
            if (next_tok.tag == .l_paren) {
                const start = self.current.loc.start;
                self.advance(); // consume 'cast'
                self.advance(); // consume '('
                const type_arg = try self.parseExpr();
                try self.expect(.r_paren);
                const operand = try self.parseUnary();
                const callee = try self.createNode(start, .{ .identifier = .{ .name = "cast" } });
                const args = try self.allocator.alloc(*Node, 2);
                args[0] = type_arg;
                args[1] = operand;
                return try self.createNode(start, .{ .call = .{ .callee = callee, .args = args } });
            }
        }
        return self.parsePostfix();
    }

    fn parsePostfix(self: *Parser) anyerror!*Node {
        var expr = try self.parsePrimary();

        while (true) {
            if (self.current.tag == .l_paren) {
                // Call
                self.advance();
                var args = std.ArrayList(*Node).empty;
                while (self.current.tag != .r_paren and self.current.tag != .eof) {
                    if (args.items.len > 0) {
                        try self.expect(.comma);
                    }
                    // Spread operator: ..expr
                    if (self.current.tag == .dot_dot) {
                        const spread_start = self.current.loc.start;
                        self.advance();
                        const operand = try self.parseExpr();
                        try args.append(self.allocator, try self.createNode(spread_start, .{ .spread_expr = .{ .operand = operand } }));
                    } else {
                        try args.append(self.allocator, try self.parseExpr());
                    }
                }
                try self.expect(.r_paren);
                expr = try self.createNode(expr.span.start, .{ .call = .{ .callee = expr, .args = try args.toOwnedSlice(self.allocator) } });
            } else if (self.current.tag == .dot) {
                self.advance();
                if (self.current.tag == .l_brace) {
                    // Struct literal: Type.{ ... }
                    if (expr.data == .identifier) {
                        // Simple name: Vec4.{ ... }
                        expr = try self.parseStructLiteral(expr.data.identifier.name, null, expr.span.start);
                    } else if (expr.data == .field_access) {
                        // Qualified name: std.Vec4.{ ... }
                        const qname = try self.reconstructQualifiedName(expr);
                        expr = try self.parseStructLiteral(qname, null, expr.span.start);
                    } else {
                        // Expression type: Vec(3, f32).{ ... }
                        expr = try self.parseStructLiteral(null, expr, expr.span.start);
                    }
                } else if (self.current.tag == .l_bracket) {
                    // Typed array/vector literal: Type.[elem, ...]
                    self.advance(); // skip '['
                    var elements = std.ArrayList(*Node).empty;
                    while (self.current.tag != .r_bracket and self.current.tag != .eof) {
                        if (elements.items.len > 0) {
                            try self.expect(.comma);
                        }
                        const elem = try self.parseExpr();
                        try elements.append(self.allocator, elem);
                    }
                    try self.expect(.r_bracket);
                    expr = try self.createNode(expr.span.start, .{ .array_literal = .{
                        .elements = try elements.toOwnedSlice(self.allocator),
                        .type_expr = expr,
                    } });
                } else if (self.current.tag == .star) {
                    // Dereference: expr.*
                    self.advance();
                    expr = try self.createNode(expr.span.start, .{ .deref_expr = .{ .operand = expr } });
                } else {
                    // Field access
                    if (self.current.tag != .identifier) {
                        return self.fail("expected field name after '.'");
                    }
                    const field = self.tokenSlice(self.current);
                    self.advance();
                    expr = try self.createNode(expr.span.start, .{ .field_access = .{ .object = expr, .field = field } });
                }
            } else if (self.current.tag == .l_bracket) {
                // Index or slice access: expr[expr] or expr[start..end]
                self.advance();
                if (self.current.tag == .dot_dot) {
                    // [..end]
                    self.advance();
                    const end_expr = try self.parseExpr();
                    try self.expect(.r_bracket);
                    expr = try self.createNode(expr.span.start, .{ .slice_expr = .{
                        .object = expr, .start = null, .end = end_expr,
                    } });
                } else {
                    const first = try self.parseExpr();
                    if (self.current.tag == .dot_dot) {
                        // [start..end] or [start..]
                        self.advance();
                        const end_expr: ?*ast.Node = if (self.current.tag != .r_bracket)
                            try self.parseExpr()
                        else
                            null;
                        try self.expect(.r_bracket);
                        expr = try self.createNode(expr.span.start, .{ .slice_expr = .{
                            .object = expr, .start = first, .end = end_expr,
                        } });
                    } else {
                        // [index] — normal index access
                        try self.expect(.r_bracket);
                        expr = try self.createNode(expr.span.start, .{ .index_expr = .{
                            .object = expr, .index = first,
                        } });
                    }
                }
            } else {
                break;
            }
        }

        return expr;
    }

    fn parsePrimary(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        switch (self.current.tag) {
            .int_literal => {
                const text = self.tokenSlice(self.current);
                const base: u8 = if (text.len > 2 and text[0] == '0' and (text[1] == 'x' or text[1] == 'X'))
                    16
                else if (text.len > 2 and text[0] == '0' and (text[1] == 'b' or text[1] == 'B'))
                    2
                else
                    10;
                const digits = if (base != 10) text[2..] else text;
                const value = std.fmt.parseInt(i64, digits, base) catch {
                    return self.fail("integer literal overflow");
                };
                self.advance();
                return try self.createNode(start, .{ .int_literal = .{ .value = value } });
            },
            .float_literal => {
                const text = self.tokenSlice(self.current);
                const value = std.fmt.parseFloat(f64, text) catch {
                    return self.fail("float literal overflow");
                };
                self.advance();
                return try self.createNode(start, .{ .float_literal = .{ .value = value } });
            },
            .string_literal => {
                // raw includes quotes
                const raw = self.tokenSlice(self.current);
                self.advance();
                return try self.createNode(start, .{ .string_literal = .{ .raw = raw[1 .. raw.len - 1] } });
            },
            .raw_string_literal => {
                // #string heredoc — token span is content only, no stripping needed
                const raw = self.tokenSlice(self.current);
                self.advance();
                return try self.createNode(start, .{ .string_literal = .{ .raw = raw, .is_raw = true } });
            },
            .kw_true => {
                self.advance();
                return try self.createNode(start, .{ .bool_literal = .{ .value = true } });
            },
            .kw_false => {
                self.advance();
                return try self.createNode(start, .{ .bool_literal = .{ .value = false } });
            },
            .kw_null => {
                self.advance();
                return try self.createNode(start, .{ .null_literal = {} });
            },
            .identifier => {
                const name = self.tokenSlice(self.current);
                // Check if this identifier is a type name (e.g. s32, u8, s128)
                if (Type.fromName(name) != null) {
                    self.advance();
                    return try self.createNode(start, .{ .type_expr = .{ .name = name } });
                }
                self.advance();
                return try self.createNode(start, .{ .identifier = .{ .name = name } });
            },
            .dot => {
                self.advance();
                // Anonymous struct literal: .{ ... }
                if (self.current.tag == .l_brace) {
                    return self.parseStructLiteral(null, null, start);
                }
                // Array literal: .[expr, expr, ...]
                if (self.current.tag == .l_bracket) {
                    self.advance(); // skip '['
                    var elements = std.ArrayList(*Node).empty;
                    while (self.current.tag != .r_bracket and self.current.tag != .eof) {
                        if (elements.items.len > 0) {
                            try self.expect(.comma);
                        }
                        const elem = try self.parseExpr();
                        try elements.append(self.allocator, elem);
                    }
                    try self.expect(.r_bracket);
                    return try self.createNode(start, .{ .array_literal = .{ .elements = try elements.toOwnedSlice(self.allocator) } });
                }
                // Enum literal: .variant_name
                if (self.current.tag != .identifier) {
                    return self.fail("expected variant name, '{', or '[' after '.'");
                }
                const name = self.tokenSlice(self.current);
                self.advance();
                // Enum literal with payload: .variant(payload) — tagged enum (formerly union literal)
                if (self.current.tag == .l_paren) {
                    self.advance(); // skip '('
                    const payload = try self.parseExpr();
                    try self.expect(.r_paren);
                    return try self.createNode(start, .{ .enum_literal = .{
                        .name = name,
                        .payload = payload,
                    } });
                }
                return try self.createNode(start, .{ .enum_literal = .{ .name = name } });
            },
            .l_paren => {
                // Lambda: (params) => expr
                if (self.isLambda()) {
                    return self.parseLambda();
                }
                // Grouped expression
                self.advance();
                const expr = try self.parseExpr();
                try self.expect(.r_paren);
                return expr;
            },
            .kw_f32, .kw_f64, .kw_Type => {
                // Type keyword used as expression (for type aliases: SOME_TYPE :: f64;)
                const name = self.tokenSlice(self.current);
                self.advance();
                return try self.createNode(start, .{ .type_expr = .{ .name = name } });
            },
            .kw_struct => {
                // Anonymous struct expression: struct { value: T; count: u32; }
                return try self.parseStructDecl("__anon", start);
            },
            .kw_enum => {
                // Anonymous enum expression: enum { variant: T; other: u32; }
                return try self.parseEnumDecl("__anon", start);
            },
            .kw_union => {
                // Anonymous C-style union expression: union { f: f32; i: s32; }
                return try self.parseUnionDecl("__anon", start);
            },
            .kw_if => {
                return self.parseIfExpr();
            },
            .kw_while => {
                return self.parseWhileExpr();
            },
            .kw_for => {
                return self.parseForExpr();
            },
            .kw_break => {
                self.advance();
                return try self.createNode(start, .{ .break_expr = {} });
            },
            .kw_continue => {
                self.advance();
                return try self.createNode(start, .{ .continue_expr = {} });
            },
            .kw_return => {
                self.advance();
                // return with optional value
                const value = if (self.current.tag != .semicolon and self.current.tag != .eof)
                    try self.parseExpr()
                else
                    null;
                return try self.createNode(start, .{ .return_stmt = .{ .value = value } });
            },
            .l_bracket => {
                // Type expression in expression position: []T.[...] or [N]T.[...]
                return try self.parseTypeExpr();
            },
            .l_brace => {
                return self.parseBlock();
            },
            .triple_minus => {
                self.advance();
                return try self.createNode(start, .{ .undef_literal = {} });
            },
            .hash_run => {
                self.advance(); // skip '#run'
                const inner = try self.parseExpr();
                return try self.createNode(start, .{ .comptime_expr = .{ .expr = inner } });
            },
            else => {
                return self.fail("unexpected token in expression");
            },
        }
    }

    fn parseIfExpr(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        self.advance(); // skip 'if'

        // Parse condition at prec 5 (arithmetic+), leaving all comparisons
        // unconsumed for manual handling with match disambiguation.
        var condition = try self.parseBinary(5);

        // Handle comparisons with chain detection and match disambiguation.
        // All comparisons (< <= > >= == !=) are at the same precedence.
        if (self.isComparisonToken()) {
            var operands = std.ArrayList(*Node).empty;
            var ops = std.ArrayList(ast.BinaryOp.Op).empty;
            try operands.append(self.allocator, condition);

            while (self.isComparisonToken()) {
                // Match disambiguation: == followed by { is a match expression
                if (self.current.tag == .equal_equal) {
                    self.advance();
                    if (self.current.tag == .l_brace) {
                        // Match expression: if expr == { case ... }
                        // Only valid as the first comparison (no chain before it)
                        if (ops.items.len == 0) {
                            return self.parseMatchBody(condition, start);
                        }
                        // Chain followed by == { is an error — fall through to
                        // regular comparison (will likely fail at parse time)
                    }
                    const rhs = try self.parseBinary(5);
                    try operands.append(self.allocator, rhs);
                    try ops.append(self.allocator, .eq);
                } else {
                    const cmp_op = self.binaryOp() orelse break;
                    self.advance();
                    const rhs = try self.parseBinary(5);
                    try operands.append(self.allocator, rhs);
                    try ops.append(self.allocator, cmp_op);
                }
            }

            if (ops.items.len == 1) {
                // Single comparison — regular binary_op
                condition = try self.createNode(condition.span.start, .{ .binary_op = .{
                    .op = ops.items[0],
                    .lhs = operands.items[0],
                    .rhs = operands.items[1],
                } });
            } else {
                // Chained comparison
                condition = try self.createNode(condition.span.start, .{ .chained_comparison = .{
                    .operands = try operands.toOwnedSlice(self.allocator),
                    .ops = try ops.toOwnedSlice(self.allocator),
                } });
            }
        }

        // Handle and/or with proper Pratt precedence
        condition = try self.parseBinaryRhs(condition, 1);

        // Inline form: if cond then expr [else expr]
        if (self.current.tag == .kw_then) {
            self.advance();
            const then_branch = try self.parseExpr();
            var else_branch: ?*Node = null;
            if (self.current.tag == .kw_else) {
                self.advance();
                else_branch = try self.parseExpr();
            }
            return try self.createNode(start, .{ .if_expr = .{
                .condition = condition,
                .then_branch = then_branch,
                .else_branch = else_branch,
                .is_inline = true,
            } });
        }

        // Block form: if cond { ... } else { ... }
        const then_branch = try self.parseBlock();
        var else_branch: ?*Node = null;
        if (self.current.tag == .kw_else) {
            self.advance();
            if (self.current.tag == .kw_if) {
                else_branch = try self.parseIfExpr();
            } else {
                else_branch = try self.parseBlock();
            }
        }
        return try self.createNode(start, .{ .if_expr = .{
            .condition = condition,
            .then_branch = then_branch,
            .else_branch = else_branch,
            .is_inline = false,
        } });
    }

    fn parseWhileExpr(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        self.advance(); // skip 'while'

        const condition = try self.parseExpr();
        const body = try self.parseBlock();

        return try self.createNode(start, .{ .while_expr = .{
            .condition = condition,
            .body = body,
        } });
    }

    fn parseForExpr(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        self.advance(); // skip 'for'

        const iterable = try self.parseExpr();
        const body = try self.parseBlock();

        return try self.createNode(start, .{ .for_expr = .{
            .iterable = iterable,
            .body = body,
        } });
    }

    fn parseMatchBody(self: *Parser, subject: *Node, start_pos: u32) anyerror!*Node {
        try self.expect(.l_brace);
        var arms = std.ArrayList(ast.MatchArm).empty;
        while (self.current.tag == .kw_case) {
            const arm_start = self.current.loc.start;
            self.advance(); // skip 'case'
            // Allow keyword tokens (struct, enum, union) as type category names in match arms
            const pattern: *Node = if (self.current.tag == .kw_struct or self.current.tag == .kw_enum or self.current.tag == .kw_union) blk: {
                const name = self.tokenSlice(self.current);
                self.advance();
                break :blk try self.createNode(arm_start, .{ .identifier = .{ .name = name } });
            } else try self.parsePrimary(); // .variant
            try self.expect(.colon);

            // Optional payload capture: (ident)
            var capture: ?[]const u8 = null;
            if (self.current.tag == .l_paren) {
                self.advance();
                if (self.current.tag != .identifier) return self.fail("expected capture name");
                capture = self.tokenSlice(self.current);
                self.advance();
                try self.expect(.r_paren);
            }

            if (self.current.tag == .kw_break) {
                self.advance();
                try self.expect(.semicolon);
                const body = try self.createNode(arm_start, .{ .block = .{ .stmts = &.{} } });
                try arms.append(self.allocator, .{ .pattern = pattern, .body = body, .is_break = true, .capture = capture });
            } else if (self.current.tag == .fat_arrow) {
                // Short form: (ident) => expr;
                self.advance();
                const expr = try self.parseExpr();
                try self.expect(.semicolon);
                const body = try self.createNode(arm_start, .{ .block = .{ .stmts = try self.allocator.dupe(*Node, &.{expr}) } });
                try arms.append(self.allocator, .{ .pattern = pattern, .body = body, .is_break = false, .capture = capture });
            } else {
                const stmts_start = self.current.loc.start;
                var stmts = std.ArrayList(*Node).empty;
                while (self.current.tag != .kw_case and self.current.tag != .kw_else and self.current.tag != .r_brace and self.current.tag != .eof) {
                    try stmts.append(self.allocator, try self.parseStmt());
                }
                const body = try self.createNode(stmts_start, .{ .block = .{ .stmts = try stmts.toOwnedSlice(self.allocator) } });
                try arms.append(self.allocator, .{ .pattern = pattern, .body = body, .is_break = false, .capture = capture });
            }
        }
        // Optional else arm (default)
        if (self.current.tag == .kw_else) {
            const else_start = self.current.loc.start;
            self.advance(); // skip 'else'
            try self.expect(.colon);
            var stmts = std.ArrayList(*Node).empty;
            while (self.current.tag != .r_brace and self.current.tag != .eof) {
                try stmts.append(self.allocator, try self.parseStmt());
            }
            const body = try self.createNode(else_start, .{ .block = .{ .stmts = try stmts.toOwnedSlice(self.allocator) } });
            try arms.append(self.allocator, .{ .pattern = null, .body = body, .is_break = false });
        }
        try self.expect(.r_brace);
        return try self.createNode(start_pos, .{ .match_expr = .{ .subject = subject, .arms = try arms.toOwnedSlice(self.allocator) } });
    }

    fn isLambda(self: *Parser) bool {
        // Peek ahead: save state, scan to matching ), check if => or -> ... => follows
        const saved_lexer = self.lexer;
        const saved_current = self.current;
        const saved_prev_end = self.prev_end;
        defer {
            self.lexer = saved_lexer;
            self.current = saved_current;
            self.prev_end = saved_prev_end;
        }

        self.advance(); // skip '('
        var depth: u32 = 1;
        while (depth > 0 and self.current.tag != .eof) {
            if (self.current.tag == .l_paren) depth += 1;
            if (self.current.tag == .r_paren) depth -= 1;
            if (depth > 0) self.advance();
        }
        if (self.current.tag == .r_paren) {
            self.advance(); // skip ')'
            if (self.current.tag == .fat_arrow) return true;
            // (params) -> ReturnType => expr
            if (self.current.tag == .arrow) {
                self.advance(); // skip '->'
                // Skip past the return type tokens until we see '=>' or something unexpected
                while (self.current.tag != .eof) {
                    if (self.current.tag == .fat_arrow) return true;
                    // Return type tokens: identifiers, dots, parens, type keywords, dollar, brackets
                    if (self.current.tag == .identifier or self.current.tag.isTypeKeyword() or
                        self.current.tag == .dot or self.current.tag == .dollar or
                        self.current.tag == .l_bracket or self.current.tag == .r_bracket or
                        self.current.tag == .l_paren or self.current.tag == .r_paren or
                        self.current.tag == .comma or self.current.tag == .int_literal)
                    {
                        self.advance();
                    } else break;
                }
            }
        }
        return false;
    }

    fn parseLambda(self: *Parser) anyerror!*Node {
        const start = self.current.loc.start;
        const params = try self.parseParams();

        // Optional return type: (params) -> Type => expr
        var return_type: ?*Node = null;
        if (self.current.tag == .arrow) {
            self.advance();
            return_type = try self.parseTypeExpr();
        }

        try self.expect(.fat_arrow);
        const body = try self.parseExpr();
        const type_params = try self.collectTypeParams(params);
        return try self.createNode(start, .{ .lambda = .{
            .params = params,
            .return_type = return_type,
            .body = body,
            .type_params = type_params,
        } });
    }

    // ---- Helpers ----

    fn isFunctionDef(self: *Parser) bool {
        // Peek ahead: save state, scan to matching ), check what follows
        const saved_lexer = self.lexer;
        const saved_current = self.current;
        const saved_prev_end = self.prev_end;
        defer {
            self.lexer = saved_lexer;
            self.current = saved_current;
            self.prev_end = saved_prev_end;
        }

        self.advance(); // skip '('
        var depth: u32 = 1;
        while (depth > 0 and self.current.tag != .eof) {
            if (self.current.tag == .l_paren) depth += 1;
            if (self.current.tag == .r_paren) depth -= 1;
            if (depth > 0) self.advance();
        }
        if (self.current.tag == .r_paren) {
            self.advance(); // skip ')'
            // Function if followed by '{', '->', '#builtin', '#foreign', or '=>'
            return self.current.tag == .l_brace or self.current.tag == .arrow or self.current.tag == .hash_builtin or self.current.tag == .hash_foreign or self.current.tag == .fat_arrow;
        }
        return false;
    }

    fn isAssignOp(self: *const Parser) bool {
        return switch (self.current.tag) {
            .equal, .plus_equal, .minus_equal, .star_equal, .slash_equal, .percent_equal => true,
            else => false,
        };
    }

    fn assignOp(self: *const Parser) ast.Assignment.Op {
        return switch (self.current.tag) {
            .equal => .assign,
            .plus_equal => .add_assign,
            .minus_equal => .sub_assign,
            .star_equal => .mul_assign,
            .slash_equal => .div_assign,
            .percent_equal => .mod_assign,
            else => unreachable,
        };
    }

    fn binaryPrec(self: *const Parser) u8 {
        return switch (self.current.tag) {
            .kw_or => 1,
            .kw_and => 2,
            .pipe => 3,
            .ampersand => 3,
            .equal_equal, .bang_equal, .less, .less_equal, .greater, .greater_equal => 4,
            .plus, .minus => 5,
            .star, .slash, .percent => 6,
            else => 0,
        };
    }

    fn binaryOp(self: *const Parser) ?ast.BinaryOp.Op {
        return switch (self.current.tag) {
            .kw_and => .and_op,
            .kw_or => .or_op,
            .pipe => .bit_or,
            .ampersand => .bit_and,
            .plus => .add,
            .minus => .sub,
            .star => .mul,
            .slash => .div,
            .percent => .mod,
            .equal_equal => .eq,
            .bang_equal => .neq,
            .less => .lt,
            .less_equal => .lte,
            .greater => .gt,
            .greater_equal => .gte,
            else => null,
        };
    }

    fn isComparisonOp(op: ast.BinaryOp.Op) bool {
        return switch (op) {
            .lt, .lte, .gt, .gte, .eq, .neq => true,
            else => false,
        };
    }

    fn isComparisonToken(self: *const Parser) bool {
        return switch (self.current.tag) {
            .less, .less_equal, .greater, .greater_equal, .equal_equal, .bang_equal => true,
            else => false,
        };
    }

    fn advance(self: *Parser) void {
        self.prev_end = self.current.loc.end;
        self.current = self.lexer.next();
    }

    fn expect(self: *Parser, tag: Tag) !void {
        if (self.current.tag != tag) {
            const expected = tag.lexeme() orelse @tagName(tag);
            return self.failFmt("expected '{s}'", .{expected});
        }
        self.advance();
    }

    fn failFmt(self: *Parser, comptime fmt: []const u8, args: anytype) error{ParseError} {
        const msg = std.fmt.allocPrint(self.allocator, fmt, args) catch return error.ParseError;
        return self.fail(msg);
    }

    fn tokenSlice(self: *const Parser, token: Token) []const u8 {
        return self.source[token.loc.start..token.loc.end];
    }

    fn fail(self: *Parser, msg: []const u8) error{ParseError} {
        self.err_msg = msg;
        self.err_offset = self.current.loc.start;
        if (self.diagnostics) |diags| {
            diags.add(.err, msg, .{ .start = self.current.loc.start, .end = self.current.loc.end });
        }
        return error.ParseError;
    }
};

test "parse minimal main" {
    const source = "main :: () { 42; }";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expect(root.data == .root);
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .fn_decl);
    try std.testing.expectEqualStrings("main", decl.data.fn_decl.name);
    const body = decl.data.fn_decl.body;
    try std.testing.expect(body.data == .block);
    try std.testing.expectEqual(@as(usize, 1), body.data.block.stmts.len);
    try std.testing.expect(body.data.block.stmts[0].data == .int_literal);
    try std.testing.expectEqual(@as(i64, 42), body.data.block.stmts[0].data.int_literal.value);
}

test "parse #run const binding" {
    const source = "x :: #run compute(5);";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .const_decl);
    try std.testing.expectEqualStrings("x", decl.data.const_decl.name);
    try std.testing.expect(decl.data.const_decl.value.data == .comptime_expr);
    // inner expr is a call
    try std.testing.expect(decl.data.const_decl.value.data.comptime_expr.expr.data == .call);
}

test "parse top-level #run" {
    const source = "#run main();";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .comptime_expr);
    // inner expr is a call
    try std.testing.expect(decl.data.comptime_expr.expr.data == .call);
}

test "parse flat import" {
    const source = "#import \"modules/std/math.sx\";";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .import_decl);
    try std.testing.expectEqualStrings("modules/std/math.sx", decl.data.import_decl.path);
    try std.testing.expect(decl.data.import_decl.name == null);
}

test "parse namespaced import" {
    const source = "std :: #import \"modules/std/std.sx\";";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .import_decl);
    try std.testing.expectEqualStrings("modules/std/std.sx", decl.data.import_decl.path);
    try std.testing.expectEqualStrings("std", decl.data.import_decl.name.?);
}

test "parse library declaration" {
    const source = "#library \"raylib\";";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .library_decl);
    try std.testing.expectEqualStrings("raylib", decl.data.library_decl.lib_name);
}

test "parse void function with builtin body" {
    const source = "foo :: () #builtin;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .fn_decl);
    try std.testing.expectEqualStrings("foo", decl.data.fn_decl.name);
    try std.testing.expect(decl.data.fn_decl.body.data == .builtin_expr);
}

test "parse void function with foreign body" {
    const source = "InitWindow :: (width: s32, height: s32, title: *u8) -> void #foreign;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .fn_decl);
    try std.testing.expectEqualStrings("InitWindow", decl.data.fn_decl.name);
    try std.testing.expect(decl.data.fn_decl.body.data == .foreign_expr);
    try std.testing.expectEqual(@as(usize, 3), decl.data.fn_decl.params.len);
}

test "parse void function with arrow body" {
    const source = "foo :: () => 42;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    try std.testing.expectEqual(@as(usize, 1), root.data.root.decls.len);
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .fn_decl);
    try std.testing.expectEqualStrings("foo", decl.data.fn_decl.name);
    try std.testing.expect(decl.data.fn_decl.body.data == .int_literal);
    try std.testing.expectEqual(@as(i64, 42), decl.data.fn_decl.body.data.int_literal.value);
}

test "parse hex and binary literals" {
    const source = "main :: () { 0xFF; 0b1010; }";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    const body = root.data.root.decls[0].data.fn_decl.body;
    try std.testing.expectEqual(@as(usize, 2), body.data.block.stmts.len);
    try std.testing.expectEqual(@as(i64, 255), body.data.block.stmts[0].data.int_literal.value);
    try std.testing.expectEqual(@as(i64, 10), body.data.block.stmts[1].data.int_literal.value);
}

test "parse array type with identifier length" {
    const source = "foo :: (arr: [N]f32) => arr;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .fn_decl);
    const param_type = decl.data.fn_decl.params[0].type_expr;
    try std.testing.expect(param_type.data == .array_type_expr);
    // length is an identifier "N", not an int literal
    try std.testing.expect(param_type.data.array_type_expr.length.data == .identifier);
    try std.testing.expectEqualStrings("N", param_type.data.array_type_expr.length.data.identifier.name);
    try std.testing.expect(param_type.data.array_type_expr.element_type.data == .type_expr);
}

test "parse lambda with generic params" {
    const source = "f :: (x: $T) => x;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .const_decl);
    const lambda = decl.data.const_decl.value;
    try std.testing.expect(lambda.data == .lambda);
    try std.testing.expectEqual(@as(usize, 1), lambda.data.lambda.params.len);
    try std.testing.expectEqualStrings("x", lambda.data.lambda.params[0].name);
    // has generic type param
    try std.testing.expectEqual(@as(usize, 1), lambda.data.lambda.type_params.len);
    try std.testing.expectEqualStrings("T", lambda.data.lambda.type_params[0].name);
}

test "parse lambda with return type" {
    const source = "f :: (x: s32) -> s32 => x;";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    const decl = root.data.root.decls[0];
    try std.testing.expect(decl.data == .const_decl);
    const lambda = decl.data.const_decl.value;
    try std.testing.expect(lambda.data == .lambda);
    try std.testing.expect(lambda.data.lambda.return_type != null);
    try std.testing.expect(lambda.data.lambda.return_type.?.data == .type_expr);
    try std.testing.expectEqualStrings("s32", lambda.data.lambda.return_type.?.data.type_expr.name);
}

test "parse match with else arm" {
    const source =
        \\main :: () {
        \\  x := 5;
        \\  if x == {
        \\    case 1: 10;
        \\    else: 99;
        \\  };
        \\}
    ;
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const root = try parser.parse();
    const body = root.data.root.decls[0].data.fn_decl.body;
    // second stmt is the match expr (after var decl)
    const match_node = body.data.block.stmts[1];
    try std.testing.expect(match_node.data == .match_expr);
    const arms = match_node.data.match_expr.arms;
    try std.testing.expectEqual(@as(usize, 2), arms.len);
    // first arm has a pattern
    try std.testing.expect(arms[0].pattern != null);
    // second arm is the else arm (null pattern)
    try std.testing.expect(arms[1].pattern == null);
}

test "integer literal overflow error" {
    const source = "main :: () { 99999999999999999999; }";
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    var parser = Parser.init(arena.allocator(), source);
    const result = parser.parse();
    try std.testing.expectError(error.ParseError, result);
    try std.testing.expectEqualStrings("integer literal overflow", parser.err_msg.?);
}