const std = @import("std");
const Lexer = @import("lexer.zig").Lexer;
const Token = @import("lexer.zig").Token;
pub const TokenType = @import("lexer.zig").TokenType;

pub const DataType = union(enum) {
    primitive: Primitive,
    slice: *DataType,
    optional: *DataType,
    custom: []const u8,

    pub fn deinit(self: DataType, allocator: std.mem.Allocator) void {
        switch (self) {
            .slice => |inner| {
                inner.deinit(allocator);
                allocator.destroy(inner);
            },
            .optional => |inner| {
                inner.deinit(allocator);
                allocator.destroy(inner);
            },
            else => {},
        }
    }

    pub const Primitive = enum {
        int,
        float,
        str,
        bool,
        void_type,
        inferred,
    };
};

pub const Expr = union(enum) {
    int_lit: i64,
    float_lit: f64,
    string_lit: []const u8,
    bool_lit: bool,
    null_lit,
    /// Enum-literal shorthand `.name` (e.g. `.cpu`). Lexeme points into the
    /// source; no allocation.
    enum_lit: []const u8,
    /// Anonymous fn literal — apply_simd kernel sugar. Heap-boxed so the
    /// union stays small; owned by the Expr.
    fn_lit: *FnLit,
    identifier: []const u8,
    array_lit: []Expr,
    struct_lit: []StructFieldInit,
    binary: struct {
        op: TokenType,
        lhs: *Expr,
        rhs: *Expr,
    },
    unary: struct {
        op: TokenType,
        operand: *Expr,
    },
    call: struct {
        callee: *Expr,
        args: []Expr,
    },
    call_directive: struct {
        name: []const u8,
        args: []Expr,
    },
    field_access: struct {
        object: *Expr,
        field: []const u8,
    },
    index_access: struct {
        object: *Expr,
        index: *Expr,
    },

    pub const StructFieldInit = struct {
        name: []const u8,
        explicit_type: ?DataType,
        value: Expr,
        key_expr: ?Expr = null,

        fn deinit(self: StructFieldInit, allocator: std.mem.Allocator) void {
            self.value.deinit(allocator);
            if (self.key_expr) |ke| ke.deinit(allocator);
            if (self.explicit_type) |t| t.deinit(allocator);
        }
    };

    pub fn deinit(self: Expr, allocator: std.mem.Allocator) void {
        switch (self) {
            .string_lit => |s| allocator.free(s),
            .enum_lit => {},
            .fn_lit => |fl| {
                fl.deinit(allocator);
                allocator.destroy(fl);
            },
            .array_lit => |items| {
                for (items) |item| item.deinit(allocator);
                allocator.free(items);
            },
            .struct_lit => |fields| {
                for (fields) |f| f.deinit(allocator);
                allocator.free(fields);
            },
            .binary => |b| {
                b.lhs.deinit(allocator);
                allocator.destroy(b.lhs);
                b.rhs.deinit(allocator);
                allocator.destroy(b.rhs);
            },
            .unary => |u| {
                u.operand.deinit(allocator);
                allocator.destroy(u.operand);
            },
            .call => |c| {
                c.callee.deinit(allocator);
                allocator.destroy(c.callee);
                for (c.args) |arg| arg.deinit(allocator);
                allocator.free(c.args);
            },
            .call_directive => |cd| {
                for (cd.args) |arg| arg.deinit(allocator);
                allocator.free(cd.args);
            },
            .field_access => |fa| {
                fa.object.deinit(allocator);
                allocator.destroy(fa.object);
            },
            .index_access => |ia| {
                ia.object.deinit(allocator);
                allocator.destroy(ia.object);
                ia.index.deinit(allocator);
                allocator.destroy(ia.index);
            },
            else => {},
        }
    }
};

pub const Stmt = union(enum) {
    var_decl: struct {
        is_const: bool,
        name: []const u8,
        type_ann: ?DataType,
        initializer: Expr,
    },
    assignment: struct {
        target: Expr,
        value: Expr,
    },
    expr_stmt: Expr,
    if_stmt: struct {
        condition: Expr,
        then_branch: Block,
        else_branch: ?Block,
    },
    while_stmt: struct {
        condition: Expr,
        body: Block,
    },
    for_stmt: struct {
        var_name: []const u8,
        iter: Expr,
        body: Block,
    },
    /// `for i in lo..hi { ... }` — numeric range loop; iterates i = lo..hi-1
    /// (inclusive of lo, exclusive of hi). With `..=` the bound is inclusive.
    range_for: struct {
        var_name: []const u8,
        lo: Expr,
        hi: Expr,
        hi_inclusive: bool,
        body: Block,
    },
    break_stmt,
    continue_stmt,
    return_stmt: ?Expr,
    defer_stmt: Expr,
    fn_decl: FnDecl,

    pub fn deinit(self: Stmt, allocator: std.mem.Allocator) void {
        switch (self) {
            .var_decl => |d| d.initializer.deinit(allocator),
            .assignment => |a| {
                a.target.deinit(allocator);
                a.value.deinit(allocator);
            },
            .expr_stmt => |e| e.deinit(allocator),
            .if_stmt => |i| {
                i.condition.deinit(allocator);
                i.then_branch.deinit(allocator);
                if (i.else_branch) |eb| eb.deinit(allocator);
            },
            .while_stmt => |w| {
                w.condition.deinit(allocator);
                w.body.deinit(allocator);
            },
            .for_stmt => |f| {
                allocator.free(f.var_name);
                f.iter.deinit(allocator);
                f.body.deinit(allocator);
            },
            .range_for => |r| {
                allocator.free(r.var_name);
                r.lo.deinit(allocator);
                r.hi.deinit(allocator);
                r.body.deinit(allocator);
            },
            .return_stmt => |r| if (r) |e| e.deinit(allocator),
            .defer_stmt => |e| e.deinit(allocator),
            .break_stmt, .continue_stmt => {},
            .fn_decl => |fd| fd.deinitFields(allocator),
        }
    }
};

