robinpath 0.2.0

use crate::ast::*;
use crate::error::RobinPathError;
use crate::lexer::{Token, TokenKind};
use crate::stream::TokenStream;
use crate::value::AttributePathSegment;

pub struct Parser {
    stream: TokenStream,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Self {
        Self {
            stream: TokenStream::new(tokens),
        }
    }

    pub fn parse(&mut self) -> Result<Vec<Stmt>, RobinPathError> {
        let mut stmts = Vec::new();
        self.stream.skip_newlines();

        while !self.stream.is_at_end() {
            if self.stream.check(TokenKind::Newline) {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Comment) {
                let tok = self.stream.current().unwrap().clone();
                let text = tok
                    .string_value()
                    .unwrap_or(&tok.text[1..])
                    .trim()
                    .to_string();
                let span = Span::from_line(tok.line, tok.column);
                stmts.push(Stmt::Comment { text, span });
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::ChunkMarker) {
                let tok = self.stream.current().unwrap().clone();
                let id = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                stmts.push(Stmt::ChunkMarker { id, span });
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::CellOpen) {
                let tok = self.stream.current().unwrap().clone();
                let cell_type = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                // Check if next token is raw text (for prompt/schema cells)
                let mut raw_text = None;
                let mut content = Vec::new();
                if self.stream.check(TokenKind::RawText) {
                    let raw_tok = self.stream.current().unwrap().clone();
                    raw_text = Some(raw_tok.string_value().unwrap_or("").to_string());
                    self.stream.advance();
                } else {
                    // Code cell — parse body until CellClose
                    self.stream.skip_newlines();
                    while !self.stream.is_at_end()
                        && !self.stream.check(TokenKind::CellClose)
                    {
                        if self.stream.check(TokenKind::Newline) {
                            self.stream.advance();
                            continue;
                        }
                        if self.stream.check(TokenKind::Comment) {
                            let ctok = self.stream.current().unwrap().clone();
                            let text = ctok.string_value().unwrap_or(&ctok.text[1..]).trim().to_string();
                            let cspan = Span::from_line(ctok.line, ctok.column);
                            content.push(Stmt::Comment { text, span: cspan });
                            self.stream.advance();
                            continue;
                        }
                        // Handle prompt blocks inside code cells
                        if self.stream.check(TokenKind::PromptOpen) {
                            let ptok = self.stream.current().unwrap().clone();
                            let pspan = Span::from_line(ptok.line, ptok.column);
                            self.stream.advance();
                            let text = if self.stream.check(TokenKind::RawText) {
                                let raw_tok = self.stream.current().unwrap().clone();
                                let t = raw_tok.string_value().unwrap_or("").to_string();
                                self.stream.advance();
                                t
                            } else { String::new() };
                            if self.stream.check(TokenKind::PromptClose) { self.stream.advance(); }
                            content.push(Stmt::PromptBlock { text, span: pspan });
                            continue;
                        }
                        let inner = self.parse_statement()?;
                        content.push(inner);
                    }
                }
                // Consume CellClose if present
                if self.stream.check(TokenKind::CellClose) {
                    self.stream.advance();
                }
                self.stream.skip_newlines();
                stmts.push(Stmt::CellBlock { cell_type, content, raw_text, span });
                continue;
            }
            if self.stream.check(TokenKind::PromptOpen) {
                let tok = self.stream.current().unwrap().clone();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                // Next should be RawText
                let text = if self.stream.check(TokenKind::RawText) {
                    let raw_tok = self.stream.current().unwrap().clone();
                    let t = raw_tok.string_value().unwrap_or("").to_string();
                    self.stream.advance();
                    t
                } else {
                    String::new()
                };
                // Consume PromptClose if present
                if self.stream.check(TokenKind::PromptClose) {
                    self.stream.advance();
                }
                self.stream.skip_newlines();
                stmts.push(Stmt::PromptBlock { text, span });
                continue;
            }
            // Skip stray CellClose/PromptClose/RawText tokens
            if self.stream.check(TokenKind::CellClose)
                || self.stream.check(TokenKind::PromptClose)
                || self.stream.check(TokenKind::RawText)
            {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Eof) {
                break;
            }

            let stmt = self.parse_statement()?;
            stmts.push(stmt);
            self.stream.skip_newlines();
        }

        Ok(stmts)
    }

    fn parse_statement(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().ok_or_else(|| {
            RobinPathError::parse("Unexpected end of input", 0, 0)
        })?;

        match tok.kind {
            TokenKind::Keyword => {
                let kw = tok.text.clone();
                match kw.as_str() {
                    "if" => self.parse_if(),
                    "iftrue" => self.parse_iftrue(),
                    "iffalse" => self.parse_iffalse(),
                    "for" => self.parse_for_loop(),
                    "def" | "define" => self.parse_define(),
                    "do" => self.parse_scope_block(),
                    "together" => self.parse_together(),
                    "on" => self.parse_on_block(),
                    "return" => self.parse_return(),
                    "break" => self.parse_break(),
                    "continue" => self.parse_continue(),
                    "set" => self.parse_set(),
                    "var" => self.parse_var_decl(),
                    "const" => self.parse_const_decl(),
                    "use" => self.parse_use(),
                    "log" | "say" => self.parse_command_call(),
                    "repeat" => self.parse_repeat_or_command(),
                    _ => {
                        // Unknown keyword — treat as command
                        self.parse_command_call()
                    }
                }
            }
            TokenKind::Variable => {
                // Could be assignment ($var = ...) or shorthand ($var used alone)
                self.parse_variable_or_assignment()
            }
            TokenKind::Identifier => {
                // Command call: funcName arg1 arg2
                self.parse_command_call()
            }
            TokenKind::Decorator => self.parse_decorator(),
            TokenKind::SubexpressionOpen => {
                // $( ... ) as a statement
                let expr = self.parse_expression(0)?;
                let span = *expr.span();
                Ok(Stmt::Command {
                    name: "_subexpr".to_string(),
                    args: vec![Arg::Expr(expr)],
                    into: None,
                    span,
                })
            }
            _ => {
                // Try to parse as expression (bare number, string, object, array, etc.)
                match tok.kind {
                    TokenKind::Number | TokenKind::String | TokenKind::Boolean
                    | TokenKind::Null | TokenKind::LBrace | TokenKind::LBracket
                    | TokenKind::LParen | TokenKind::TemplateLiteral | TokenKind::Minus => {
                        let expr = self.parse_expression(0)?;
                        let span = *expr.span();
                        Ok(Stmt::Command {
                            name: "_expr".to_string(),
                            args: vec![Arg::Expr(expr)],
                            into: None,
                            span,
                        })
                    }
                    _ => {
                        let line = tok.line;
                        let col = tok.column;
                        let text = tok.text.clone();
                        let kind = tok.kind;
                        self.stream.advance();
                        Err(RobinPathError::parse(
                            format!("Unexpected token: {:?} '{}'", kind, text),
                            line,
                            col,
                        ))
                    }
                }
            }
        }
    }

