"""
Nexa Language — Recursive Descent Parser v25
New: var, fn-lambda, native assert, trait/impl, use, union types, /= assign, Rust-style errors
"""
from __future__ import annotations
from typing import Any, List, Optional

from .lexer import Token, TT, Lexer
from .ast_nodes import *


class ParseError(Exception):
    pass

_IDENT_LIKE = (
    TT.IDENT, TT.AUTOML, TT.TRAIN, TT.PREDICT,
    TT.MODEL, TT.LAYER, TT.ASYNC, TT.AWAIT, TT.SELF,
    TT.MATCH, TT.CASE, TT.DEFAULT, TT.TEST, TT.ASSERT,
    TT.YIELD, TT.STRUCT, TT.ENUM,
    TT.TRAIT, TT.IMPL, TT.USE, TT.VAR, TT.FN,   # v25
    TT.NEWTYPE, TT.TEMPLATE, TT.SPAWN,            # v27: promoted from IDENT
)


class Parser:
    def __init__(self, tokens: List[Token], source: str = ""):
        self.tokens = [t for t in tokens if t.type not in (TT.NEWLINE, TT.SEMICOLON)]
        self.pos = 0
        self._source_lines = source.splitlines() if source else []

    def _peek(self) -> Token:       return self.tokens[self.pos]
    def _advance(self) -> Token:
        t = self.tokens[self.pos]; self.pos += 1; return t
    def _check(self, *types: TT) -> bool: return self._peek().type in types
    def _match(self, *types: TT) -> Optional[Token]:
        if self._check(*types): return self._advance()
    def _expect(self, tt: TT, msg: str = "") -> Token:
        if not self._check(tt):
            tok = self._peek()
            line_no = tok.line
            src_line = self._source_lines[line_no - 1] if 0 < line_no <= len(self._source_lines) else ""
            caret = "    ^" + " " * max(0, len(src_line.rstrip()) - 1)
            hint = f"\n    | {src_line.rstrip()}\n{caret}" if src_line else ""
            raise ParseError(
                f"\n  Parse Error at line {line_no}: expected {tt.name}, got {tok.type.name} ({tok.value!r}){hint}\n  {msg}"
            )
        return self._advance()
    def _expect_ident(self) -> Token:
        tok = self._peek()
        if tok.type in _IDENT_LIKE or tok.type == TT.IDENT:
            return self._advance()
        line_no = tok.line
        src_line = self._source_lines[line_no - 1] if 0 < line_no <= len(self._source_lines) else ""
        hint = f"\n    | {src_line.rstrip()}\n    ^" if src_line else ""
        raise ParseError(f"\n  Parse Error at line {line_no}: expected identifier, got {tok.type.name}{hint}")

    def parse(self) -> Program:
        stmts = []
        while not self._check(TT.EOF):
            stmts.append(self._statement())
        return Program(stmts)

    # ── Statements ────────────────────────────────────────────────────────
    def _statement(self) -> Any:
        line = self._peek().line if self.pos < len(self.tokens) else -1
        decs = self._decorators()
        if self._check(TT.CLASS):   node = self._class_def(decs)
        elif self._check(TT.STRUCT):
            if decs: raise ParseError("Decorators not supported on structs.")
            node = self._struct_def()
        elif self._check(TT.ASYNC):   node = self._async_func_def(decs)
        elif self._check(TT.DEF):     node = self._func_def(decs)
        elif self._check(TT.ENUM):
            if decs: raise ParseError("Decorators not supported on enums.")
            node = self._enum_def()
        elif self._check(TT.INTERFACE):
            if decs: raise ParseError("Decorators not supported on interfaces.")
            node = self._interface_def()
        elif decs: raise ParseError("Decorators only allowed on def or class.")

        elif self._check(TT.LET):     node = self._let_stmt()
        elif self._check(TT.VAR):     node = self._var_stmt()           # v25
        elif self._check(TT.CONST):   node = self._const_stmt()
        elif self._check(TT.RETURN):  node = self._return_stmt()
        elif self._check(TT.YIELD):   node = self._yield_stmt()
        elif self._check(TT.BREAK):   node = self._break_stmt()
        elif self._check(TT.CONTINUE): node = self._continue_stmt()
        # v24: if let / while let check via lookahead
        elif self._check(TT.IF) and self._v24_is_if_let():   node = self._if_let_stmt()
        elif self._check(TT.WHILE) and self._v24_is_while_let(): node = self._while_let_stmt()
        elif self._check(TT.WHILE):   node = self._while_stmt()
        elif self._check(TT.FOR):     node = self._for_stmt()
        elif self._check(TT.LOOP):    node = self._loop_stmt()
        elif self._check(TT.PRIVATE, TT.PUBLIC): node = self._func_def()
        elif self._check(TT.ABSTRACT): node = self._abstract_stmt()
        elif self._check(TT.EXPORT):  node = self._export_stmt()
        elif self._check(TT.IMPORT):  node = self._import_stmt()
        elif self._check(TT.FROM):    node = self._from_import_stmt()
        elif self._check(TT.USE):     node = self._use_stmt()           # v25
        elif self._check(TT.MODEL):   node = self._model_def()
        elif self._check(TT.TRAIN):   node = self._train_stmt()
        elif self._check(TT.TEST):    node = self._test_block()
        elif self._check(TT.BENCH):   node = self._bench_block()
        elif self._check(TT.ASSERT) and self._v25_is_native_assert(): node = self._native_assert_stmt()  # v25
        elif self._check(TT.ASSERT):  node = self._assert_stmt()
        # v25 new statements
        elif self._check(TT.TRAIT):   node = self._trait_def()          # v25
        elif self._check(TT.IMPL):    node = self._trait_impl()          # v25
        # v24 new statements
        elif self._check(TT.SEALED):  node = self._sealed_class_def()
        elif self._check(TT.MACRO):   node = self._macro_def()
        elif self._check(TT.TYPE):    node = self._type_alias()
        # v27 new statements
        elif self._check(TT.NEWTYPE): node = self._newtype_stmt()
        elif self._check(TT.IDENT) and self._v24_is_macro_call(): node = ExprStatement(self._macro_call())
        else: node = ExprStatement(self._expr())

        try:
            if not hasattr(node, 'line') or node.line == 0:
                node.line = line
        except AttributeError:
            pass
        return node

    def _loop_stmt(self) -> LoopStatement:
        self._expect(TT.LOOP)
        return LoopStatement(self._block())

    def _newtype_stmt(self) -> 'NewtypeStatement':
        """newtype Celsius = Float  — distinct branded type (v27)."""
        self._expect(TT.NEWTYPE)
        name = self._expect_ident().value
        self._expect(TT.ASSIGN)
        base_type = self._expect_ident().value
        return NewtypeStatement(name, base_type)

    def _const_stmt(self) -> ConstStatement:
        """const PI = 3.14  — immutable binding"""
        self._expect(TT.CONST)
        name = self._expect_ident().value
        type_hint = None
        if self._match(TT.COLON):
            type_hint = self._advance().value if self._check(TT.IDENT) else None
        self._expect(TT.ASSIGN)
        return ConstStatement(name, self._expr(), type_hint)

    def _break_stmt(self):
        # Forward to final definition below (which handles break if expr)
        # NOTE: This is overridden by the definition at _while_stmt section.
        pass  # placeholder — actual impl below


    def _enum_def(self) -> EnumDef:
        """enum Option<T> { Some(T)  None  Error(str) }"""
        self._expect(TT.ENUM)
        name = self._expect_ident().value
        type_params = self._parse_type_params()
        self._expect(TT.LBRACE)
        variants = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            vname = self._expect_ident().value
            fields = []
            if self._check(TT.LPAREN):
                self._advance()
                while not self._check(TT.RPAREN) and not self._check(TT.EOF):
                    fields.append(self._expect_ident().value)
                    self._match(TT.COMMA)
                self._expect(TT.RPAREN)
            variants.append(EnumVariant(vname, fields))
            self._match(TT.COMMA)  # optional comma
        self._expect(TT.RBRACE)
        return EnumDef(name, type_params, variants)

    def _abstract_stmt(self):
        """abstract class Foo { } or abstract def method()"""
        self._expect(TT.ABSTRACT)
        if self._check(TT.CLASS):
            node = self._class_def()
            node.is_abstract = True
            return node
        elif self._check(TT.DEF):
            self._expect(TT.DEF)
            name = self._expect_ident().value
            self._expect(TT.LPAREN)
            params = []
            while not self._check(TT.RPAREN) and not self._check(TT.EOF):
                params.append(self._expect_ident().value)
                self._match(TT.COMMA)
            self._expect(TT.RPAREN)
            if self._check(TT.LBRACE):
                self._block()
            return AbstractMethodDef(name, params)
        raise ParseError("Expected 'class' or 'def' after 'abstract'")

    def _export_stmt(self) -> ExportStatement:
        self._expect(TT.EXPORT)
        self._expect(TT.LBRACE)
        names = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            names.append(self._expect_ident().value)
            self._match(TT.COMMA)
        self._expect(TT.RBRACE)
        return ExportStatement(names)

    def _parse_type_params(self) -> List[str]:
        """Parse <T, U, V> generic type params after a name."""
        type_params = []
        if self._check(TT.LT):
            self._advance()  # consume <
            while not self._check(TT.GT) and not self._check(TT.EOF):
                type_params.append(self._expect_ident().value)
                self._match(TT.COMMA)
            self._expect(TT.GT)
        return type_params

    # ── Struct ────────────────────────────────────────────────────────────
    def _struct_def(self) -> StructDef:
        self._expect(TT.STRUCT)
        name = self._expect_ident().value
        self._expect(TT.LBRACE)
        fields = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            field_name = self._expect_ident().value
            type_hint = None
            if self._match(TT.COLON):
                type_hint = self._expect_ident().value
            fields.append((field_name, type_hint))
            self._match(TT.COMMA)
        self._expect(TT.RBRACE)
        return StructDef(name, fields)

    # ── Decorators and Enums ──────────────────────────────────────────────
    def _decorators(self) -> List[Decorator]:
        decs = []
        while self._match(TT.AT):
            decs.append(Decorator(self._call_or_access()))
        return decs

    # NOTE: _enum_def is defined above (v23 version with variants+fields)
    # This stub is kept only as a placeholder to avoid class layout changes

    # ── Definitions ───────────────────────────────────────────────────────
    def _class_def(self, decorators=None) -> ClassDef:
        self._expect(TT.CLASS)
        name = self._expect_ident().value
        base = None
        implements = []
        if self._match(TT.LPAREN):
            base = self._expect_ident().value; self._expect(TT.RPAREN)
        elif self._match(TT.COLON) or self._match(TT.LT):
            base = self._expect_ident().value
        if self._match(TT.IMPLEMENTS):
            implements.append(self._expect_ident().value)
            while self._match(TT.COMMA):
                implements.append(self._expect_ident().value)
        self._expect(TT.LBRACE)
        body = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            decs = self._decorators()
            if self._check(TT.DEF):     body.append(self._func_def(decs))
            elif self._check(TT.ASYNC): body.append(self._async_func_def(decs))
            elif self._check(TT.PRIVATE, TT.PUBLIC):
                vis = self._advance().value
                fn = self._func_def(decs)
                fn.visibility = vis
                body.append(fn)
            elif self._check(TT.ABSTRACT):
                body.append(self._abstract_stmt())
            else:
                if decs: raise ParseError("Decorators only allow on methods inside class")
                body.append(self._statement())
        self._expect(TT.RBRACE)
        node = ClassDef(name, body, base, decorators or [], implements)
        node.type_params = self._parse_type_params.__doc__ and [] or []
        return node

    def _interface_def(self) -> InterfaceDef:
        self._expect(TT.INTERFACE)
        name = self._expect_ident().value
        self._expect(TT.LBRACE)
        method_names = []
        defaults = {}
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            self._match(TT.DEF)  # optional keyword
            m_name = self._expect_ident().value
            method_names.append(m_name)
            # consume optional param signature
            params = []
            if self._check(TT.LPAREN):
                self._expect(TT.LPAREN)
                params, _, _ = self._typed_param_list()
                self._expect(TT.RPAREN)
            # consume optional return type like `-> String`
            if self._match(TT.ARROW):
                self._expect_ident()
            # consume optional default method body like `{ ... }`
            if self._check(TT.LBRACE):
                defaults[m_name] = FunctionDef(m_name, params, self._block())
        self._expect(TT.RBRACE)
        return InterfaceDef(name, method_names, defaults)

    # ── Functions with docstrings ──────────────────────────────────────────
    def _async_func_def(self, decorators=None) -> FunctionDef:
        self._expect(TT.ASYNC)
        f = self._func_def(decorators)
        f.is_async = True
        return f

    def _func_def(self, decorators=None) -> FunctionDef:
        self._expect(TT.DEF)
        name = self._expect_ident().value
        type_params = self._parse_type_params()  # v22: <T, U>
        self._expect(TT.LPAREN)
        params = []; param_types = {}; defaults = {}
        while not self._check(TT.RPAREN): # Changed from TT.EOF to TT.RPAREN for parameter parsing
            p_name = self._expect_ident().value
            params.append(p_name)
            if self._match(TT.COLON):
                param_types[p_name] = self._expect_ident().value
            if self._match(TT.ASSIGN):
                defaults[p_name] = self._expr()
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        ret_type = self._optional_return_type()
        body, docstring = self._block_with_docstring()
        return FunctionDef(name, params, body, False, param_types, ret_type, defaults, docstring, decorators or [])

    def _typed_param_list(self):
        params, param_types, defaults = [], {}, {}
        while not self._check(TT.RPAREN):
            name_tok = self._expect_ident()
            name = str(name_tok.value)
            if self._match(TT.COLON):
                param_types[name] = self._expect_ident().value
            if self._match(TT.ASSIGN):
                defaults[name] = self._expr()
            params.append(name)
            if not self._match(TT.COMMA): break
        return params, param_types, defaults

    def _optional_return_type(self) -> Optional[str]:
        if self._check(TT.ARROW):
            self._advance(); return self._expect_ident().value

    def _block_with_docstring(self):
        """Parse a {...} block, extracting a leading triple-string docstring."""
        self._expect(TT.LBRACE)
        docstring = None
        if self._check(TT.TRIPLE_STRING):
            docstring = self._advance().value
        stmts = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            stmts.append(self._statement())
        self._expect(TT.RBRACE)
        return stmts, docstring

    def _block(self) -> List[Any]:
        body, _ = self._block_with_docstring()
        return body

    # ── yield ───────────────────────────────────────────────────────
    # NOTE: v0.2 — duplicate _yield_stmt removed. Single definition at line ~402.

    # ── Other statements ──────────────────────────────────────────────────
    def _let_stmt(self) -> Any:
        self._expect(TT.LET)
        if self._check(TT.LBRACKET):
            self._advance()
            targets = []
            while not self._check(TT.RBRACKET):
                targets.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RBRACKET); self._expect(TT.ASSIGN)
            return DestructuredAssignment(targets, self._expr(), is_dict=False)
        if self._check(TT.LBRACE):
            self._advance()
            targets = []
            while not self._check(TT.RBRACE):
                targets.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RBRACE); self._expect(TT.ASSIGN)
            return DestructuredAssignment(targets, self._expr(), is_dict=True)
            
        name = self._expect_ident().value
        hint = None
        if self._match(TT.COLON): hint = self._expect_ident().value
        self._expect(TT.ASSIGN)
        line = self._peek().line if self.pos < len(self.tokens) else 0
        # v0.2: debug print removed (was: print("DEBUG IN PARSER let_stmt:", line))
        expr = self._expr()
        return LetStatement(name, expr, hint, line=line)


    def _return_stmt(self) -> ReturnStatement:
        tok = self._expect(TT.RETURN)
        value = self._expr() if not self._check(TT.RBRACE) else NullLiteral()
        return ReturnStatement(value, line=tok.line)

    def _yield_stmt(self) -> YieldStatement:
        """yield expr  — single canonical definition (v0.2: duplicate removed)."""
        self._expect(TT.YIELD)
        value = self._expr() if not self._check(TT.RBRACE) else NullLiteral()
        return YieldStatement(value)

    def _if_expr(self) -> IfExpr:
        self._expect(TT.IF)
        has_paren = self._match(TT.LPAREN) is not None
        cond = self._expr()
        if has_paren: self._expect(TT.RPAREN)
        then_b = self._block(); else_b = []
        if self._match(TT.ELSE):
            else_b = [ExprStatement(self._if_expr())] if self._check(TT.IF) else self._block()
        return IfExpr(cond, then_b, else_b)

    def _for_stmt(self) -> ForStatement:
        self._expect(TT.FOR)
        has_paren = self._match(TT.LPAREN) is not None
        var = self._expect_ident().value; self._expect(TT.IN)
        iterable = self._expr()
        if has_paren: self._expect(TT.RPAREN)
        return ForStatement(var, iterable, self._block())

    def _while_stmt(self) -> WhileStatement:
        self._expect(TT.WHILE)
        has_paren = self._match(TT.LPAREN) is not None
        cond = self._expr()
        if has_paren: self._expect(TT.RPAREN)
        return WhileStatement(cond, self._block())

    def _break_stmt(self):
        """break  or  break if expr  — conditional inline break (v23)."""
        self._expect(TT.BREAK)
        # Peek at next token: if it's IF, parse as conditional break
        if self.pos < len(self.tokens) and self.tokens[self.pos].type == TT.IF:
            self._advance()  # consume 'if'
            cond = self._expr()
            return BreakIfStatement(cond)
        return BreakStatement()

    def _continue_stmt(self) -> ContinueStatement:
        self._expect(TT.CONTINUE)
        return ContinueStatement()

    def _import_stmt(self) -> ImportStatement:
        """v0.2: deprecated — prefer 'use module' or 'use module::{item}'."""
        import warnings
        warnings.warn(
            "Deprecation: 'import X' is deprecated. Use 'use X' instead. "
            "This syntax will be removed in Nexa v0.4.",
            DeprecationWarning, stacklevel=2
        )
        self._expect(TT.IMPORT)
        if self._check(TT.STRING):
            path = self._advance().value; alias = None
            if self._match(TT.AS): alias = self._expect_ident().value
            return ImportStatement(path, alias, is_file=True)
        module = self._expect_ident().value; alias = None
        if self._match(TT.AS): alias = self._expect_ident().value
        return ImportStatement(module, alias, is_file=False)

    def _from_import_stmt(self) -> FromImportStatement:
        """v0.2: deprecated — prefer 'use module::{item1, item2}'."""
        import warnings
        warnings.warn(
            "Deprecation: 'from X import {Y}' is deprecated. "
            "Use 'use X::{Y}' instead. Removed in Nexa v0.4.",
            DeprecationWarning, stacklevel=2
        )
        self._expect(TT.FROM)
        if self._check(TT.STRING):
            module = self._advance().value
        else:
            module = self._expect_ident().value
        self._expect(TT.IMPORT)
        self._expect(TT.LBRACE)
        names = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            names.append(self._expect_ident().value)
            self._match(TT.COMMA)
        self._expect(TT.RBRACE)
        return FromImportStatement(module, names)

    def _model_def(self) -> ModelDef:
        self._expect(TT.MODEL)
        name = self._expect_ident().value; self._expect(TT.LBRACE)
        layers, config = [], {}
        while not self._check(TT.RBRACE):
            key = self._expect_ident().value; self._expect(TT.COLON)
            val = self._expr()
            if key == "layers": layers = val
            else: config[key] = val
            self._match(TT.COMMA)
        self._expect(TT.RBRACE)
        return ModelDef(name, layers, config)

    def _train_stmt(self) -> TrainStatement:
        self._expect(TT.TRAIN)
        model_name = self._expect_ident().value
        if self._check(TT.IDENT) and self._peek().value == "on":
            self._advance()
        data = self._expr()
        config = self._expr()
        return TrainStatement(model_name, data, config)

    def _try_expr(self) -> TryExpr:
        self._expect(TT.TRY)
        try_block = self._block()
        catch_var, catch_block, finally_block = None, [], []
        if self._match(TT.CATCH):
            self._expect(TT.LPAREN); catch_var = self._expect_ident().value
            self._expect(TT.RPAREN); catch_block = self._block()
        if self._match(TT.FINALLY): finally_block = self._block()
        return TryExpr(try_block, catch_var, catch_block, finally_block)

    def _throw_stmt(self) -> ThrowStatement:
        self._expect(TT.THROW); return ThrowStatement(self._expr())

    def _match_expr(self) -> MatchExpr:
        self._expect(TT.MATCH); subject = self._expr()
        self._expect(TT.LBRACE)
        cases = []; default_body = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            if self._match(TT.DEFAULT):
                self._expect(TT.COLON); default_body = self._case_body()
            elif self._match(TT.CASE):
                pattern = self._parse_case_pattern()
                self._expect(TT.COLON)
                cases.append(CaseClause(pattern, self._case_body()))
        self._expect(TT.RBRACE)
        return MatchExpr(subject, cases, default_body)

    def _parse_case_pattern(self):
        """Parse a match case pattern — expr, {x,y}, [h,...t], or Cls(x)."""
        # Dict destructure: case {x, y}  or  case {"x", "y"}
        if self._check(TT.LBRACE):
            self._advance()
            fields = []
            while not self._check(TT.RBRACE) and not self._check(TT.EOF):
                if self._check(TT.STRING):
                    # String key — use the string value as the bind name
                    fields.append(self._advance().value)
                else:
                    fields.append(self._expect_ident().value)
                self._match(TT.COMMA)
            self._expect(TT.RBRACE)
            return DestructurePattern('dict', fields)
        # List destructure: case [head, ...tail]
        if self._check(TT.LBRACKET):
            self._advance()
            fields = []
            while not self._check(TT.RBRACKET) and not self._check(TT.EOF):
                if self._match(TT.ELLIPSIS):
                    fields.append('...' + self._expect_ident().value)
                else:
                    fields.append(self._expect_ident().value)
                self._match(TT.COMMA)
            self._expect(TT.RBRACKET)
            return DestructurePattern('list', fields)
        # Enum variant pattern: case Option.Some(val) or case Option.None
        if (self._check(TT.IDENT) and (self.pos + 1) < len(self.tokens)
                and self.tokens[self.pos+1].type == TT.DOT):
            # Could be Enum.Variant or Enum.Variant(binds)
            saved_pos = self.pos
            enum_name = self._advance().value     # "Option"
            self._advance()                        # "."
            if self._check(TT.IDENT):
                variant_name = self._advance().value  # "Some" or "None"
                binds = []
                if self._check(TT.LPAREN):
                    self._advance()
                    while not self._check(TT.RPAREN) and not self._check(TT.EOF):
                        binds.append(self._expect_ident().value)
                        self._match(TT.COMMA)
                    self._expect(TT.RPAREN)
                return EnumVariantPattern(enum_name, variant_name, binds)
            else:
                self.pos = saved_pos  # rollback, not an enum pattern
        # v24: CLASS destructure with field names and typed match pattern
        # Handle: case String s  or  case String s if guard:
        if self._check(TT.IDENT):
            next_is_ident = (self.pos + 1 < len(self.tokens) and
                             self.tokens[self.pos + 1].type == TT.IDENT)
            if next_is_ident:
                # Could be typed match: type_name bind_name [if guard]
                saved_pos = self.pos
                type_name = self._advance().value   # e.g. "String"
                bind_name = self._advance().value   # e.g. "s"
                guard = None
                if self._match(TT.IF):
                    guard = self._expr()
                return TypedMatchCase(type_name, bind_name, guard, [])
        # Class destructure: case Point(x, y)
        if (self._check(TT.IDENT) and (self.pos + 1) < len(self.tokens)
                and self.tokens[self.pos+1].type == TT.LPAREN):
            cls_name = self._advance().value
            self._expect(TT.LPAREN)
            fields = []
            while not self._check(TT.RPAREN) and not self._check(TT.EOF):
                fields.append(self._expect_ident().value)
                self._match(TT.COMMA)
            self._expect(TT.RPAREN)
            return DestructurePattern('class', fields, cls_name)
        # Default: literal / expression pattern
        return self._expr()

    def _case_body(self) -> List[Any]:
        stmts = []
        while not self._check(TT.CASE, TT.DEFAULT, TT.RBRACE, TT.EOF):
            stmts.append(self._statement())
        return stmts

    def _test_block(self) -> TestBlock:
        self._expect(TT.TEST); name = self._expect(TT.STRING).value
        return TestBlock(name, self._block())

    def _bench_block(self):
        from nexa_lang.ast_nodes import BenchBlock
        self._expect(TT.BENCH); name = self._expect(TT.STRING).value
        return BenchBlock(name, self._block())

    def _assert_stmt(self) -> AssertStatement:
        self._expect(TT.ASSERT)
        func_name = self._expect_ident().value
        kind_map = {"assert_eq": "eq", "assert_ne": "ne", "assert_true": "true",
                    "assert_false": "false", "assert_raises": "raises",
                    "assert_contains": "contains"}
        kind = kind_map.get(func_name, "true")
        self._expect(TT.LPAREN); args = []
        while not self._check(TT.RPAREN):
            args.append(self._expr())
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        return AssertStatement(kind, args)

    # ── Expressions ───────────────────────────────────────────────────────
    def _expr(self) -> Any:
        return self._pipeline()

    def _pipeline(self) -> Any:
        left = self._assignment()
        if self._check(TT.PIPE):
            stages = [left]
            while self._match(TT.PIPE): stages.append(self._assignment())
            return PipelineExpr(stages)
        return left

    def _assignment(self) -> Any:
        left = self._ternary()
        if self._check(TT.PLUS_EQ, TT.MINUS_EQ, TT.STAR_EQ, TT.SLASH_EQ,
                        TT.AMPERSAND_EQ, TT.VBAR_EQ, TT.CARET_EQ, TT.LSHIFT_EQ, TT.RSHIFT_EQ) and isinstance(left, Identifier):
            op_tok = self._advance()
            op_map = {TT.PLUS_EQ: "+", TT.MINUS_EQ: "-", TT.STAR_EQ: "*", TT.SLASH_EQ: "/",
                      TT.AMPERSAND_EQ: "&", TT.VBAR_EQ: "|", TT.CARET_EQ: "^",
                      TT.LSHIFT_EQ: "<<", TT.RSHIFT_EQ: ">>"}
            return CompoundAssignment(left.name, op_map[op_tok.type], self._assignment())
        if self._check(TT.ASSIGN):
            self._advance(); value = self._assignment()
            if isinstance(left, Identifier):   return Assignment(left.name, value)
            if isinstance(left, MemberAccess): return MemberAssignment(left.obj, left.member, value)
            if isinstance(left, IndexAccess):  return IndexAssignment(left.obj, left.index, value)
        return left

    def _ternary(self) -> Any:
        left = self._nullish()
        if self._match(TT.QUESTION):
            true_expr = self._expr()
            self._expect(TT.COLON)
            false_expr = self._ternary()
            return TernaryOp(left, true_expr, false_expr)
        return left

    def _nullish(self) -> Any:
        left = self._or_expr()
        while self._match(TT.QUESTION_QUESTION):
            left = NullishCoalescing(left, self._or_expr())
        return left

    def _or_expr(self) -> Any:
        left = self._and_expr()
        while self._match(TT.OR): left = BinaryOp("or", left, self._and_expr())
        return left

    def _and_expr(self) -> Any:
        left = self._equality()
        while self._match(TT.AND): left = BinaryOp("and", left, self._equality())
        return left

    def _equality(self) -> Any:
        left = self._comparison()
        while self._check(TT.EQ, TT.NEQ):
            op = self._advance().type
            left = BinaryOp("==" if op == TT.EQ else "!=", left, self._comparison())
        return left

    def _comparison(self) -> Any:
        left = self._bitwise_or()
        MAP = {TT.LT: "<", TT.GT: ">", TT.LTE: "<=", TT.GTE: ">="}
        while self._check(*MAP):
            left = BinaryOp(MAP[self._advance().type], left, self._bitwise_or())
        return left

    def _bitwise_or(self) -> Any:
        left = self._bitwise_xor()
        while self._match(TT.VBAR): left = BinaryOp("|", left, self._bitwise_xor())
        return left

    def _bitwise_xor(self) -> Any:
        left = self._bitwise_and()
        while self._match(TT.CARET): left = BinaryOp("^", left, self._bitwise_and())
        return left

    def _bitwise_and(self) -> Any:
        left = self._shift()
        while self._match(TT.AMPERSAND): left = BinaryOp("&", left, self._shift())
        return left

    def _shift(self) -> Any:
        left = self._addition()
        while self._check(TT.LSHIFT, TT.RSHIFT):
            op = "<<" if self._advance().type == TT.LSHIFT else ">>"
            left = BinaryOp(op, left, self._addition())
        return left

    def _addition(self) -> Any:
        left = self._multiplication()
        while self._check(TT.PLUS, TT.MINUS):
            op = self._advance().value
            left = BinaryOp(op, left, self._multiplication())
        # Range literals: 0..10  and  0..=10
        if self._check(TT.DOTDOT, TT.DOTDOT_EQ):
            inclusive = self._advance().type == TT.DOTDOT_EQ
            return RangeExpr(left, self._addition(), inclusive)
        return left

    def _multiplication(self) -> Any:
        left = self._power()
        while self._check(TT.STAR, TT.SLASH, TT.MODULO):
            op = self._advance().value
            left = BinaryOp(op, left, self._power())
        return left

    def _power(self) -> Any:
        base = self._unary()
        if self._match(TT.POW): return BinaryOp("**", base, self._power())
        return base

    def _unary(self) -> Any:
        if self._match(TT.NOT):   return UnaryOp("not", self._unary())
        if self._match(TT.MINUS): return UnaryOp("-", self._unary())
        if self._match(TT.TILDE): return UnaryOp("~", self._unary())
        if self._match(TT.AWAIT): return AwaitExpr(self._unary())
        return self._call_or_access()

    def _call_or_access(self) -> Any:
        node = self._primary()
        while True:
            if self._match(TT.DOT):
                node = MemberAccess(node, str(self._expect_ident().value))
            elif self._match(TT.QUESTION_DOT):
                node = OptionalAccess(node, str(self._expect_ident().value))
            elif self._check(TT.LPAREN):
                node = self._finish_call(node)
            elif self._check(TT.LBRACKET):
                self._advance()
                start = None
                if not self._check(TT.COLON) and not self._check(TT.RBRACKET):
                    start = self._expr()
                
                if self._match(TT.COLON):
                    stop = None
                    if not self._check(TT.COLON) and not self._check(TT.RBRACKET):
                        stop = self._expr()
                    
                    step = None
                    if self._match(TT.COLON):
                        if not self._check(TT.RBRACKET):
                            step = self._expr()
                    
                    self._expect(TT.RBRACKET)
                    node = IndexAccess(node, SliceExpr(start, stop, step))
                else:
                    self._expect(TT.RBRACKET)
                    node = IndexAccess(node, start)
            else: break
        # v24: postfix ? operator — wrap in QuestionExpr for Err/None propagation
        if self._match(TT.QUESTION):
            node = QuestionExpr(node)
        return node

    def _finish_call(self, callee: Any) -> FunctionCall:
        self._expect(TT.LPAREN)
        args, kwargs = [], {}
        while not self._check(TT.RPAREN):
            if self._check(TT.IDENT) and self.pos + 1 < len(self.tokens) \
                    and self.tokens[self.pos + 1].type == TT.ASSIGN:
                key = self._advance().value; self._advance()
                kwargs[key] = self._expr()
            else: args.append(self._expr())
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        return FunctionCall(callee, args, kwargs)

    def _primary(self) -> Any:
        tok = self._peek()
        if tok.type == TT.NUMBER:        self._advance(); return NumberLiteral(tok.value)
        if tok.type == TT.STRING:        self._advance(); return StringLiteral(tok.value)
        if tok.type == TT.TRIPLE_STRING: self._advance(); return StringLiteral(tok.value)
        if tok.type == TT.FSTRING:       self._advance(); return self._parse_fstring(tok.value)
        if tok.type == TT.BOOL:          self._advance(); return BoolLiteral(tok.value)
        if tok.type == TT.NULL:          self._advance(); return NullLiteral()
        if tok.type == TT.SELF:          self._advance(); return SelfRef()
        if tok.type == TT.SUPER:
            self._advance()
            if self._match(TT.LPAREN): self._expect(TT.RPAREN)
            return SuperRef()
        if tok.type == TT.NEW:           return self._new_expr()
        if tok.type == TT.PREDICT:
            self._advance()
            model_name = self._expect_ident().value
            self._expect(TT.LPAREN)
            input_data = self._expr()
            self._expect(TT.RPAREN)
            return PredictExpr(model_name, input_data)
        if tok.type == TT.LBRACKET:      return self._list_or_comp()
        if tok.type == TT.DEF:
            self._advance()
            self._expect(TT.LPAREN)
            params = []
            while not self._check(TT.RPAREN):
                if self._match(TT.ELLIPSIS):
                    params.append("..." + self._expect_ident().value)
                    break
                params.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RPAREN)
            if self._match(TT.FAT_ARROW):
                body = [ReturnStatement(self._expr())]
            else:
                body = self._block()
            return LambdaExpr(params, body)
            
        # v25: Expression-oriented blocks and control flow
        if tok.type == TT.LBRACE:        return self._dict_or_set_or_block()
        if tok.type == TT.IF:            return self._if_expr()
        if tok.type == TT.MATCH:         return self._match_expr()
        if tok.type == TT.TRY:           return self._try_expr()
        
        if tok.type in _IDENT_LIKE or tok.type == TT.IDENT:
            self._advance(); return Identifier(str(tok.value))
        if tok.type == TT.LPAREN:
            self._advance(); expr = self._expr(); self._expect(TT.RPAREN); return expr
        if self._match(TT.ELLIPSIS):     return SpreadExpr(self._expr())
        raise ParseError(f"Unexpected {tok.type.name} ({tok.value!r}) L{tok.line}")

    def _list_or_comp(self) -> Any:
        self._expect(TT.LBRACKET)
        if self._check(TT.RBRACKET): self._advance(); return ListLiteral([])
        first = self._expr()
        if self._match(TT.FOR):
            var = self._expect_ident().value; self._expect(TT.IN)
            iterable = self._expr(); condition = None
            if self._match(TT.IF): condition = self._expr()
            self._expect(TT.RBRACKET)
            return ListComp(first, var, iterable, condition)
        elems = [first]
        while self._match(TT.COMMA) and not self._check(TT.RBRACKET):
            elems.append(self._expr())
        self._expect(TT.RBRACKET)
        return ListLiteral(elems)

    def _parse_fstring(self, parts: list) -> FString:
        ast_parts = []
        for kind, val in parts:
            if kind == "lit": ast_parts.append(("lit", val))
            else:
                tokens = Lexer(val).tokenize()
                ast_parts.append(("expr", Parser(tokens)._expr()))
        return FString(ast_parts)

    def _new_expr(self) -> NewExpr:
        self._expect(TT.NEW); class_name = self._expect_ident().value
        self._expect(TT.LPAREN); args, kwargs = [], {}
        while not self._check(TT.RPAREN):
            if self._check(TT.IDENT) and self.pos+1 < len(self.tokens) \
                    and self.tokens[self.pos+1].type == TT.ASSIGN:
                key = self._advance().value; self._advance()
                kwargs[key] = self._expr()
            else: args.append(self._expr())
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        return NewExpr(class_name, args, kwargs)

    def _dict_or_set_or_block(self) -> Any:
        # A BlockExpr if it contains statements that aren't dict/set syntax
        # We peek to decide if it's a dict, set, or block.
        # Simple heuristic: if the first token is an expression followed by a comma or colon, or it's empty, dict/set
        # NOTE: For now, we reuse the dict/set logic but if it fails, we fall back to a block.
        # A simpler way: peek ahead.
        self._expect(TT.LBRACE)
        if self._check(TT.RBRACE):
            self._advance(); return DictLiteral([])
            
        # Peek at token after next to distinguish dict/set vs block
        # For simplicity, if we see a LET, CONST, IF, WHILE, FOR, RETURN etc. it's a block.
        first_type = self._peek().type
        if first_type in (TT.LET, TT.CONST, TT.IF, TT.WHILE, TT.FOR, TT.RETURN, TT.BREAK, TT.CONTINUE):
            # It's a block!
            stmts = []
            while not self._check(TT.RBRACE) and not self._check(TT.EOF):
                stmts.append(self._statement())
            self._expect(TT.RBRACE)
            return BlockExpr(stmts)

        # Fallback to dict/set literal parsing
        elements = []
        has_colon = False
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            expr = self._expr()
            if isinstance(expr, SpreadExpr):
                elements.append((expr, None))
            elif self._match(TT.COLON):
                elements.append((expr, self._expr()))
                has_colon = True
            else:
                elements.append((expr, None))
            # If no comma is present and it's not the end, it might be a block after all!
            if not self._check(TT.COMMA) and not self._check(TT.RBRACE):
                # Backtrack logic would be needed here for true expressions-as-statements, 
                # but Nexa semantics require braces for Dicts/Sets to use commas. 
                # So if no comma, treat rest as a block... but we already consumed self._expr().
                # For this pass, we'll enforce that blocks without control-flow words start with a clear expr.
                break
            self._match(TT.COMMA)
            
        self._expect(TT.RBRACE)
        
        if has_colon or all(isinstance(k, SpreadExpr) for k, v in elements):
            return DictLiteral([(k, v) for k, v in elements])
        else:
            return SetLiteral([k for k, v in elements])

    # ── v24 Lookahead helpers ─────────────────────────────────────────────────

    def _v24_is_if_let(self) -> bool:
        """Return True when current token is IF and next is LET."""
        return (self.pos + 1 < len(self.tokens) and
                self.tokens[self.pos + 1].type == TT.LET)

    def _v24_is_while_let(self) -> bool:
        """Return True when current token is WHILE and next is LET."""
        return (self.pos + 1 < len(self.tokens) and
                self.tokens[self.pos + 1].type == TT.LET)

    def _v24_is_macro_call(self) -> bool:
        """Return True when we have IDENT followed by BANG and LPAREN — macro call."""
        return (self.pos + 1 < len(self.tokens) and
                self.tokens[self.pos + 1].type == TT.BANG)

    # ── v24: if let / while let ───────────────────────────────────────────────

    def _if_let_stmt(self) -> IfLetStatement:
        """if let Pattern = expr { then } else { else }"""
        self._expect(TT.IF)
        self._expect(TT.LET)
        pattern = self._parse_case_pattern()
        self._expect(TT.ASSIGN)
        expr = self._expr()
        then_body = self._block()
        else_body = []
        if self._match(TT.ELSE):
            else_body = self._block()
        return IfLetStatement(pattern, expr, then_body, else_body)

    def _while_let_stmt(self) -> WhileLetStatement:
        """while let Pattern = expr { body }"""
        self._expect(TT.WHILE)
        self._expect(TT.LET)
        pattern = self._parse_case_pattern()
        self._expect(TT.ASSIGN)
        expr = self._expr()
        body = self._block()
        return WhileLetStatement(pattern, expr, body)

    # ── v24: Sealed Classes ───────────────────────────────────────────────────

    def _sealed_class_def(self) -> SealedClassDef:
        """sealed class Expr { class Num(value: Int)  class Add(left, right) }"""
        self._expect(TT.SEALED)
        self._expect(TT.CLASS)
        name = self._expect_ident().value
        self._expect(TT.LBRACE)
        variants = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            self._expect(TT.CLASS)
            vname = self._expect_ident().value
            fields = []
            if self._match(TT.LPAREN):
                while not self._check(TT.RPAREN) and not self._check(TT.EOF):
                    fname = self._expect_ident().value
                    ftype = None
                    if self._match(TT.COLON):
                        ftype = self._expect_ident().value
                    fields.append((fname, ftype))
                    if not self._match(TT.COMMA): break
                self._expect(TT.RPAREN)
            variants.append(SealedVariant(vname, fields))
        self._expect(TT.RBRACE)
        return SealedClassDef(name, variants)

    # ── v24: Macros ───────────────────────────────────────────────────────────

    def _macro_def(self) -> MacroDef:
        """macro name!(a, b) { body... }"""
        self._expect(TT.MACRO)
        name = self._expect_ident().value
        self._expect(TT.BANG)
        self._expect(TT.LPAREN)
        params = []
        while not self._check(TT.RPAREN) and not self._check(TT.EOF):
            params.append(self._expect_ident().value)
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        body = self._block()
        return MacroDef(name, params, body)

    def _macro_call(self) -> MacroCall:
        """name!(arg1, arg2)"""
        name = self._expect_ident().value
        self._expect(TT.BANG)
        self._expect(TT.LPAREN)
        args = []
        while not self._check(TT.RPAREN) and not self._check(TT.EOF):
            args.append(self._expr())
            if not self._match(TT.COMMA): break
        self._expect(TT.RPAREN)
        return MacroCall(name, args)

    # ── v24: Type Aliases ─────────────────────────────────────────────────────

    def _type_alias(self) -> TypeAlias:
        """type UserId = Int"""
        self._expect(TT.TYPE)
        name = self._expect_ident().value
        self._expect(TT.ASSIGN)
        # Type target can be any expression, not just an identifier
        target = self._expr()
        return TypeAlias(name, target)

    # ── v25: New Statements ───────────────────────────────────────────────────

    def _var_stmt(self) -> Any:
        """var x: Int = 5  — explicit mutable binding (semantically same as let for now)."""
        self._expect(TT.VAR)
        if self._check(TT.LBRACKET):
            # var [a, b] = [1, 2]  — array destructure
            self._advance()
            targets = []
            while not self._check(TT.RBRACKET):
                targets.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RBRACKET); self._expect(TT.ASSIGN)
            return DestructuredAssignment(targets, self._expr(), is_dict=False)
        name = self._expect_ident().value
        hint = self._parse_type_annotation()
        self._expect(TT.ASSIGN)
        return VarStatement(name, self._expr(), hint)

    def _parse_type_annotation(self) -> Optional[Any]:
        """Parse optional ': TypeName' or ': Int | String' union type."""
        if not self._match(TT.COLON):
            return None
        first = self._expect_ident().value
        
        # Check for nullable shorthand: String?
        if self._match(TT.QUESTION):
            return UnionType([first, "None"])
            
        # Check for union: Int | String | None
        if self._check(TT.VBAR):
            types = [first]
            while self._match(TT.VBAR):
                types.append(self._expect_ident().value)
            return UnionType(types)
        return first

    def _v25_is_native_assert(self) -> bool:
        """Lookahead: ASSERT followed by non-IDENT (to distinguish from old assert assert_eq(..))."""
        if self.pos + 1 >= len(self.tokens):
            return False
        next_tok = self.tokens[self.pos + 1]
        # If next token after ASSERT is not an identifier that looks like assert_eq/assert_true
        val = next_tok.value if next_tok.type == TT.IDENT else ""
        if isinstance(val, str) and val.startswith("assert_"):
            return False
        return True

    def _native_assert_stmt(self) -> NativeAssertStatement:
        """assert expr  or  assert expr, "message" """
        self._expect(TT.ASSERT)
        cond = self._expr()
        msg = None
        if self._match(TT.COMMA):
            msg = self._expr()
        return NativeAssertStatement(cond, msg)

    def _trait_def(self) -> TraitDef:
        """trait Printable { def print(self) -> String; def describe(self) { ... } }"""
        self._expect(TT.TRAIT)
        name = self._expect_ident().value
        self._expect(TT.LBRACE)
        methods = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            if self._check(TT.DEF):
                methods.append(self._func_def())
            elif self._check(TT.ASYNC):
                methods.append(self._async_func_def())
            else:
                self._advance()  # skip unknown tokens gracefully
        self._expect(TT.RBRACE)
        return TraitDef(name, methods)

    def _trait_impl(self) -> TraitImpl:
        """impl Printable for Dog { def print(self) { return "woof" } }"""
        self._expect(TT.IMPL)
        trait_name = self._expect_ident().value
        # Expect 'for' keyword (parsed as IDENT since 'for' is FOR token)
        if self._check(TT.FOR):
            self._advance()
        else:
            self._expect_ident()  # allow bare word
        class_name = self._expect_ident().value
        self._expect(TT.LBRACE)
        methods = []
        while not self._check(TT.RBRACE) and not self._check(TT.EOF):
            decs = self._decorators()
            if self._check(TT.DEF):
                methods.append(self._func_def(decs))
            elif self._check(TT.ASYNC):
                methods.append(self._async_func_def(decs))
            else:
                self._advance()
        self._expect(TT.RBRACE)
        return TraitImpl(trait_name, class_name, methods)

    def _use_stmt(self) -> UseStatement:
        """use std::math::sin  or  use std::math::{sin, cos} as m"""
        self._expect(TT.USE)
        path = [self._expect_ident().value]
        while self.pos < len(self.tokens) and self.tokens[self.pos].type == TT.COLON:
            # consume :: (two colons)
            self._advance()  # first :
            if self.pos < len(self.tokens) and self.tokens[self.pos].type == TT.COLON:
                self._advance()  # second :
            if self._check(TT.LBRACE):
                # use std::math::{sin, cos}
                self._advance()
                names = []
                while not self._check(TT.RBRACE) and not self._check(TT.EOF):
                    names.append(self._expect_ident().value)
                    self._match(TT.COMMA)
                self._expect(TT.RBRACE)
                alias = None
                if self._match(TT.AS):
                    alias = self._expect_ident().value
                return UseStatement(path, names, alias)
            else:
                path.append(self._expect_ident().value)
        alias = None
        if self._match(TT.AS):
            alias = self._expect_ident().value
        return UseStatement(path, [], alias)

    # Override _let_stmt to support union type annotations
    def _let_stmt(self) -> Any:
        self._expect(TT.LET)
        if self._check(TT.LBRACKET):
            self._advance()
            targets = []
            while not self._check(TT.RBRACKET):
                targets.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RBRACKET); self._expect(TT.ASSIGN)
            return DestructuredAssignment(targets, self._expr(), is_dict=False)
        if self._check(TT.LBRACE):
            self._advance()
            targets = []
            while not self._check(TT.RBRACE):
                targets.append(self._expect_ident().value)
                if not self._match(TT.COMMA): break
            self._expect(TT.RBRACE); self._expect(TT.ASSIGN)
            return DestructuredAssignment(targets, self._expr(), is_dict=True)
        name = self._expect_ident().value
        hint = self._parse_type_annotation()       # v25: union-aware
        self._expect(TT.ASSIGN)
        return LetStatement(name, self._expr(), hint if isinstance(hint, str) else None)