/// A `{ ... }` statement block plus parallel arrays of the source line
/// each statement started on and its starting column (for diagnostics).
pub const Block = struct {
    stmts: []Stmt,
    lines: []u32,
    cols: []u32,

    pub fn deinit(self: Block, allocator: std.mem.Allocator) void {
        for (self.stmts) |s| s.deinit(allocator);
        allocator.free(self.stmts);
        allocator.free(self.lines);
        allocator.free(self.cols);
    }
};

pub const FnParam = struct {
    name: []const u8,
    param_type: DataType,
};

/// An anonymous `fn(x: T) T { ... }` literal. Only legal as the kernel
/// argument of a tensor `apply_simd(...)` method call (enforced by the
/// compiler, which lowers the literal into a hidden function chunk).
pub const FnLit = struct {
    params: []FnParam,
    body: Block,
    /// Lane hint captured from a `@vector(N, T)` parameter annotation.
    /// 0 means unspecified (the VM falls back to scalar execution; the
    /// native backend uses it to pick vector widths).
    vector_width: u64 = 0,

    pub fn deinit(self: *FnLit, allocator: std.mem.Allocator) void {
        for (self.params) |p| p.param_type.deinit(allocator);
        allocator.free(self.params);
        self.body.deinit(allocator);
    }
};

pub const Decorator = struct {
    name: []const u8,
    args: []Expr,

    pub fn deinit(self: Decorator, allocator: std.mem.Allocator) void {
        for (self.args) |arg| arg.deinit(allocator);
        allocator.free(self.args);
    }
};

pub const FnDecl = struct {
    is_pub: bool,
    name: []const u8,
    params: []FnParam,
    return_type: DataType,
    decorators: []Decorator,
    body: Block,

    pub fn deinitFields(self: FnDecl, allocator: std.mem.Allocator) void {
        for (self.params) |p| p.param_type.deinit(allocator);
        allocator.free(self.params);
        self.return_type.deinit(allocator);
        for (self.decorators) |dec| dec.deinit(allocator);
        allocator.free(self.decorators);
        self.body.deinit(allocator);
    }

    pub fn deinit(self: FnDecl, allocator: std.mem.Allocator) void {
        self.deinitFields(allocator);
    }
};

pub const Directive = union(enum) {
    import_directive: []const u8,
    include_directive: []Expr,
    generic: GenericDirective,

    pub const GenericDirective = struct {
        name: []const u8,
        args: []Expr,

        pub fn deinit(self: GenericDirective, allocator: std.mem.Allocator) void {
            for (self.args) |arg| arg.deinit(allocator);
            allocator.free(self.args);
        }
    };
};

pub const Program = struct {
    directives: []Directive,
    fn_decls: []FnDecl,
    global_stmts: Block,

    pub fn deinit(self: Program, allocator: std.mem.Allocator) void {
        for (self.directives) |d| switch (d) {
            .include_directive => |exprs| {
                for (exprs) |e| e.deinit(allocator);
                allocator.free(exprs);
            },
            .generic => |g| g.deinit(allocator),
            else => {},
        };
        allocator.free(self.directives);
        for (self.fn_decls) |fd| fd.deinitFields(allocator);
        allocator.free(self.fn_decls);
        self.global_stmts.deinit(allocator);
    }
};