    // ==================== Expression parsing (precedence climbing) ====================

    fn parse_expression(&mut self, min_prec: u8) -> Result<Expr, RobinPathError> {
        let mut left = self.parse_unary()?;

        while let Some(tok) = self.stream.current() {
            let op = match tok.kind {
                TokenKind::Plus => Some(BinaryOp::Add),
                TokenKind::Minus => Some(BinaryOp::Sub),
                TokenKind::Multiply => Some(BinaryOp::Mul),
                TokenKind::Divide => Some(BinaryOp::Div),
                TokenKind::Modulo => Some(BinaryOp::Mod),
                TokenKind::Eq => Some(BinaryOp::Eq),
                TokenKind::Ne => Some(BinaryOp::Ne),
                TokenKind::Lt => Some(BinaryOp::Lt),
                TokenKind::Lte => Some(BinaryOp::Lte),
                TokenKind::Gt => Some(BinaryOp::Gt),
                TokenKind::Gte => Some(BinaryOp::Gte),
                TokenKind::And => Some(BinaryOp::And),
                TokenKind::Or => Some(BinaryOp::Or),
                _ => None,
            };

            let Some(op) = op else { break };

            if op.precedence() < min_prec {
                break;
            }

            self.stream.advance(); // consume operator
            let right = self.parse_expression(op.precedence() + 1)?;
            let span = Span::new(
                left.span().start_row,
                left.span().start_col,
                right.span().end_row,
                right.span().end_col,
            );
            left = Expr::Binary {
                op,
                left: Box::new(left),
                right: Box::new(right),
                span,
            };
        }

        Ok(left)
    }

    fn parse_unary(&mut self) -> Result<Expr, RobinPathError> {
        if let Some(tok) = self.stream.current() {
            match tok.kind {
                TokenKind::Not => {
                    let span_start = Span::from_line(tok.line, tok.column);
                    self.stream.advance();
                    let arg = self.parse_unary()?;
                    let span = Span::new(
                        span_start.start_row,
                        span_start.start_col,
                        arg.span().end_row,
                        arg.span().end_col,
                    );
                    return Ok(Expr::Unary {
                        op: UnaryOp::Not,
                        argument: Box::new(arg),
                        span,
                    });
                }
                TokenKind::Minus => {
                    // Could be unary negation
                    let span_start = Span::from_line(tok.line, tok.column);
                    self.stream.advance();
                    let arg = self.parse_primary()?;
                    let span = Span::new(
                        span_start.start_row,
                        span_start.start_col,
                        arg.span().end_row,
                        arg.span().end_col,
                    );
                    return Ok(Expr::Unary {
                        op: UnaryOp::Neg,
                        argument: Box::new(arg),
                        span,
                    });
                }
                _ => {}
            }
        }
        self.parse_primary()
    }

    fn parse_primary(&mut self) -> Result<Expr, RobinPathError> {
        let tok = self.stream.current().ok_or_else(|| {
            RobinPathError::parse("Unexpected end of input in expression", 0, 0)
        })?;
        let span = Span::from_line(tok.line, tok.column);

        match tok.kind {
            TokenKind::Number => {
                let n = tok.number_value().unwrap_or(0.0);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Number(n),
                    span,
                })
            }
            TokenKind::String => {
                let s = tok.string_value().unwrap_or("").to_string();
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::String(s),
                    span,
                })
            }
            TokenKind::Boolean => {
                let b = tok.bool_value().unwrap_or(false);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Bool(b),
                    span,
                })
            }
            TokenKind::Null => {
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Null,
                    span,
                })
            }
            TokenKind::Variable => {
                let _var_text = tok.text.clone();
                let var_value = tok.string_value().unwrap_or("").to_string();
                self.stream.advance();

                if var_value.is_empty() {
                    // bare $
                    return Ok(Expr::LastValue { span });
                }

                let (name, path) = parse_variable_path(&var_value);
                Ok(Expr::Variable { name, path, span })
            }
            TokenKind::LParen => {
                self.stream.advance(); // (
                let expr = self.parse_expression(0)?;
                self.stream.expect(TokenKind::RParen)?;
                Ok(expr)
            }
            TokenKind::LBrace => self.parse_object_literal(),
            TokenKind::LBracket => self.parse_array_literal(),
            TokenKind::SubexpressionOpen => self.parse_subexpression(),
            TokenKind::TemplateLiteral => {
                let raw = tok.string_value().unwrap_or("").to_string();
                self.stream.advance();
                Ok(Expr::TemplateLiteral { raw, span })
            }
            TokenKind::Identifier => {
                // Could be a function call used as expression value like range(1, 5)
                let name = tok.text.clone();
                self.stream.advance();

                // Check for parenthesized call: identifier(args)
                if self.stream.check(TokenKind::LParen) {
                    self.stream.advance(); // consume (
                    let mut args = Vec::new();
                    while !self.stream.is_at_end() && !self.stream.check(TokenKind::RParen) {
                        if self.stream.check(TokenKind::Comma) || self.stream.check(TokenKind::Newline) {
                            self.stream.advance();
                            continue;
                        }
                        if self.stream.check(TokenKind::Comment) {
                            self.stream.advance();
                            continue;
                        }
                        let arg = self.parse_command_arg()?;
                        args.push(arg);
                    }
                    if self.stream.check(TokenKind::RParen) {
                        self.stream.advance(); // consume )
                    }
                    return Ok(Expr::Call {
                        callee: name,
                        args,
                        span,
                    });
                }

                // Not a call — treat as string literal
                Ok(Expr::Literal {
                    value: LiteralValue::String(name),
                    span,
                })
            }
            _ => {
                let line = tok.line;
                let col = tok.column;
                let text = tok.text.clone();
                let kind = tok.kind;
                Err(RobinPathError::parse(
                    format!("Unexpected token in expression: {:?} '{}'", kind, text),
                    line,
                    col,
                ))
            }
        }
    }

    fn parse_object_literal(&mut self) -> Result<Expr, RobinPathError> {
        let start = self.stream.current().unwrap();
        let span_start = Span::from_line(start.line, start.column);
        self.stream.advance(); // {

        let mut properties = Vec::new();

        while !self.stream.is_at_end() && !self.stream.check(TokenKind::RBrace) {
            self.stream.skip_newlines();

            if self.stream.check(TokenKind::RBrace) {
                break;
            }

            // Parse key
            let key_tok = self.stream.current().ok_or_else(|| {
                RobinPathError::parse("Expected object key", 0, 0)
            })?;
            let key_span = Span::from_line(key_tok.line, key_tok.column);
            let key = match key_tok.kind {
                TokenKind::Identifier | TokenKind::Keyword => {
                    let k = key_tok.text.clone();
                    self.stream.advance();
                    k
                }
                TokenKind::String => {
                    let k = key_tok.string_value().unwrap_or("").to_string();
                    self.stream.advance();
                    k
                }
                TokenKind::Number => {
                    let k = key_tok.text.clone();
                    self.stream.advance();
                    k
                }
                TokenKind::LBracket => {
                    // Computed property: [expr]
                    self.stream.advance(); // consume [
                    let key_expr = self.parse_expression(0)?;
                    self.stream.expect(TokenKind::RBracket)?;
                    // We'll use a special marker for computed keys
                    // Store the expression result as key — for now, evaluate at runtime
                    // Use __computed:N as placeholder key
                    let computed_key = format!("__computed_{}", properties.len());
                    self.stream.expect(TokenKind::Colon)?;
                    let value = self.parse_expression(0)?;
                    // For computed properties, we store the key expression as a special property
                    properties.push(ObjectProperty {
                        key: computed_key,
                        value: Expr::ArrayLiteral {
                            elements: vec![key_expr, value],
                            span: key_span,
                        },
                        span: key_span,
                    });
                    self.stream.skip_newlines();
                    if self.stream.check(TokenKind::Comma) {
                        self.stream.advance();
                        self.stream.skip_newlines();
                    }
                    continue;
                }
                _ => {
                    return Err(RobinPathError::parse(
                        format!("Expected object key, got {:?}", key_tok.kind),
                        key_tok.line,
                        key_tok.column,
                    ));
                }
            };

            self.stream.expect(TokenKind::Colon)?;

            let value = self.parse_expression(0)?;

            properties.push(ObjectProperty {
                key,
                value,
                span: key_span,
            });

            self.stream.skip_newlines();

            if self.stream.check(TokenKind::Comma) {
                self.stream.advance();
                self.stream.skip_newlines();
            }
        }

        let end = self.stream.current();
        let end_line = end.map_or(span_start.end_row, |t| t.line);
        let end_col = end.map_or(span_start.end_col, |t| t.column);
        self.stream.expect(TokenKind::RBrace)?;

        Ok(Expr::ObjectLiteral {
            properties,
            span: Span::new(span_start.start_row, span_start.start_col, end_line, end_col),
        })
    }

    fn parse_array_literal(&mut self) -> Result<Expr, RobinPathError> {
        let start = self.stream.current().unwrap();
        let span_start = Span::from_line(start.line, start.column);
        self.stream.advance(); // [

        let mut elements = Vec::new();

        while !self.stream.is_at_end() && !self.stream.check(TokenKind::RBracket) {
            self.stream.skip_newlines();

            if self.stream.check(TokenKind::RBracket) {
                break;
            }

            let elem = self.parse_expression(0)?;
            elements.push(elem);

            self.stream.skip_newlines();

            if self.stream.check(TokenKind::Comma) {
                self.stream.advance();
                self.stream.skip_newlines();
            }
        }

        let end = self.stream.current();
        let end_line = end.map_or(span_start.end_row, |t| t.line);
        let end_col = end.map_or(span_start.end_col, |t| t.column);
        self.stream.expect(TokenKind::RBracket)?;

        Ok(Expr::ArrayLiteral {
            elements,
            span: Span::new(span_start.start_row, span_start.start_col, end_line, end_col),
        })
    }

    fn parse_subexpression(&mut self) -> Result<Expr, RobinPathError> {
        let start = self.stream.current().unwrap();
        let span_start = Span::from_line(start.line, start.column);
        self.stream.advance(); // $(

        let mut body = Vec::new();
        let mut depth = 1;

        // We need to parse statements until matching )
        while !self.stream.is_at_end() && depth > 0 {
            if self.stream.check(TokenKind::RParen) {
                depth -= 1;
                if depth == 0 {
                    break;
                }
            }
            if self.stream.check(TokenKind::SubexpressionOpen) {
                depth += 1;
            }
            if self.stream.check(TokenKind::Newline) {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Comment) {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Eof) {
                break;
            }

            let stmt = self.parse_statement()?;
            body.push(stmt);
            self.stream.skip_newlines();
        }

        let end = self.stream.current();
        let end_line = end.map_or(span_start.end_row, |t| t.line);
        let end_col = end.map_or(span_start.end_col, |t| t.column);
        self.stream.expect(TokenKind::RParen)?;

        Ok(Expr::Subexpression {
            body,
            span: Span::new(span_start.start_row, span_start.start_col, end_line, end_col),
        })
    }

    // ==================== Statement parsers ====================

    fn parse_variable_or_assignment(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        let var_value = tok.string_value().unwrap_or("").to_string();

        if var_value.is_empty() {
            // Bare $ — treat as expression statement
            self.stream.advance();
            return Ok(Stmt::Command {
                name: "_var".to_string(),
                args: vec![Arg::Expr(Expr::LastValue { span })],
                into: None,
                span,
            });
        }

        let (name, path) = parse_variable_path(&var_value);

        // Look ahead for = (assignment)
        let saved = self.stream.save();
        self.stream.advance(); // consume variable

        if self.stream.check(TokenKind::Assign) {
            self.stream.advance(); // consume =
            let value = self.parse_assignment_value()?;

            return Ok(Stmt::Assignment {
                target_name: name,
                target_path: path,
                value,
                is_set: false,
                is_var: false,
                is_const: false,
                fallback: None,
                span,
            });
        }

        // Not an assignment — this is a variable reference on its own line
        // Restore and treat as expression
        self.stream.restore(saved);
        let expr = self.parse_expression(0)?;
        Ok(Stmt::Command {
            name: "_var".to_string(),
            args: vec![Arg::Expr(expr)],
            into: None,
            span,
        })
    }

    fn parse_command_call(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        let name = tok.text.clone();
        self.stream.advance();

        let mut args = Vec::new();
        let mut into = None;

        // Handle parenthesized call syntax: command(arg1, arg2) or command(arg1 arg2)
        if self.stream.check(TokenKind::LParen) {
            self.stream.advance(); // consume (
            while !self.stream.is_at_end() && !self.stream.check(TokenKind::RParen) {
                if self.stream.check(TokenKind::Comma) || self.stream.check(TokenKind::Newline) {
                    self.stream.advance();
                    continue;
                }
                if self.stream.check(TokenKind::Comment) {
                    self.stream.advance();
                    continue;
                }
                // Check for $param=value named arg pattern inside parens
                if self.stream.check(TokenKind::Variable) {
                    let var_tok = self.stream.current().unwrap().clone();
                    let var_text = var_tok.string_value().unwrap_or("").to_string();
                    if !var_text.is_empty() {
                        let saved = self.stream.save();
                        self.stream.advance();
                        if self.stream.check(TokenKind::Assign) {
                            self.stream.advance();
                            let value = self.parse_command_arg()?;
                            args.push(Arg::Named { name: var_text, value });
                            continue;
                        }
                        self.stream.restore(saved);
                    }
                }
                // Check for key=value inside parens
                if self.stream.check(TokenKind::Identifier) {
                    let id_tok = self.stream.current().unwrap().clone();
                    let id_text = id_tok.text.clone();
                    let saved = self.stream.save();
                    self.stream.advance();
                    if self.stream.check(TokenKind::Assign) {
                        self.stream.advance();
                        let value = self.parse_command_arg()?;
                        args.push(Arg::Named { name: id_text, value });
                        continue;
                    }
                    self.stream.restore(saved);
                }
                let arg = self.parse_command_arg()?;
                args.push(Arg::Expr(arg));
            }
            if self.stream.check(TokenKind::RParen) {
                self.stream.advance(); // consume )
            }
            // Check for into after parens
            if self.stream.check_keyword("into") {
                self.stream.advance();
                let var_tok = self.stream.expect(TokenKind::Variable)?;
                let var_val = var_tok.string_value().unwrap_or("").to_string();
                let (var_name, var_path) = parse_variable_path(&var_val);
                into = Some(IntoTarget {
                    name: var_name,
                    path: var_path,
                    span: Span::from_line(var_tok.line, var_tok.column),
                });
            }
            return Ok(Stmt::Command { name, args, into, span });
        }

        while !self.stream.is_at_end()
            && !self.stream.check(TokenKind::Newline)
            && !self.stream.check(TokenKind::Eof)
            && !self.stream.check(TokenKind::Comment)
            && !self.stream.check(TokenKind::RParen)
        {
            // Check for "into" keyword
            if self.stream.check_keyword("into") {
                self.stream.advance(); // consume "into"
                let var_tok = self.stream.expect(TokenKind::Variable)?;
                let var_val = var_tok.string_value().unwrap_or("").to_string();
                let (var_name, var_path) = parse_variable_path(&var_val);
                into = Some(IntoTarget {
                    name: var_name,
                    path: var_path,
                    span: Span::from_line(var_tok.line, var_tok.column),
                });
                break;
            }

            // Check for "then" keyword (inline if)
            if self.stream.check_keyword("then") {
                break;
            }

            // Check for $param=value named arg pattern
            if self.stream.check(TokenKind::Variable) {
                let var_tok = self.stream.current().unwrap().clone();
                let var_text = var_tok.string_value().unwrap_or("").to_string();
                if !var_text.is_empty() {
                    let saved = self.stream.save();
                    self.stream.advance(); // consume variable
                    if self.stream.check(TokenKind::Assign) {
                        self.stream.advance(); // consume =
                        let value = self.parse_command_arg()?;
                        args.push(Arg::Named { name: var_text, value });
                        continue;
                    }
                    self.stream.restore(saved);
                }
            }

            // Check for key=value named arg pattern (identifier)
            if self.stream.check(TokenKind::Identifier) {
                let id_tok = self.stream.current().unwrap().clone();
                let id_text = id_tok.text.clone();
                let saved = self.stream.save();
                self.stream.advance(); // consume identifier
                if self.stream.check(TokenKind::Assign) {
                    self.stream.advance(); // consume =
                    let value = self.parse_command_arg()?;
                    args.push(Arg::Named { name: id_text, value });
                    continue;
                }
                self.stream.restore(saved);
            }

            let arg = self.parse_command_arg()?;
            args.push(Arg::Expr(arg));
        }

        Ok(Stmt::Command {
            name,
            args,
            into,
            span,
        })
    }

    /// Parse a single argument in a command call context
    fn parse_command_arg(&mut self) -> Result<Expr, RobinPathError> {
        let tok = self.stream.current().ok_or_else(|| {
            RobinPathError::parse("Expected argument", 0, 0)
        })?;

        match tok.kind {
            TokenKind::String => {
                let s = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::String(s),
                    span,
                })
            }
            TokenKind::Number => {
                let n = tok.number_value().unwrap_or(0.0);
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Number(n),
                    span,
                })
            }
            TokenKind::Boolean => {
                let b = tok.bool_value().unwrap_or(false);
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Bool(b),
                    span,
                })
            }
            TokenKind::Null => {
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                Ok(Expr::Literal {
                    value: LiteralValue::Null,
                    span,
                })
            }
            TokenKind::Variable => {
                let var_text = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();

                if var_text.is_empty() {
                    Ok(Expr::LastValue { span })
                } else {
                    let (name, path) = parse_variable_path(&var_text);
                    Ok(Expr::Variable { name, path, span })
                }
            }
            TokenKind::LBrace => self.parse_object_literal(),
            TokenKind::LBracket => self.parse_array_literal(),
            TokenKind::SubexpressionOpen => self.parse_subexpression(),
            TokenKind::TemplateLiteral => {
                let raw = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                Ok(Expr::TemplateLiteral { raw, span })
            }
            TokenKind::LParen => {
                // Parenthesized expression in command args
                self.parse_expression(0)
            }
            TokenKind::Identifier | TokenKind::Keyword => {
                // Bare identifier as string argument (like assert message text)
                // Only treat as string if it's not a reserved control keyword
                let text = tok.text.clone();
                let span = Span::from_line(tok.line, tok.column);

                // Don't consume keywords that end blocks
                if matches!(
                    text.as_str(),
                    "into" | "then" | "endif" | "enddo" | "endfor" | "enddef" | "endon"
                        | "endtogether" | "endwith" | "else" | "elseif"
                ) {
                    return Err(RobinPathError::parse(
                        format!("Unexpected keyword '{}'", text),
                        span.start_row,
                        span.start_col,
                    ));
                }

                self.stream.advance();

                // Handle named argument: key=value or key="value"
                if self.stream.check(TokenKind::Assign) {
                    self.stream.advance(); // consume =
                    let value = self.parse_command_arg()?;
                    // Return as an object literal with single property
                    return Ok(Expr::ObjectLiteral {
                        properties: vec![ObjectProperty {
                            key: text,
                            value,
                            span,
                        }],
                        span,
                    });
                }
                Ok(Expr::Literal {
                    value: LiteralValue::String(text),
                    span,
                })
            }
            TokenKind::Minus => {
                // Negative number
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                if let Some(num_tok) = self.stream.current() {
                    if num_tok.kind == TokenKind::Number {
                        let n = num_tok.number_value().unwrap_or(0.0);
                        self.stream.advance();
                        return Ok(Expr::Literal {
                            value: LiteralValue::Number(-n),
                            span,
                        });
                    }
                }
                Ok(Expr::Unary {
                    op: UnaryOp::Neg,
                    argument: Box::new(self.parse_primary()?),
                    span,
                })
            }
            _ => self.parse_expression(0),
        }
    }

    fn parse_if(&mut self) -> Result<Stmt, RobinPathError> {
        let if_tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(if_tok.line, if_tok.column);
        self.stream.advance(); // consume "if"

        let condition = self.parse_expression(0)?;

        // Check for inline if: if cond then command [else command]
        if self.stream.check_keyword("then") {
            self.stream.advance(); // consume "then"

            // If followed by newline, this is a block-if with explicit "then"
            if self.stream.check(TokenKind::Newline) || self.stream.check(TokenKind::Eof) {
                // Fall through to block-if parsing below
            } else {
                let command = self.parse_statement()?;

                let else_command = if self.stream.check_keyword("else") {
                    self.stream.advance();
                    Some(Box::new(self.parse_statement()?))
                } else {
                    None
                };

                return Ok(Stmt::InlineIf {
                    condition,
                    command: Box::new(command),
                    else_command,
                    span,
                });
            }
        }

        // Block if
        self.stream.skip_newlines();

        let then_branch = self.parse_block_until(&["elseif", "else", "endif"])?;

        let mut elseif_branches = Vec::new();
        while self.stream.check_keyword("elseif") {
            let ei_tok = self.stream.current().unwrap().clone();
            let ei_span = Span::from_line(ei_tok.line, ei_tok.column);
            self.stream.advance(); // consume "elseif"
            let cond = self.parse_expression(0)?;
            self.stream.skip_newlines();
            let body = self.parse_block_until(&["elseif", "else", "endif"])?;
            elseif_branches.push(ElseIfBranch {
                condition: cond,
                body,
                span: ei_span,
            });
        }

        let else_branch = if self.stream.check_keyword("else") {
            self.stream.advance();
            self.stream.skip_newlines();
            Some(self.parse_block_until(&["endif"])?)
        } else {
            None
        };

        self.stream.expect_keyword("endif")?;

        Ok(Stmt::IfBlock {
            condition,
            then_branch,
            elseif_branches,
            else_branch,
            span,
        })
    }

    fn parse_iftrue(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "iftrue"
        let command = self.parse_statement()?;
        Ok(Stmt::IfTrue {
            command: Box::new(command),
            span,
        })
    }

    fn parse_iffalse(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "iffalse"
        let command = self.parse_statement()?;
        Ok(Stmt::IfFalse {
            command: Box::new(command),
            span,
        })
    }

    fn parse_for_loop(&mut self) -> Result<Stmt, RobinPathError> {
        let for_tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(for_tok.line, for_tok.column);
        self.stream.advance(); // consume "for"

        // Parse loop variable
        let var_tok = self.stream.expect(TokenKind::Variable)?;
        let var_name = var_tok.string_value().unwrap_or("").to_string();

        let mut key_var_name = None;

        // Check for "key" before "in" or "from"
        // Not standard, but check if key comes first

        if self.stream.check_keyword("in") {
            self.stream.advance(); // consume "in"
            let iterable = self.parse_assignment_value()?;

            // Check for optional "key $var"
            if self.stream.check_keyword("key") {
                self.stream.advance();
                let key_tok = self.stream.expect(TokenKind::Variable)?;
                key_var_name = Some(key_tok.string_value().unwrap_or("").to_string());
            }

            self.stream.skip_newlines();
            let body = self.parse_block_until(&["endfor"])?;
            self.stream.expect_keyword("endfor")?;

            return Ok(Stmt::ForLoop {
                var_name,
                iterable: Some(iterable),
                from: None,
                to: None,
                step: None,
                key_var_name,
                body,
                span,
            });
        }

        if self.stream.check_keyword("from") {
            self.stream.advance(); // consume "from"
            let from = self.parse_expression(0)?;
            self.stream.expect_keyword("to")?;
            let to = self.parse_expression(0)?;

            let step = if self.stream.check_keyword("by") || self.stream.check_keyword("step") {
                self.stream.advance();
                Some(self.parse_expression(0)?)
            } else {
                None
            };

            // Check for optional "key $var"
            if self.stream.check_keyword("key") {
                self.stream.advance();
                let key_tok = self.stream.expect(TokenKind::Variable)?;
                key_var_name = Some(key_tok.string_value().unwrap_or("").to_string());
            }

            self.stream.skip_newlines();
            let body = self.parse_block_until(&["endfor"])?;
            self.stream.expect_keyword("endfor")?;

            return Ok(Stmt::ForLoop {
                var_name,
                iterable: None,
                from: Some(from),
                to: Some(to),
                step,
                key_var_name,
                body,
                span,
            });
        }

        Err(RobinPathError::parse(
            "Expected 'in' or 'from' after for variable",
            var_tok.line,
            var_tok.column,
        ))
    }

    fn parse_define(&mut self) -> Result<Stmt, RobinPathError> {
        let def_tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(def_tok.line, def_tok.column);
        self.stream.advance(); // consume "def" or "define"

        // Function name
        let name_tok = self.stream.expect(TokenKind::Identifier)?;
        let name = name_tok.text.clone();

        // Parameter names
        let mut param_names = Vec::new();
        while self.stream.check(TokenKind::Variable)
            && !self.stream.is_at_end()
            && !self.stream.check(TokenKind::Newline)
        {
            let p = self.stream.expect(TokenKind::Variable)?;
            param_names.push(p.string_value().unwrap_or("").to_string());
        }

        // Optional "as" keyword (ignored)
        self.stream.match_keyword("as");

        self.stream.skip_newlines();
        let body = self.parse_block_until(&["enddef"])?;
        self.stream.expect_keyword("enddef")?;

        Ok(Stmt::Define {
            name,
            param_names,
            body,
            span,
        })
    }

    fn parse_scope_block(&mut self) -> Result<Stmt, RobinPathError> {
        let do_tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(do_tok.line, do_tok.column);
        self.stream.advance(); // consume "do"

        // Check for "into" before body: do into $var
        let mut into = None;
        if self.stream.check_keyword("into") {
            self.stream.advance(); // consume "into"
            let var_tok = self.stream.expect(TokenKind::Variable)?;
            let var_val = var_tok.string_value().unwrap_or("").to_string();
            let (var_name, var_path) = parse_variable_path(&var_val);
            into = Some(IntoTarget {
                name: var_name,
                path: var_path,
                span: Span::from_line(var_tok.line, var_tok.column),
            });
        }

        // Optional parameters
        let mut param_names = Vec::new();
        while self.stream.check(TokenKind::Variable)
            && !self.stream.check(TokenKind::Newline)
            && !self.stream.is_at_end()
        {
            // Peek ahead: if this variable is followed by `into`, it's a param, then into target
            // But we can't know yet — just collect it
            let p = self.stream.expect(TokenKind::Variable)?;
            let pname = p.string_value().unwrap_or("").to_string();
            if !pname.is_empty() {
                param_names.push(pname);
            } else {
                break;
            }
        }

        // Check for "into" after params: do $a $b into $result
        if into.is_none() && self.stream.check_keyword("into") {
            self.stream.advance(); // consume "into"
            let var_tok = self.stream.expect(TokenKind::Variable)?;
            let var_val = var_tok.string_value().unwrap_or("").to_string();
            let (var_name, var_path) = parse_variable_path(&var_val);
            into = Some(IntoTarget {
                name: var_name,
                path: var_path,
                span: Span::from_line(var_tok.line, var_tok.column),
            });
        }

        self.stream.skip_newlines();
        let body = self.parse_block_until(&["enddo"])?;
        self.stream.expect_keyword("enddo")?;

        // Also check for "into" after enddo (alternative syntax)
        if into.is_none() {
            into = self.parse_optional_into()?;
        }

        Ok(Stmt::ScopeBlock {
            param_names,
            body,
            into,
            span,
        })
    }

    fn parse_repeat_or_command(&mut self) -> Result<Stmt, RobinPathError> {
        // peek ahead to see if this is "repeat N with ... endwith" or just "repeat N value"
        let saved = self.stream.save();
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "repeat"

        // Parse the count expression
        if self.stream.is_at_end() || self.stream.check(TokenKind::Newline) || self.stream.check(TokenKind::Eof) {
            self.stream.restore(saved);
            return self.parse_command_call();
        }

        let count = self.parse_command_arg()?;

        // Check if next token is "with" keyword
        if self.stream.check_keyword("with") {
            self.stream.advance(); // consume "with"

            // Check for "into" after "with"
            let mut into = None;
            if self.stream.check_keyword("into") {
                self.stream.advance();
                let var_tok = self.stream.expect(TokenKind::Variable)?;
                let var_val = var_tok.string_value().unwrap_or("").to_string();
                let (var_name, var_path) = parse_variable_path(&var_val);
                into = Some(IntoTarget {
                    name: var_name,
                    path: var_path,
                    span: Span::from_line(var_tok.line, var_tok.column),
                });
            }

            self.stream.skip_newlines();
            let body = self.parse_block_until(&["endwith"])?;
            self.stream.expect_keyword("endwith")?;

            // Check for "into" after "endwith"
            if into.is_none() {
                into = self.parse_optional_into()?;
            }

            return Ok(Stmt::RepeatBlock {
                count,
                body,
                into,
                span,
            });
        }

        // Not a "with" block — restore and parse as command
        self.stream.restore(saved);
        self.parse_command_call()
    }

    fn parse_together(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "together"
        self.stream.skip_newlines();

        let mut blocks = Vec::new();

        while !self.stream.is_at_end() && !self.stream.check_keyword("endtogether") {
            if self.stream.check(TokenKind::Newline) || self.stream.check(TokenKind::Comment) {
                self.stream.advance();
                continue;
            }
            if self.stream.check_keyword("do") {
                let block = self.parse_scope_block()?;
                blocks.push(block);
            } else {
                // Skip unexpected tokens
                self.stream.advance();
            }
            self.stream.skip_newlines();
        }

        self.stream.expect_keyword("endtogether")?;

        Ok(Stmt::TogetherBlock { blocks, span })
    }

    fn parse_on_block(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "on"

        let event_name = if self.stream.check(TokenKind::String) {
            let event_tok = self.stream.expect(TokenKind::String)?;
            event_tok.string_value().unwrap_or("").to_string()
        } else {
            let event_tok = self.stream.expect(TokenKind::Identifier)?;
            event_tok.text.clone()
        };

        self.stream.skip_newlines();
        // on blocks can be closed by "endon" or auto-closed at another "on" / EOF
        let body = self.parse_block_until(&["endon", "on"])?;
        if self.stream.check_keyword("endon") {
            self.stream.advance(); // consume endon
        }
        // If we stopped at "on", don't consume it — it's the start of the next handler

        Ok(Stmt::OnBlock {
            event_name,
            body,
            span,
        })
    }

    fn parse_return(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "return"

        let value =
            if !self.stream.is_at_end()
                && !self.stream.check(TokenKind::Newline)
                && !self.stream.check(TokenKind::Eof)
                && !self.stream.check(TokenKind::Comment)
            {
                Some(self.parse_assignment_value()?)
            } else {
                None
            };

        Ok(Stmt::Return { value, span })
    }

    fn parse_break(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance();
        Ok(Stmt::Break { span })
    }

    fn parse_continue(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance();
        Ok(Stmt::Continue { span })
    }

    fn parse_set(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "set"

        // set $var [as] value [fallback]
        let var_tok = self.stream.expect(TokenKind::Variable)?;
        let var_val = var_tok.string_value().unwrap_or("").to_string();
        let (name, path) = parse_variable_path(&var_val);

        // Optional "as" keyword or "=" sign
        if !self.stream.match_keyword("as") {
            // Also try consuming = sign
            if self.stream.check(TokenKind::Assign) {
                self.stream.advance();
            }
        }

        let value = self.parse_assignment_value()?;

        // Optional fallback
        let fallback =
            if !self.stream.is_at_end()
                && !self.stream.check(TokenKind::Newline)
                && !self.stream.check(TokenKind::Eof)
                && !self.stream.check(TokenKind::RParen)
                && !self.stream.check(TokenKind::Comment)
            {
                Some(self.parse_command_arg()?)
            } else {
                None
            };

        Ok(Stmt::Assignment {
            target_name: name,
            target_path: path,
            value,
            is_set: true,
            is_var: false,
            is_const: false,
            fallback,
            span,
        })
    }

    fn parse_var_decl(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "var"

        let var_tok = self.stream.expect(TokenKind::Variable)?;
        let var_val = var_tok.string_value().unwrap_or("").to_string();
        let (name, path) = parse_variable_path(&var_val);

        let value =
            if self.stream.check(TokenKind::Assign) {
                self.stream.advance();
                self.parse_command_arg()?
            } else if !self.stream.is_at_end()
                && !self.stream.check(TokenKind::Newline)
                && !self.stream.check(TokenKind::Eof)
            {
                self.parse_command_arg()?
            } else {
                Expr::Literal {
                    value: LiteralValue::Null,
                    span,
                }
            };

        Ok(Stmt::Assignment {
            target_name: name,
            target_path: path,
            value,
            is_set: false,
            is_var: true,
            is_const: false,
            fallback: None,
            span,
        })
    }

    fn parse_const_decl(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "const"

        let var_tok = self.stream.expect(TokenKind::Variable)?;
        let var_val = var_tok.string_value().unwrap_or("").to_string();
        let (name, path) = parse_variable_path(&var_val);

        // Optional = sign
        self.stream.match_kind(TokenKind::Assign);

        let value = self.parse_command_arg()?;

        Ok(Stmt::Assignment {
            target_name: name,
            target_path: path,
            value,
            is_set: false,
            is_var: false,
            is_const: true,
            fallback: None,
            span,
        })
    }

    fn parse_use(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        self.stream.advance(); // consume "use"

        let module_tok = self.stream.expect(TokenKind::Identifier)?;
        let module_name = module_tok.text.clone();

        Ok(Stmt::UseModule { module_name, span })
    }

    fn parse_decorator(&mut self) -> Result<Stmt, RobinPathError> {
        let tok = self.stream.current().unwrap().clone();
        let span = Span::from_line(tok.line, tok.column);
        let name = tok.string_value().unwrap_or("").to_string();
        self.stream.advance(); // consume @name

        // Optional decorator args on same line
        let mut args = Vec::new();
        while !self.stream.is_at_end()
            && !self.stream.check(TokenKind::Newline)
            && !self.stream.check(TokenKind::Eof)
        {
            let arg = self.parse_command_arg()?;
            args.push(Arg::Expr(arg));
        }

        Ok(Stmt::Decorator { name, args, span })
    }

    /// Parse the right-hand side of an assignment.
    /// RobinPath allows `$var = command arg1 arg2` in addition to `$var = expr`.
    /// If the RHS starts with an identifier that looks like a command, treat it as a call.
    fn parse_assignment_value(&mut self) -> Result<Expr, RobinPathError> {
        let tok = self.stream.current().ok_or_else(|| {
            RobinPathError::parse("Expected value after =", 0, 0)
        })?;

        // If it's an identifier (potential command call like `range 1 3` or `json.parse "..."`),
        // check if it's followed by arguments
        if tok.kind == TokenKind::Identifier {
            let name = tok.text.clone();
            let span = Span::from_line(tok.line, tok.column);

            // Check if there are arguments after the identifier (on the same line)
            let saved = self.stream.save();
            self.stream.advance(); // consume identifier

            // If next token is a value token on the same line, this is a command call
            if let Some(next) = self.stream.current() {
                // Handle parenthesized call: command(arg1, arg2) or command(arg1 arg2)
                if next.kind == TokenKind::LParen {
                    self.stream.advance(); // consume (
                    let mut args = Vec::new();
                    while !self.stream.is_at_end() && !self.stream.check(TokenKind::RParen) {
                        if self.stream.check(TokenKind::Comma) || self.stream.check(TokenKind::Newline) {
                            self.stream.advance();
                            continue;
                        }
                        if self.stream.check(TokenKind::Comment) {
                            self.stream.advance();
                            continue;
                        }
                        let arg = self.parse_command_arg()?;
                        args.push(arg);
                    }
                    if self.stream.check(TokenKind::RParen) {
                        self.stream.advance(); // consume )
                    }
                    return Ok(Expr::Call {
                        callee: name,
                        args,
                        span,
                    });
                }

                if next.kind != TokenKind::Newline
                    && next.kind != TokenKind::Eof
                    && next.kind != TokenKind::Comment
                    && !next.is_keyword("into")
                {
                    // It has arguments — parse as command call expression
                    let mut args = Vec::new();
                    while !self.stream.is_at_end()
                        && !self.stream.check(TokenKind::Newline)
                        && !self.stream.check(TokenKind::Eof)
                        && !self.stream.check(TokenKind::Comment)
                        && !self.stream.check_keyword("into")
                    {
                        let arg = self.parse_command_arg()?;
                        args.push(arg);
                    }

                    return Ok(Expr::Call {
                        callee: name,
                        args,
                        span,
                    });
                }
            }

            // No arguments — restore and parse as expression
            self.stream.restore(saved);
        }

        self.parse_expression(0)
    }

    // ==================== Helpers ====================

    fn parse_block_until(&mut self, end_keywords: &[&str]) -> Result<Vec<Stmt>, RobinPathError> {
        let mut stmts = Vec::new();
        let mut iterations = 0;
        let max_iterations = 100_000;

        while !self.stream.is_at_end() {
            iterations += 1;
            if iterations > max_iterations {
                return Err(RobinPathError::parse(
                    "Block parsing exceeded maximum iterations",
                    self.stream.current_line(),
                    self.stream.current_column(),
                ));
            }

            if self.stream.check(TokenKind::Newline) {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Comment) {
                let tok = self.stream.current().unwrap().clone();
                let text = tok
                    .string_value()
                    .unwrap_or(&tok.text[1..])
                    .trim()
                    .to_string();
                let span = Span::from_line(tok.line, tok.column);
                stmts.push(Stmt::Comment { text, span });
                self.stream.advance();
                continue;
            }
            // Skip cell/prompt-related tokens that may appear inside blocks
            if self.stream.check(TokenKind::ChunkMarker) {
                let tok = self.stream.current().unwrap().clone();
                let id = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                stmts.push(Stmt::ChunkMarker { id, span });
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::CellOpen) {
                let tok = self.stream.current().unwrap().clone();
                let cell_type = tok.string_value().unwrap_or("").to_string();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                let mut raw_text = None;
                let mut content = Vec::new();
                if self.stream.check(TokenKind::RawText) {
                    let raw_tok = self.stream.current().unwrap().clone();
                    raw_text = Some(raw_tok.string_value().unwrap_or("").to_string());
                    self.stream.advance();
                } else {
                    // Code cell inside a block
                    self.stream.skip_newlines();
                    while !self.stream.is_at_end() && !self.stream.check(TokenKind::CellClose) {
                        if self.stream.check(TokenKind::Newline) { self.stream.advance(); continue; }
                        if self.stream.check(TokenKind::Comment) {
                            let ctok = self.stream.current().unwrap().clone();
                            let text = ctok.string_value().unwrap_or(&ctok.text[1..]).trim().to_string();
                            let cspan = Span::from_line(ctok.line, ctok.column);
                            content.push(Stmt::Comment { text, span: cspan });
                            self.stream.advance();
                            continue;
                        }
                        // Handle prompt blocks inside code cells
                        if self.stream.check(TokenKind::PromptOpen) {
                            let ptok = self.stream.current().unwrap().clone();
                            let pspan = Span::from_line(ptok.line, ptok.column);
                            self.stream.advance();
                            let text = if self.stream.check(TokenKind::RawText) {
                                let raw_tok = self.stream.current().unwrap().clone();
                                let t = raw_tok.string_value().unwrap_or("").to_string();
                                self.stream.advance();
                                t
                            } else { String::new() };
                            if self.stream.check(TokenKind::PromptClose) { self.stream.advance(); }
                            content.push(Stmt::PromptBlock { text, span: pspan });
                            continue;
                        }
                        let inner = self.parse_statement()?;
                        content.push(inner);
                    }
                }
                if self.stream.check(TokenKind::CellClose) { self.stream.advance(); }
                stmts.push(Stmt::CellBlock { cell_type, content, raw_text, span });
                continue;
            }
            if self.stream.check(TokenKind::PromptOpen) {
                let tok = self.stream.current().unwrap().clone();
                let span = Span::from_line(tok.line, tok.column);
                self.stream.advance();
                let text = if self.stream.check(TokenKind::RawText) {
                    let raw_tok = self.stream.current().unwrap().clone();
                    let t = raw_tok.string_value().unwrap_or("").to_string();
                    self.stream.advance();
                    t
                } else { String::new() };
                if self.stream.check(TokenKind::PromptClose) { self.stream.advance(); }
                stmts.push(Stmt::PromptBlock { text, span });
                continue;
            }
            if self.stream.check(TokenKind::CellClose)
                || self.stream.check(TokenKind::PromptClose)
                || self.stream.check(TokenKind::RawText)
            {
                self.stream.advance();
                continue;
            }
            if self.stream.check(TokenKind::Eof) {
                break;
            }

            // Check if current token is one of the end keywords
            if let Some(tok) = self.stream.current() {
                if tok.kind == TokenKind::Keyword {
                    for kw in end_keywords {
                        if tok.text == *kw {
                            return Ok(stmts);
                        }
                    }
                }
            }

            let stmt = self.parse_statement()?;
            stmts.push(stmt);
        }

        Ok(stmts)
    }

    fn parse_optional_into(&mut self) -> Result<Option<IntoTarget>, RobinPathError> {
        if self.stream.check_keyword("into") {
            self.stream.advance();
            let var_tok = self.stream.expect(TokenKind::Variable)?;
            let var_val = var_tok.string_value().unwrap_or("").to_string();
            let (name, path) = parse_variable_path(&var_val);
            Ok(Some(IntoTarget {
                name,
                path,
                span: Span::from_line(var_tok.line, var_tok.column),
            }))
        } else {
            Ok(None)
        }
    }
}