pub const Parser = struct {
    lexer: Lexer,
    allocator: std.mem.Allocator,
    current_token: Token,
    /// Diagnostics for the most recent failure (set by `fail`).
    err_line: usize = 0,
    err_col: usize = 0,
    err_msg: ?[]u8 = null,

    pub fn init(allocator: std.mem.Allocator, source: []const u8) Parser {
        var lex = Lexer.init(source);
        const tok = lex.nextToken();
        return .{ .lexer = lex, .allocator = allocator, .current_token = tok };
    }

    pub fn deinit(self: *Parser) void {
        if (self.err_msg) |m| self.allocator.free(m);
        self.err_msg = null;
    }

    fn advance(self: *Parser) void {
        self.current_token = self.lexer.nextToken();
    }

    fn match(self: *Parser, tag: TokenType) bool {
        if (self.current_token.tag == tag) {
            self.advance();
            return true;
        }
        return false;
    }

    fn expect(self: *Parser, tag: TokenType) ParseError!Token {
        if (self.current_token.tag == tag) {
            const tok = self.current_token;
            self.advance();
            return tok;
        }
        return self.fail(error.UnexpectedToken, "expected '{s}', found '{s}'", .{ @tagName(tag), @tagName(self.current_token.tag) });
    }

    /// Records the current token position plus a formatted message for
    /// diagnostics, then returns `e`.
    fn fail(self: *Parser, e: ParseError, comptime fmt: []const u8, args: anytype) ParseError {
        self.err_line = self.current_token.line;
        self.err_col = self.current_token.column;
        self.err_msg = std.fmt.allocPrint(self.allocator, fmt ++ " (at line {d}:{d})", args ++ .{ self.current_token.line, self.current_token.column }) catch null;
        return e;
    }

    pub fn parseProgram(self: *Parser) ParseError!Program {
        var directives: std.ArrayList(Directive) = .empty;
        var fn_decls: std.ArrayList(FnDecl) = .empty;
        var global_stmts: std.ArrayList(Stmt) = .empty;
        var global_lines: std.ArrayList(u32) = .empty;
        var global_cols: std.ArrayList(u32) = .empty;
        var pending_decorators: std.ArrayList(Decorator) = .empty;
        defer pending_decorators.deinit(self.allocator);

        while (self.current_token.tag != .eof) {
            if (self.current_token.tag == .at_directive) {
                const name = self.current_token.directiveName();

                if (std.mem.eql(u8, name, "import")) {
                    try self.parseImport(&directives);
                    continue;
                } else if (std.mem.eql(u8, name, "include")) {
                    try self.parseInclude(&directives);
                    continue;
                }

                // Generic directive or decorator: `@name(args...)`
                const dir_name = name;
                const dir_line = self.current_token.line;
                const dir_col: u32 = @intCast(self.current_token.column);
                self.advance();
                var args: std.ArrayList(Expr) = .empty;
                if (self.match(.l_paren)) {
                    while (self.current_token.tag != .r_paren and self.current_token.tag != .eof) {
                        try args.append(self.allocator, try self.parseOr());
                        if (!self.match(.comma)) break;
                    }
                    _ = try self.expect(.r_paren);
                }
                if (self.current_token.tag == .kw_pub or self.current_token.tag == .kw_fn) {
                    // Decorator directly preceding a function declaration.
                    try pending_decorators.append(self.allocator, .{ .name = dir_name, .args = try args.toOwnedSlice(self.allocator) });
                    continue;
                }
                if (pending_decorators.items.len > 0) {
                    return error.UnexpectedToken;
                }
                _ = self.match(.semicolon);
                if (!isModuleDirective(dir_name)) {
                    // Executable builtin call at global scope, e.g.
                    // `@stdout("hi")` — compile it as an expression stmt.
                    try global_lines.append(self.allocator, @intCast(dir_line));
                    try global_cols.append(self.allocator, dir_col);
                    try global_stmts.append(self.allocator, .{ .expr_stmt = .{ .call_directive = .{ .name = dir_name, .args = try args.toOwnedSlice(self.allocator) } } });
                    continue;
                }
                try directives.append(self.allocator, .{ .generic = .{ .name = dir_name, .args = try args.toOwnedSlice(self.allocator) } });
            } else if (self.current_token.tag == .kw_pub or self.current_token.tag == .kw_fn) {
                var decl = try self.parseFnDecl();
                decl.decorators = try mergeDecorators(self.allocator, &pending_decorators, decl.decorators);
                try fn_decls.append(self.allocator, decl);
            } else {
                try global_lines.append(self.allocator, @intCast(self.current_token.line));
                try global_cols.append(self.allocator, @intCast(self.current_token.column));
                try global_stmts.append(self.allocator, try self.parseStmt());
            }
        }

        return .{
            .directives = try directives.toOwnedSlice(self.allocator),
            .fn_decls = try fn_decls.toOwnedSlice(self.allocator),
            .global_stmts = .{
                .stmts = try global_stmts.toOwnedSlice(self.allocator),
                .lines = try global_lines.toOwnedSlice(self.allocator),
                .cols = try global_cols.toOwnedSlice(self.allocator),
            },
        };
    }

    fn parseImport(self: *Parser, directives: *std.ArrayList(Directive)) ParseError!void {
        self.advance();
        _ = try self.expect(.assign);
        if (self.match(.period)) {
            _ = try self.expect(.l_brace);
            while (self.current_token.tag != .r_brace and self.current_token.tag != .eof) {
                if (self.current_token.tag == .identifier) {
                    const mod_tok = self.current_token;
                    self.advance();
                    try directives.append(self.allocator, .{ .import_directive = mod_tok.lexeme });
                }
                if (!self.match(.comma)) break;
            }
            _ = try self.expect(.r_brace);
        } else if (self.match(.l_brace)) {
            while (self.current_token.tag != .r_brace and self.current_token.tag != .eof) {
                if (self.current_token.tag == .identifier) {
                    const mod_tok = self.current_token;
                    self.advance();
                    try directives.append(self.allocator, .{ .import_directive = mod_tok.lexeme });
                }
                if (!self.match(.comma)) break;
            }
            _ = try self.expect(.r_brace);
        } else if (self.current_token.tag == .string_literal) {
            const str_tok = self.current_token;
            self.advance();
            try directives.append(self.allocator, .{ .import_directive = str_tok.lexeme });
        }
        _ = self.match(.semicolon);
    }

    fn parseInclude(self: *Parser, directives: *std.ArrayList(Directive)) ParseError!void {
        self.advance();
        _ = self.match(.colon);
        _ = self.match(.l_bracket);
        _ = self.match(.r_bracket);
        _ = self.match(.type_str);
        _ = try self.expect(.assign);
        const list_expr = try self.parseOr();
        if (list_expr == .array_lit) {
            try directives.append(self.allocator, .{ .include_directive = list_expr.array_lit });
        } else {
            list_expr.deinit(self.allocator);
        }
        _ = self.match(.semicolon);
    }

    fn mergeDecorators(allocator: std.mem.Allocator, pending: *std.ArrayList(Decorator), existing: []Decorator) ![]Decorator {
        if (pending.items.len == 0) return existing;
        var merged: std.ArrayList(Decorator) = .empty;
        try merged.appendSlice(allocator, pending.items);
        try merged.appendSlice(allocator, existing);
        pending.clearRetainingCapacity();
        if (existing.len > 0) allocator.free(existing);
        return merged.toOwnedSlice(allocator);
    }

    /// Directive names that describe the module (types, build, imports)
    /// rather than perform runtime work. Anything else written as a
    /// statement — e.g. `@stdout(...)` — is executable code.
    fn isModuleDirective(name: []const u8) bool {
        const module_level = [_][]const u8{
            "import_c", "load_lib", "comptime", "build",
            "tensor",   "fetch",    "fetch_pkg",
        };
        for (module_level) |m| {
            if (std.mem.eql(u8, name, m)) return true;
        }
        return false;
    }

    fn parseFnDecl(self: *Parser) ParseError!FnDecl {
        var decorators: std.ArrayList(Decorator) = .empty;

        const is_pub = self.match(.kw_pub);
        _ = try self.expect(.kw_fn);
        const name_tok = try self.expect(.identifier);

        _ = try self.expect(.l_paren);
        var params: std.ArrayList(FnParam) = .empty;
        while (self.current_token.tag != .r_paren and self.current_token.tag != .eof) {
            const p_name = try self.expect(.identifier);
            _ = try self.expect(.colon);
            try params.append(self.allocator, .{ .name = p_name.lexeme, .param_type = try self.parseDataType() });
            if (!self.match(.comma)) break;
        }
        _ = try self.expect(.r_paren);

        // Return type: optional leading colon (`-> int` style) or bare type
        // token as in `pub fn main(args: []str) u32 { ... }`.
        var ret_type: DataType = .{ .primitive = .void_type };
        if (self.match(.colon)) {
            ret_type = try self.parseDataType();
        } else if (self.startsType()) {
            ret_type = try self.parseDataType();
        }

        const body = try self.parseBlock();

        return .{
            .is_pub = is_pub,
            .name = name_tok.lexeme,
            .params = try params.toOwnedSlice(self.allocator),
            .return_type = ret_type,
            .decorators = try decorators.toOwnedSlice(self.allocator),
            .body = body,
        };
    }

    fn startsType(self: *Parser) bool {
        return switch (self.current_token.tag) {
            .type_int, .type_float, .type_str, .type_bool, .l_bracket, .question, .identifier, .at_directive => true,
            else => false,
        };
    }

    /// Parameter-type parser used by anonymous fn literals. Identical to
    /// `parseDataType`, except that a `@vector(N, T)` annotation records its
    /// leading integer N into `out_width` as a SIMD lane hint.
    fn parseParamTypeCapturingWidth(self: *Parser, out_width: *u64) ParseError!DataType {
        if (self.current_token.tag == .at_directive and
            std.mem.eql(u8, self.current_token.directiveName(), "vector"))
        {
            const dir_tok = self.current_token;
            self.advance();
            if (self.match(.l_paren)) {
                var depth: usize = 1;
                while (depth > 0 and self.current_token.tag != .eof) {
                    if (out_width.* == 0 and self.current_token.tag == .int_literal) {
                        out_width.* = std.fmt.parseInt(u64, self.current_token.lexeme, 10) catch 0;
                    }
                    if (self.current_token.tag == .l_paren) depth += 1;
                    if (self.current_token.tag == .r_paren) depth -= 1;
                    self.advance();
                }
            }
            return .{ .custom = dir_tok.lexeme };
        }
        return self.parseDataType();
    }

    fn parseDataType(self: *Parser) ParseError!DataType {
        if (self.current_token.tag == .at_directive) {
            // Directive types such as `@vector(4, f32)` are opaque to the
            // Phase 1 compiler; consume a balanced parenthesized argument
            // list and treat the result as a custom type.
            const dir_tok = self.current_token;
            self.advance();
            if (self.match(.l_paren)) {
                var depth: usize = 1;
                while (depth > 0 and self.current_token.tag != .eof) {
                    if (self.current_token.tag == .l_paren) depth += 1;
                    if (self.current_token.tag == .r_paren) depth -= 1;
                    self.advance();
                }
            }
            return .{ .custom = dir_tok.lexeme };
        }
        if (self.match(.question)) {
            const inner = try self.allocator.create(DataType);
            inner.* = try self.parseDataType();
            return .{ .optional = inner };
        }
        if (self.match(.l_bracket)) {
            // `[]T`, `[*]T`, `[N]T` — all treated as slices in Phase 1.
            while (self.current_token.tag != .r_bracket and self.current_token.tag != .eof) {
                self.advance();
            }
            _ = try self.expect(.r_bracket);
            const inner = try self.allocator.create(DataType);
            inner.* = try self.parseDataType();
            return .{ .slice = inner };
        }
        if (self.match(.type_int)) return .{ .primitive = .int };
        if (self.match(.type_float)) return .{ .primitive = .float };
        if (self.match(.type_str)) return .{ .primitive = .str };
        if (self.match(.type_bool)) return .{ .primitive = .bool };
        if (self.current_token.tag == .identifier) {
            const name = self.current_token.lexeme;
            self.advance();
            return .{ .custom = name };
        }
        return .{ .primitive = .inferred };
    }

    fn parseBlock(self: *Parser) ParseError!Block {
        _ = try self.expect(.l_brace);
        var stmts: std.ArrayList(Stmt) = .empty;
        var lines: std.ArrayList(u32) = .empty;
        var cols: std.ArrayList(u32) = .empty;
        while (self.current_token.tag != .r_brace and self.current_token.tag != .eof) {
            try lines.append(self.allocator, @intCast(self.current_token.line));
            try cols.append(self.allocator, @intCast(self.current_token.column));
            try stmts.append(self.allocator, try self.parseStmt());
        }
        _ = try self.expect(.r_brace);
        return .{
            .stmts = try stmts.toOwnedSlice(self.allocator),
            .lines = try lines.toOwnedSlice(self.allocator),
            .cols = try cols.toOwnedSlice(self.allocator),
        };
    }

    fn parseStmt(self: *Parser) ParseError!Stmt {
        if (self.current_token.tag == .kw_fn) {
            return .{ .fn_decl = try self.parseFnDecl() };
        }

        if (self.current_token.tag == .kw_const or self.current_token.tag == .kw_var) {
            const is_const = (self.current_token.tag == .kw_const);
            self.advance();
            const name_tok = try self.expect(.identifier);
            var type_ann: ?DataType = null;
            if (self.match(.colon)) {
                type_ann = try self.parseDataType();
            }
            _ = try self.expect(.assign);
            const init_expr = try self.parseOr();
            _ = self.match(.semicolon);
            return .{ .var_decl = .{ .is_const = is_const, .name = name_tok.lexeme, .type_ann = type_ann, .initializer = init_expr } };
        }

        // @defer — deferred expression execution (Phase 2 will wire real
        // resource semantics; Phase 1 executes defers at scope exit).
        if (self.current_token.tag == .at_directive and std.mem.eql(u8, self.current_token.directiveName(), "defer")) {
            self.advance();
            const expr = try self.parseOr();
            _ = self.match(.semicolon);
            return .{ .defer_stmt = expr };
        }
        if (self.current_token.tag == .kw_defer) {
            self.advance();
            const expr = try self.parseOr();
            _ = self.match(.semicolon);
            return .{ .defer_stmt = expr };
        }

        if (self.match(.kw_if)) {
            return try self.parseIfTail();
        }

        if (self.match(.kw_while)) {
            const cond = try self.parseOr();
            const body = try self.parseBlock();
            return .{ .while_stmt = .{ .condition = cond, .body = body } };
        }

        // for NAME | ITER { BODY }  — sugar over an index loop; `continue`
        // still advances the iteration. Also supports the range forms
        // `for i in lo..hi` and `for i in lo..=hi`.
        if (self.match(.kw_for)) {
            const name_tok = self.expect(.identifier) catch return error.UnexpectedToken;
            const name = self.allocator.dupe(u8, name_tok.lexeme) catch return error.OutOfMemory;
            if (self.match(.kw_in)) {
                const lo = try self.parseOr();
                const inclusive = if (self.match(.dot_dot_equal))
                    true
                else if (self.match(.dot_dot))
                    false
                else
                    return error.UnexpectedToken;
                const hi = try self.parseOr();
                const body = try self.parseBlock();
                return .{ .range_for = .{ .var_name = name, .lo = lo, .hi = hi, .hi_inclusive = inclusive, .body = body } };
            }
            _ = self.expect(.pipe) catch return error.UnexpectedToken;
            const iter = try self.parseOr();
            const body = try self.parseBlock();
            return .{ .for_stmt = .{ .var_name = name, .iter = iter, .body = body } };
        }

        if (self.match(.kw_break)) {
            _ = self.match(.semicolon);
            return .break_stmt;
        }

        if (self.match(.kw_continue)) {
            _ = self.match(.semicolon);
            return .continue_stmt;
        }

        if (self.match(.kw_return)) {
            var val: ?Expr = null;
            if (self.current_token.tag != .semicolon and self.current_token.tag != .r_brace) {
                val = try self.parseOr();
            }
            _ = self.match(.semicolon);
            return .{ .return_stmt = val };
        }

        const expr = try self.parseOr();
        if (self.match(.assign)) {
            switch (expr) {
                .identifier, .index_access => {},
                else => return error.InvalidAssignmentTarget,
            }
            const val = try self.parseOr();
            _ = self.match(.semicolon);
            return .{ .assignment = .{ .target = expr, .value = val } };
        }
        _ = self.match(.semicolon);
        return .{ .expr_stmt = expr };
    }

    fn parseIfTail(self: *Parser) ParseError!Stmt {
        const cond = try self.parseOr();
        const then_branch = try self.parseBlock();

        var else_branch: ?Block = null;
        if (self.match(.kw_else)) {
            if (self.match(.kw_if)) {
                // `else if` chains recurse into an inner if statement.
                const inner_line: u32 = @intCast(self.current_token.line);
                const inner_col: u32 = @intCast(self.current_token.column);
                const inner = try self.parseIfTail();
                var one_stmts: std.ArrayList(Stmt) = .empty;
                var one_lines: std.ArrayList(u32) = .empty;
                var one_cols: std.ArrayList(u32) = .empty;
                try one_stmts.append(self.allocator, inner);
                try one_lines.append(self.allocator, inner_line);
                try one_cols.append(self.allocator, inner_col);
                else_branch = .{
                    .stmts = try one_stmts.toOwnedSlice(self.allocator),
                    .lines = try one_lines.toOwnedSlice(self.allocator),
                    .cols = try one_cols.toOwnedSlice(self.allocator),
                };
            } else {
                else_branch = try self.parseBlock();
            }
        }

        return .{ .if_stmt = .{ .condition = cond, .then_branch = then_branch, .else_branch = else_branch } };
    }

    pub const ParseError = error{ UnexpectedToken, InvalidAssignmentTarget, OutOfMemory, Overflow, InvalidCharacter, InvalidExpression };

    pub fn parseExpr(self: *Parser) ParseError!Expr {
        return try self.parseOr();
    }

    // ---- Expression precedence chain ----
    // or -> and -> equality -> comparison -> additive -> multiplicative -> unary -> postfix

    fn parseOr(self: *Parser) ParseError!Expr {
        var lhs = try self.parseAnd();
        while (self.current_token.tag == .kw_or) {
            self.advance();
            const rhs = try self.parseAnd();
            lhs = try self.makeBinary(.kw_or, lhs, rhs);
        }
        return lhs;
    }

    fn parseAnd(self: *Parser) ParseError!Expr {
        var lhs = try self.parseEquality();
        while (self.current_token.tag == .kw_and) {
            self.advance();
            const rhs = try self.parseEquality();
            lhs = try self.makeBinary(.kw_and, lhs, rhs);
        }
        return lhs;
    }

    fn parseEquality(self: *Parser) ParseError!Expr {
        var lhs = try self.parseComparison();
        while (self.current_token.tag == .equal_equal or self.current_token.tag == .not_equal) {
            const op = self.current_token.tag;
            self.advance();
            const rhs = try self.parseComparison();
            lhs = try self.makeBinary(op, lhs, rhs);
        }
        return lhs;
    }

    fn parseComparison(self: *Parser) ParseError!Expr {
        var lhs = try self.parseAdditive();
        while (self.current_token.tag == .less or self.current_token.tag == .less_equal or self.current_token.tag == .greater or self.current_token.tag == .greater_equal) {
            const op = self.current_token.tag;
            self.advance();
            const rhs = try self.parseAdditive();
            lhs = try self.makeBinary(op, lhs, rhs);
        }
        return lhs;
    }

    fn parseAdditive(self: *Parser) ParseError!Expr {
        var lhs = try self.parseMultiplicative();
        while (self.current_token.tag == .plus or self.current_token.tag == .minus) {
            const op = self.current_token.tag;
            self.advance();
            const rhs = try self.parseMultiplicative();
            lhs = try self.makeBinary(op, lhs, rhs);
        }
        return lhs;
    }

    fn parseMultiplicative(self: *Parser) ParseError!Expr {
        var lhs = try self.parseUnary();
        while (self.current_token.tag == .star or self.current_token.tag == .slash or self.current_token.tag == .percent) {
            const op = self.current_token.tag;
            self.advance();
            const rhs = try self.parseUnary();
            lhs = try self.makeBinary(op, lhs, rhs);
        }
        return lhs;
    }

    fn makeBinary(self: *Parser, op: TokenType, lhs: Expr, rhs: Expr) ParseError!Expr {
        const lhs_box = try self.allocator.create(Expr);
        lhs_box.* = lhs;
        const rhs_box = try self.allocator.create(Expr);
        rhs_box.* = rhs;
        return .{ .binary = .{ .op = op, .lhs = lhs_box, .rhs = rhs_box } };
    }

    fn makeUnary(self: *Parser, op: TokenType, operand: Expr) ParseError!Expr {
        const operand_box = try self.allocator.create(Expr);
        operand_box.* = operand;
        return .{ .unary = .{ .op = op, .operand = operand_box } };
    }

    fn parseUnary(self: *Parser) ParseError!Expr {
        // `not` keyword and `!` both compile to logical not.
        if (self.match(.kw_not)) {
            const operand = try self.parseUnary();
            return try self.makeUnaryNot(operand);
        }
        if (self.match(.bang)) {
            const operand = try self.parseUnary();
            return try self.makeUnaryNot(operand);
        }
        if (self.match(.minus)) {
            const operand = try self.parseUnary();
            return try self.makeUnary(.minus, operand);
        }
        return try self.parsePostfix();
    }

    fn makeUnaryNot(self: *Parser, operand: Expr) ParseError!Expr {
        return try self.makeUnary(.kw_not, operand);
    }

    fn parsePostfix(self: *Parser) ParseError!Expr {
        var primary = try self.parsePrimary();

        while (true) {
            if (self.match(.period)) {
                const field_tok = try self.expect(.identifier);
                const obj_box = try self.allocator.create(Expr);
                obj_box.* = primary;
                primary = .{ .field_access = .{ .object = obj_box, .field = field_tok.lexeme } };
            } else if (self.match(.l_paren)) {
                var args: std.ArrayList(Expr) = .empty;
                while (self.current_token.tag != .r_paren and self.current_token.tag != .eof) {
                    try args.append(self.allocator, try self.parseOr());
                    if (!self.match(.comma)) break;
                }
                _ = try self.expect(.r_paren);
                const callee_box = try self.allocator.create(Expr);
                callee_box.* = primary;
                primary = .{ .call = .{ .callee = callee_box, .args = try args.toOwnedSlice(self.allocator) } };
            } else if (self.match(.l_bracket)) {
                const idx = try self.parseOr();
                _ = try self.expect(.r_bracket);
                const obj_box = try self.allocator.create(Expr);
                obj_box.* = primary;
                const idx_box = try self.allocator.create(Expr);
                idx_box.* = idx;
                primary = .{ .index_access = .{ .object = obj_box, .index = idx_box } };
            } else {
                break;
            }
        }

        return primary;
    }

    fn unescapeString(self: *Parser, raw: []const u8) ParseError![]const u8 {
        // raw includes surrounding quotes. The result is always an owned
        // copy so Expr.deinit can free it uniformly.
        if (raw.len < 2) return raw;
        const body = raw[1 .. raw.len - 1];
        if (std.mem.indexOfScalar(u8, body, '\\') == null) {
            return self.allocator.dupe(u8, body) catch return error.OutOfMemory;
        }

        var out: std.ArrayList(u8) = .empty;
        var i: usize = 0;
        while (i < body.len) : (i += 1) {
            if (body[i] == '\\' and i + 1 < body.len) {
                i += 1;
                const esc: u8 = switch (body[i]) {
                    'n' => '\n',
                    't' => '\t',
                    'r' => '\r',
                    '0' => 0,
                    '"' => '"',
                    '\\' => '\\',
                    else => body[i],
                };
                out.append(self.allocator, esc) catch return error.OutOfMemory;
            } else {
                out.append(self.allocator, body[i]) catch return error.OutOfMemory;
            }
        }
        return out.toOwnedSlice(self.allocator) catch return error.OutOfMemory;
    }

    fn parsePrimary(self: *Parser) ParseError!Expr {
        var primary: Expr = undefined;

        if (self.current_token.tag == .int_literal) {
            const val = try std.fmt.parseInt(i64, self.current_token.lexeme, 10);
            self.advance();
            primary = .{ .int_lit = val };
        } else if (self.current_token.tag == .float_literal) {
            const val = try std.fmt.parseFloat(f64, self.current_token.lexeme);
            self.advance();
            primary = .{ .float_lit = val };
        } else if (self.current_token.tag == .string_literal) {
            const lex = self.current_token.lexeme;
            self.advance();
            primary = .{ .string_lit = try self.unescapeString(lex) };
        } else if (self.current_token.tag == .kw_true) {
            self.advance();
            primary = .{ .bool_lit = true };
        } else if (self.current_token.tag == .kw_false) {
            self.advance();
            primary = .{ .bool_lit = false };
        } else if (self.current_token.tag == .kw_null) {
            self.advance();
            primary = .null_lit;
        } else if (self.current_token.tag == .at_directive) {
            // Open-ended directive call: `@stdout("hi")`, `@env("HOME")`,
            // `@tensor(f32, ...)`. The directive name is resolved by the
            // compiler/backend, keeping the lexer O(1) and future-proof.
            const name = self.current_token.directiveName();
            self.advance();
            var args: std.ArrayList(Expr) = .empty;
            if (self.match(.l_paren)) {
                while (self.current_token.tag != .r_paren and self.current_token.tag != .eof) {
                    try args.append(self.allocator, try self.parseOr());
                    if (!self.match(.comma)) break;
                }
                _ = try self.expect(.r_paren);
            }
            primary = .{ .call_directive = .{ .name = name, .args = try args.toOwnedSlice(self.allocator) } };
        } else if (self.current_token.tag == .kw_fn) {
            // Anonymous fn literal: `fn(x: f32) f32 { ... }`. The compiler
            // restricts these to apply_simd kernel position.
            self.advance();
            _ = try self.expect(.l_paren);
            var params: std.ArrayList(FnParam) = .empty;
            var vector_width: u64 = 0;
            errdefer {
                for (params.items) |*p| p.param_type.deinit(self.allocator);
                params.deinit(self.allocator);
            }
            while (self.current_token.tag != .r_paren and self.current_token.tag != .eof) {
                const p_name = try self.expect(.identifier);
                _ = try self.expect(.colon);
                const p_type = try self.parseParamTypeCapturingWidth(&vector_width);
                try params.append(self.allocator, .{ .name = p_name.lexeme, .param_type = p_type });
                if (!self.match(.comma)) break;
            }
            _ = try self.expect(.r_paren);

            // Optional return type, same rules as named fn declarations.
            if (self.match(.colon)) {
                _ = try self.parseDataType();
            } else if (self.startsType()) {
                _ = try self.parseDataType();
            }

            const body = try self.parseBlock();
            const fl = try self.allocator.create(FnLit);
            fl.* = .{ .params = try params.toOwnedSlice(self.allocator), .body = body, .vector_width = vector_width };
            primary = .{ .fn_lit = fl };
        } else if (self.current_token.tag == .identifier or
            self.current_token.tag == .type_int or
            self.current_token.tag == .type_float or
            self.current_token.tag == .type_str or
            self.current_token.tag == .type_bool)
        {
            // Primitive type names double as conversion builtins when they
            // appear in expression position, e.g. `int("42")`.
            const id = self.current_token.lexeme;
            self.advance();
            primary = .{ .identifier = id };
        } else if (self.match(.period)) {
            if (self.current_token.tag == .identifier) {
                // Enum-literal shorthand: `.cpu`, `.f32`. The lexer emits
                // float literals for `.5`-style digit starts, so there is no
                // ambiguity with numeric literals here.
                const name = self.current_token.lexeme;
                self.advance();
                primary = .{ .enum_lit = name };
            } else {
                // `.{...}` bundle: plain items form a tuple/array literal,
                // `.name = value` entries form a struct literal (option
                // bundles like `@tensor(..., .{ .device = .cpu })`).
                _ = try self.expect(.l_brace);
                var items: std.ArrayList(Expr) = .empty;
                var fields: std.ArrayList(Expr.StructFieldInit) = .empty;
                errdefer {
                    for (items.items) |*it| it.deinit(self.allocator);
                    items.deinit(self.allocator);
                    for (fields.items) |*f| f.deinit(self.allocator);
                    fields.deinit(self.allocator);
                }
                while (self.current_token.tag != .r_brace and self.current_token.tag != .eof) {
                    if (self.current_token.tag == .period) {
                        self.advance();
                        const f_name = try self.expect(.identifier);
                        _ = try self.expect(.assign);
                        try fields.append(self.allocator, .{
                            .name = f_name.lexeme,
                            .explicit_type = null,
                            .value = try self.parseOr(),
                            .key_expr = null,
                        });
                    } else {
                        try items.append(self.allocator, try self.parseOr());
                    }
                    if (!self.match(.comma)) break;
                }
                _ = try self.expect(.r_brace);
                if (fields.items.len > 0) {
                    if (items.items.len > 0) return error.UnexpectedToken;
                    primary = .{ .struct_lit = try fields.toOwnedSlice(self.allocator) };
                } else {
                    primary = .{ .array_lit = try items.toOwnedSlice(self.allocator) };
                }
            }
        } else if (self.match(.l_bracket)) {
            var items: std.ArrayList(Expr) = .empty;
            while (self.current_token.tag != .r_bracket and self.current_token.tag != .eof) {
                try items.append(self.allocator, try self.parseOr());
                if (!self.match(.comma)) break;
            }
            _ = try self.expect(.r_bracket);
            primary = .{ .array_lit = try items.toOwnedSlice(self.allocator) };
        } else if (self.match(.l_paren)) {
            primary = try self.parseOr();
            _ = try self.expect(.r_paren);
        } else if (self.match(.l_brace)) {
            primary = try self.parseStructLitBody();
        } else {
            return error.InvalidExpression;
        }

        return primary;
    }

    fn parseStructLitBody(self: *Parser) ParseError!Expr {
        var fields: std.ArrayList(Expr.StructFieldInit) = .empty;
        while (self.current_token.tag != .r_brace and self.current_token.tag != .eof) {
            if (self.match(.period)) {
                const f_name = try self.expect(.identifier);
                var explicit_t: ?DataType = null;
                if (self.match(.colon)) {
                    explicit_t = try self.parseDataType();
                }
                _ = try self.expect(.assign);
                try fields.append(self.allocator, .{ .name = f_name.lexeme, .explicit_type = explicit_t, .value = try self.parseOr(), .key_expr = null });
            } else {
                const key = try self.parseOr();
                if (self.match(.colon)) {
                    const val = try self.parseOr();
                    try fields.append(self.allocator, .{ .name = "", .explicit_type = null, .value = val, .key_expr = key });
                } else {
                    key.deinit(self.allocator);
                    return error.UnexpectedToken;
                }
            }
            if (!self.match(.comma)) break;
        }
        _ = try self.expect(.r_brace);
        return .{ .struct_lit = try fields.toOwnedSlice(self.allocator) };
    }
};

// ---- Tests ----

test "Parser parses basic script" {
    const source =
        \\const tab_size: int = 4;
        \\pub fn main(args: []str) int {
        \\    @stdout(tab_size);
        \\}
    ;
    var p = Parser.init(std.testing.allocator, source);
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 1), prog.fn_decls.len);
    try std.testing.expectEqualSlices(u8, "main", prog.fn_decls[0].name);
    try std.testing.expect(prog.fn_decls[0].return_type == .primitive);
    try std.testing.expectEqual(DataType.Primitive.int, prog.fn_decls[0].return_type.primitive);
}

test "Parser handles nested fn declarations" {
    var p = Parser.init(std.testing.allocator,
        \\fn helper() {
        \\    @stdout("toggle");
        \\}
        \\
        \\pub fn main(args: []str) {
        \\    helper();
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 2), prog.fn_decls.len);
    try std.testing.expectEqualStrings("helper", prog.fn_decls[0].name);
}

test "Parser handles map literal entries" {
    var p = Parser.init(std.testing.allocator,
        \\pub fn main() {
        \\    bind_group("normal", {
        \\        key("<leader>w"): ":write",
        \\        key("Ctrl-h"): "window-left",
        \\    });
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 1), prog.fn_decls.len);
}

test "Parser handles @import without leading dot" {
    var p = Parser.init(std.testing.allocator,
        \\@import = { std, ai };
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 2), prog.directives.len);
}