/// Parse a variable path string like "obj.prop[0][$key]" into name + path segments
pub fn parse_variable_path(input: &str) -> (String, Vec<AttributePathSegment>) {
    let mut path = Vec::new();
    let mut chars = input.chars().peekable();
    let mut name = String::new();

    // Read variable name (up to first . or [)
    while let Some(&ch) = chars.peek() {
        if ch == '.' || ch == '[' {
            break;
        }
        name.push(ch);
        chars.next();
    }

    // Read path segments
    while let Some(&ch) = chars.peek() {
        if ch == '.' {
            chars.next(); // skip .
            let mut prop = String::new();
            while let Some(&c) = chars.peek() {
                if c == '.' || c == '[' {
                    break;
                }
                prop.push(c);
                chars.next();
            }
            if !prop.is_empty() {
                path.push(AttributePathSegment::Property { name: prop });
            }
        } else if ch == '[' {
            chars.next(); // skip [
            let mut inner = String::new();
            while let Some(&c) = chars.peek() {
                if c == ']' {
                    break;
                }
                inner.push(c);
                chars.next();
            }
            if chars.peek() == Some(&']') {
                chars.next(); // skip ]
            }

            // Check if it's a numeric index
            if let Ok(idx) = inner.parse::<i64>() {
                path.push(AttributePathSegment::Index { index: idx });
            } else if inner.starts_with('$') {
                // Dynamic key: [$var] or [$var.path]
                let var_content = &inner[1..];
                let (dyn_name, dyn_path) = if var_content.contains('.') || var_content.contains('[')
                {
                    let (n, p) = parse_variable_path(var_content);
                    (n, Some(p))
                } else {
                    (var_content.to_string(), None)
                };
                path.push(AttributePathSegment::DynamicKey {
                    variable: dyn_name,
                    var_path: dyn_path,
                });
            } else {
                // String key
                path.push(AttributePathSegment::Property {
                    name: inner.trim_matches('"').trim_matches('\'').to_string(),
                });
            }
        } else {
            break;
        }
    }

    (name, path)
}