test "Parser precedence: additive vs multiplicative" {
    var p = Parser.init(std.testing.allocator, "1 + 2 * 3");
    const e = try p.parseExpr();
    defer e.deinit(std.testing.allocator);
    try std.testing.expect(e == .binary);
    try std.testing.expectEqual(TokenType.plus, e.binary.op);
    try std.testing.expect(e.binary.rhs.* == .binary);
    try std.testing.expectEqual(TokenType.star, e.binary.rhs.binary.op);
}

test "Parser comparison and equality chains" {
    var p = Parser.init(std.testing.allocator, "a + 1 < b * 2 == true");
    const e = try p.parseExpr();
    defer e.deinit(std.testing.allocator);
    try std.testing.expect(e == .binary);
    try std.testing.expectEqual(TokenType.equal_equal, e.binary.op);
}

test "Parser index access and assignment" {
    var p = Parser.init(std.testing.allocator,
        \\fn f() {
        \\    xs[0] = xs[1] + 2;
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    const body = prog.fn_decls[0].body.stmts;
    try std.testing.expect(body[0] == .assignment);
    try std.testing.expect(body[0].assignment.target == .index_access);
}

test "Parser while loop with break and continue" {
    var p = Parser.init(std.testing.allocator,
        \\fn f() {
        \\    while x < 10 {
        \\        if x == 3 { break; }
        \\        continue;
        \\    }
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    const body = prog.fn_decls[0].body.stmts;
    try std.testing.expect(body[0] == .while_stmt);
    try std.testing.expect(body[0].while_stmt.body.stmts.len == 2);
}

test "Parser else-if chain" {
    var p = Parser.init(std.testing.allocator,
        \\fn f(x: int) {
        \\    if x == 1 { @stdout("one"); }
        \\    else if x == 2 { @stdout("two"); }
        \\    else { @stdout("many"); }
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    const s = prog.fn_decls[0].body.stmts[0];
    try std.testing.expect(s == .if_stmt);
    const els = s.if_stmt.else_branch.?;
    try std.testing.expectEqual(@as(usize, 1), els.stmts.len);
    try std.testing.expect(els.stmts[0] == .if_stmt);
}

test "Parser generic directive statement" {
    var p = Parser.init(std.testing.allocator,
        \\@fetch_pkg("https://example.com/pkg.rg", "1.0.0");
        \\fn f() {}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 1), prog.directives.len);
    try std.testing.expect(prog.directives[0] == .generic);
    try std.testing.expectEqualStrings("fetch_pkg", prog.directives[0].generic.name);
    try std.testing.expectEqual(@as(usize, 2), prog.directives[0].generic.args.len);
}

test "Parser decorator attaches to fn" {
    var p = Parser.init(std.testing.allocator,
        \\@export("fib")
        \\pub fn fib(n: int) int {
        \\    return n;
        \\}
    );
    const prog = try p.parseProgram();
    defer prog.deinit(std.testing.allocator);
    try std.testing.expectEqual(@as(usize, 1), prog.fn_decls.len);
    try std.testing.expectEqual(@as(usize, 1), prog.fn_decls[0].decorators.len);
    try std.testing.expectEqualStrings("export", prog.fn_decls[0].decorators[0].name);
}

test "Parser string escapes" {
    var p = Parser.init(std.testing.allocator, "\"a\\nb\\tc\"");
    const e = try p.parseExpr();
    defer e.deinit(std.testing.allocator);
    try std.testing.expectEqualStrings("a\nb\tc", e.string_lit);
}

test "Parser parenthesized expression" {
    var p = Parser.init(std.testing.allocator, "(1 + 2) * 3");
    const e = try p.parseExpr();
    defer e.deinit(std.testing.allocator);
    try std.testing.expect(e == .binary);
    try std.testing.expectEqual(TokenType.star, e.binary.op);
    try std.testing.expectEqual(@as(i64, 1), e.binary.lhs.binary.lhs.int_lit);
}

test "Parser diagnostics record failing position" {
    var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena_state.deinit();
    const arena = arena_state.allocator();

    // Line 3, col 9: `end` instead of a closing brace.
    var p = Parser.init(arena,
        \\fn ok() int {
        \\    return 1;
        \\}
        \\
        \\fn broken() int { end
    );
    try std.testing.expectError(error.UnexpectedToken, p.parseProgram());
    try std.testing.expectEqual(@as(usize, 5), p.err_line);
    try std.testing.expect(p.err_msg != null);
    try std.testing.expect(std.mem.indexOf(u8, p.err_msg.?, "line 5:") != null);
    try std.testing.expect(std.mem.indexOf(u8, p.err_msg.?, "expected 'r_brace'") != null);
}

test "Block records per-statement source lines" {
    var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena_state.deinit();
    const arena = arena_state.allocator();

    var p = Parser.init(arena,
        \\const a: int = 1;
        \\fn f() void {
        \\    @stdout("x\n", .{});
        \\    const b: int = 2;
        \\}
    );
    const prog = try p.parseProgram();

    try std.testing.expectEqual(@as(u32, 1), prog.global_stmts.lines[0]);
    try std.testing.expectEqual(@as(usize, 2), prog.fn_decls[0].body.stmts.len);
    try std.testing.expectEqual(@as(u32, 3), prog.fn_decls[0].body.lines[0]);
    try std.testing.expectEqual(@as(u32, 4), prog.fn_decls[0].body.lines[1]);
}