dreamstack/compiler/ds-parser/src/parser.rs

/// DreamStack Parser — recursive descent parser producing AST from tokens.
use crate::ast::*;
use crate::lexer::{Token, TokenKind};

pub struct Parser {
    tokens: Vec<Token>,
    pos: usize,
}

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

    pub fn parse_program(&mut self) -> Result<Program, ParseError> {
        let mut declarations = Vec::new();
        self.skip_newlines();

        while !self.is_at_end() {
            let decl = self.parse_declaration()?;
            declarations.push(decl);
            self.skip_newlines();
        }

        Ok(Program { declarations })
    }

    // ── Helpers ──────────────────────────────────────────

    fn peek(&self) -> &TokenKind {
        self.tokens
            .get(self.pos)
            .map(|t| &t.kind)
            .unwrap_or(&TokenKind::Eof)
    }

    fn current_token(&self) -> &Token {
        &self.tokens[self.pos.min(self.tokens.len() - 1)]
    }

    fn advance(&mut self) -> &Token {
        let tok = &self.tokens[self.pos.min(self.tokens.len() - 1)];
        self.pos += 1;
        tok
    }

    fn expect(&mut self, expected: &TokenKind) -> Result<&Token, ParseError> {
        if std::mem::discriminant(self.peek()) == std::mem::discriminant(expected) {
            Ok(self.advance())
        } else {
            Err(self.error(format!("expected {expected:?}, got {:?}", self.peek())))
        }
    }

    fn expect_ident(&mut self) -> Result<String, ParseError> {
        match self.peek().clone() {
            TokenKind::Ident(name) => {
                self.advance();
                Ok(name)
            }
            // Also accept keywords that can be used as identifiers in some contexts
            _ => Err(self.error(format!("expected identifier, got {:?}", self.peek()))),
        }
    }

    fn check(&self, kind: &TokenKind) -> bool {
        std::mem::discriminant(self.peek()) == std::mem::discriminant(kind)
    }

    fn is_at_end(&self) -> bool {
        matches!(self.peek(), TokenKind::Eof)
    }

    fn skip_newlines(&mut self) {
        while matches!(self.peek(), TokenKind::Newline) {
            self.advance();
        }
    }

    fn error(&self, msg: String) -> ParseError {
        let tok = self.current_token();
        ParseError {
            message: msg,
            line: tok.line,
            col: tok.col,
        }
    }

    // ── Declarations ────────────────────────────────────

    fn parse_declaration(&mut self) -> Result<Declaration, ParseError> {
        match self.peek() {
            TokenKind::Let => self.parse_let_decl(),
            TokenKind::View => self.parse_view_decl(),
            TokenKind::Effect => self.parse_effect_decl(),
            TokenKind::On => self.parse_on_handler(),
            TokenKind::Component => self.parse_component_decl(),
            TokenKind::Route => self.parse_route_decl(),
            TokenKind::Constrain => self.parse_constrain_decl(),
            TokenKind::Stream => self.parse_stream_decl(),
            TokenKind::Every => self.parse_every_decl(),
            TokenKind::Import => self.parse_import_decl(),
            TokenKind::Export => self.parse_export_decl(),
            // Expression statement: `log("hello")`, `push(items, x)`
            TokenKind::Ident(_) => {
                let expr = self.parse_expr()?;
                Ok(Declaration::ExprStatement(expr))
            }
            _ => Err(self.error(format!(
                "expected declaration (let, view, effect, on, component, route, constrain, stream, every), got {:?}",
                self.peek()
            ))),
        }
    }

    fn parse_every_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'every'
        let interval = self.parse_expr()?;
        self.expect(&TokenKind::Arrow)?;
        self.skip_newlines();
        let body = self.parse_expr()?;
        Ok(Declaration::Every(EveryDecl {
            interval_ms: interval,
            body,
            span: Span { start: 0, end: 0, line },
        }))
    }

    // import { Card, Button } from "./components"
    fn parse_import_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'import'

        // Parse the import names: { name1, name2 }
        self.expect(&TokenKind::LBrace)?;
        let mut names = Vec::new();
        loop {
            self.skip_newlines();
            if self.check(&TokenKind::RBrace) {
                break;
            }
            match self.peek().clone() {
                TokenKind::Ident(name) => {
                    names.push(name);
                    self.advance();
                }
                // Allow keywords to be imported as names (e.g., component names)
                _ => {
                    let tok = self.peek().clone();
                    return Err(self.error(format!("expected identifier in import, got {:?}", tok)));
                }
            }
            self.skip_newlines();
            if self.check(&TokenKind::Comma) {
                self.advance(); // consume ','
            }
        }
        self.expect(&TokenKind::RBrace)?;

        // Parse "from"
        self.expect(&TokenKind::From)?;

        // Parse the source path
        let source = match self.peek().clone() {
            TokenKind::StringFragment(s) => {
                self.advance();
                // Consume the StringEnd if present
                if self.check(&TokenKind::StringEnd) {
                    self.advance();
                }
                s
            }
            TokenKind::StringEnd => {
                self.advance();
                String::new()
            }
            _ => return Err(self.error("expected string after 'from'".to_string())),
        };

        Ok(Declaration::Import(ImportDecl {
            names,
            source,
            span: Span { start: 0, end: 0, line },
        }))
    }

    // export let count = 0 / export component Card(...) = ...
    fn parse_export_decl(&mut self) -> Result<Declaration, ParseError> {
        self.advance(); // consume 'export'

        // Parse the inner declaration
        let inner = self.parse_declaration()?;

        // Extract the name being exported
        let name = match &inner {
            Declaration::Let(d) => d.name.clone(),
            Declaration::View(d) => d.name.clone(),
            Declaration::Component(d) => d.name.clone(),
            Declaration::Effect(d) => d.name.clone(),
            _ => return Err(self.error("can only export let, view, component, or effect".to_string())),
        };

        Ok(Declaration::Export(name, Box::new(inner)))
    }

    fn parse_let_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'let'
        let name = self.expect_ident()?;
        self.expect(&TokenKind::Eq)?;
        let value = self.parse_expr()?;

        Ok(Declaration::Let(LetDecl {
            name,
            value,
            span: Span { start: 0, end: 0, line },
        }))
    }

    fn parse_view_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'view'
        let name = self.expect_ident()?;

        // Optional parameters
        let params = if self.check(&TokenKind::LParen) {
            self.parse_params()?
        } else {
            Vec::new()
        };

        self.expect(&TokenKind::Eq)?;
        self.skip_newlines();
        let body = self.parse_expr()?;

        Ok(Declaration::View(ViewDecl {
            name,
            params,
            body,
            span: Span { start: 0, end: 0, line },
        }))
    }

    fn parse_effect_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'effect'
        let name = self.expect_ident()?;
        let params = self.parse_params()?;
        self.expect(&TokenKind::Colon)?;
        let return_type = self.parse_type_expr()?;

        Ok(Declaration::Effect(EffectDecl {
            name,
            params,
            return_type,
            span: Span { start: 0, end: 0, line },
        }))
    }

    fn parse_on_handler(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'on'
        let event = self.expect_ident()?;

        // Optional parameter: `on drag(event) ->`
        let param = if self.check(&TokenKind::LParen) {
            self.advance();
            let p = self.expect_ident()?;
            self.expect(&TokenKind::RParen)?;
            Some(p)
        } else {
            None
        };

        self.expect(&TokenKind::Arrow)?;
        self.skip_newlines();
        let body = self.parse_expr()?;

        Ok(Declaration::OnHandler(OnHandler {
            event,
            param,
            body,
            span: Span { start: 0, end: 0, line },
        }))
    }

    /// `component Card(title: String) = column [ ... ]`
    fn parse_component_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'component'
        let name = self.expect_ident()?;

        // Props
        let props = if self.check(&TokenKind::LParen) {
            self.parse_params()?
        } else {
            Vec::new()
        };

        self.expect(&TokenKind::Eq)?;
        self.skip_newlines();
        let body = self.parse_expr()?;

        Ok(Declaration::Component(ComponentDecl {
            name,
            props,
            body,
            span: Span { start: 0, end: 0, line },
        }))
    }

    /// `route "/users/:id" -> column [ ... ]`
    fn parse_route_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'route'

        // Path string: consume "path" as StringStart + StringFragment + StringEnd
        let path = if matches!(self.peek(), TokenKind::StringFragment(_)) {
            // Simple string token (no interpolation)
            if let TokenKind::StringFragment(s) = self.peek() {
                let p = s.clone();
                self.advance(); // consume fragment
                // Consume trailing StringEnd if present
                if matches!(self.peek(), TokenKind::StringEnd) {
                    self.advance();
                }
                p
            } else {
                return Err(self.error("expected path string after 'route'".to_string()));
            }
        } else {
            return Err(self.error("expected path string after 'route'".to_string()));
        };

        self.expect(&TokenKind::Arrow)?;
        self.skip_newlines();
        let body = self.parse_expr()?;

        Ok(Declaration::Route(RouteDecl {
            path,
            body,
            span: Span { start: 0, end: 0, line },
        }))
    }

    /// `constrain sidebar.width = clamp(200, viewport.width * 0.2, 350)`
    fn parse_constrain_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume 'constrain'

        let element = self.expect_ident()?;
        self.expect(&TokenKind::Dot)?;
        let prop = self.expect_ident()?;

        self.expect(&TokenKind::Eq)?;
        self.skip_newlines();
        let expr = self.parse_expr()?;

        Ok(Declaration::Constrain(ConstrainDecl {
            element,
            prop,
            expr,
            span: Span { start: 0, end: 0, line },
        }))
    }

    /// Parse: `stream <view_name> on <url_expr> { mode: pixel | delta | signal }`
    fn parse_stream_decl(&mut self) -> Result<Declaration, ParseError> {
        let line = self.current_token().line;
        self.advance(); // consume `stream`

        let view_name = self.expect_ident()?;

        // Expect `on`
        match self.peek() {
            TokenKind::On => { self.advance(); }
            _ => return Err(self.error("Expected 'on' after stream view name".into())),
        }

        let relay_url = self.parse_expr()?;

        // Optional mode block: `{ mode: signal, transport: webrtc }`
        let mut mode = StreamMode::Signal;
        let mut transport = StreamTransport::WebSocket;

        if self.check(&TokenKind::LBrace) {
            self.advance(); // {
            while !self.check(&TokenKind::RBrace) && !self.is_at_end() {
                self.skip_newlines();
                let key = self.expect_ident()?;
                self.expect(&TokenKind::Colon)?;
                if key == "mode" {
                    match self.peek() {
                        TokenKind::Pixel => { mode = StreamMode::Pixel; self.advance(); }
                        TokenKind::Delta => { mode = StreamMode::Delta; self.advance(); }
                        TokenKind::Signals => { mode = StreamMode::Signal; self.advance(); }
                        TokenKind::Ident(s) if s == "signal" => { mode = StreamMode::Signal; self.advance(); }
                        _ => return Err(self.error("Expected pixel, delta, or signal".into())),
                    }
                } else if key == "transport" {
                    match self.peek() {
                        TokenKind::Ident(s) if s == "webrtc" || s == "WebRTC" => {
                            transport = StreamTransport::WebRTC;
                            self.advance();
                        }
                        TokenKind::Ident(s) if s == "websocket" || s == "ws" => {
                            transport = StreamTransport::WebSocket;
                            self.advance();
                        }
                        _ => return Err(self.error("Expected webrtc or websocket".into())),
                    }
                }
                if self.check(&TokenKind::Comma) { self.advance(); }
                self.skip_newlines();
            }
            self.expect(&TokenKind::RBrace)?;
        }

        Ok(Declaration::Stream(StreamDecl {
            view_name,
            relay_url,
            mode,
            transport,
            span: Span { start: 0, end: 0, line },
        }))
    }

    fn parse_params(&mut self) -> Result<Vec<Param>, ParseError> {
        self.expect(&TokenKind::LParen)?;
        let mut params = Vec::new();

        while !self.check(&TokenKind::RParen) && !self.is_at_end() {
            let name = self.expect_ident()?;
            let type_annotation = if self.check(&TokenKind::Colon) {
                self.advance();
                Some(self.parse_type_expr()?)
            } else {
                None
            };
            params.push(Param { name, type_annotation });
            if self.check(&TokenKind::Comma) {
                self.advance();
            }
        }

        self.expect(&TokenKind::RParen)?;
        Ok(params)
    }

    fn parse_type_expr(&mut self) -> Result<TypeExpr, ParseError> {
        let name = self.expect_ident()?;
        if self.check(&TokenKind::Lt) {
            self.advance();
            let mut type_args = Vec::new();
            while !self.check(&TokenKind::Gt) && !self.is_at_end() {
                type_args.push(self.parse_type_expr()?);
                if self.check(&TokenKind::Comma) {
                    self.advance();
                }
            }
            self.expect(&TokenKind::Gt)?;
            Ok(TypeExpr::Generic(name, type_args))
        } else {
            Ok(TypeExpr::Named(name))
        }
    }

    // ── Expressions ─────────────────────────────────────

    fn parse_expr(&mut self) -> Result<Expr, ParseError> {
        self.parse_pipe_expr()
    }

    /// Pipe: `expr | operator`
    fn parse_pipe_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_assignment()?;

        while self.check(&TokenKind::Pipe) {
            self.advance();
            self.skip_newlines();
            let right = self.parse_assignment()?;
            left = Expr::Pipe(Box::new(left), Box::new(right));
        }

        Ok(left)
    }

    /// Assignment: `x = value`, `x += 1`
    fn parse_assignment(&mut self) -> Result<Expr, ParseError> {
        let expr = self.parse_or()?;

        match self.peek() {
            TokenKind::Eq => {
                self.advance();
                self.skip_newlines();
                let value = self.parse_expr()?;
                Ok(Expr::Assign(Box::new(expr), AssignOp::Set, Box::new(value)))
            }
            TokenKind::PlusEq => {
                self.advance();
                let value = self.parse_expr()?;
                Ok(Expr::Assign(Box::new(expr), AssignOp::AddAssign, Box::new(value)))
            }
            TokenKind::MinusEq => {
                self.advance();
                let value = self.parse_expr()?;
                Ok(Expr::Assign(Box::new(expr), AssignOp::SubAssign, Box::new(value)))
            }
            _ => Ok(expr),
        }
    }

    /// `||`
    fn parse_or(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_and()?;
        while self.check(&TokenKind::Or) {
            self.advance();
            let right = self.parse_and()?;
            left = Expr::BinOp(Box::new(left), BinOp::Or, Box::new(right));
        }
        Ok(left)
    }

    /// `&&`
    fn parse_and(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_comparison()?;
        while self.check(&TokenKind::And) {
            self.advance();
            let right = self.parse_comparison()?;
            left = Expr::BinOp(Box::new(left), BinOp::And, Box::new(right));
        }
        Ok(left)
    }

    /// `==`, `!=`, `<`, `>`, `<=`, `>=`
    fn parse_comparison(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_additive()?;
        loop {
            let op = match self.peek() {
                TokenKind::EqEq => BinOp::Eq,
                TokenKind::Neq => BinOp::Neq,
                TokenKind::Lt => BinOp::Lt,
                TokenKind::Gt => BinOp::Gt,
                TokenKind::Lte => BinOp::Lte,
                TokenKind::Gte => BinOp::Gte,
                _ => break,
            };
            self.advance();
            let right = self.parse_additive()?;
            left = Expr::BinOp(Box::new(left), op, Box::new(right));
        }
        Ok(left)
    }

    /// `+`, `-`
    fn parse_additive(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_multiplicative()?;
        loop {
            let op = match self.peek() {
                TokenKind::Plus => BinOp::Add,
                TokenKind::Minus => BinOp::Sub,
                _ => break,
            };
            self.advance();
            let right = self.parse_multiplicative()?;
            left = Expr::BinOp(Box::new(left), op, Box::new(right));
        }
        Ok(left)
    }

    /// `*`, `/`, `%`
    fn parse_multiplicative(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_unary()?;
        loop {
            let op = match self.peek() {
                TokenKind::Star => BinOp::Mul,
                TokenKind::Slash => BinOp::Div,
                TokenKind::Percent => BinOp::Mod,
                _ => break,
            };
            self.advance();
            let right = self.parse_unary()?;
            left = Expr::BinOp(Box::new(left), op, Box::new(right));
        }
        Ok(left)
    }

    /// `-x`, `!flag`
    fn parse_unary(&mut self) -> Result<Expr, ParseError> {
        match self.peek() {
            TokenKind::Minus => {
                self.advance();
                let expr = self.parse_unary()?;
                Ok(Expr::UnaryOp(UnaryOp::Neg, Box::new(expr)))
            }
            TokenKind::Not => {
                self.advance();
                let expr = self.parse_unary()?;
                Ok(Expr::UnaryOp(UnaryOp::Not, Box::new(expr)))
            }
            _ => self.parse_postfix(),
        }
    }

    /// Dot access: `user.name`, function calls: `clamp(a, b)`
    fn parse_postfix(&mut self) -> Result<Expr, ParseError> {
        let mut expr = self.parse_primary()?;

        loop {
            match self.peek() {
                TokenKind::Dot => {
                    self.advance();
                    let field = self.expect_ident()?;
                    expr = Expr::DotAccess(Box::new(expr), field);
                }
                TokenKind::LBracket => {
                    self.advance(); // consume [
                    let index = self.parse_expr()?;
                    self.expect(&TokenKind::RBracket)?;
                    expr = Expr::Index(Box::new(expr), Box::new(index));
                }
                _ => break,
            }
        }

        Ok(expr)
    }

    /// Primary expressions: literals, identifiers, containers, etc.
    /// Parse a string literal, handling interpolation: `"hello {name}, you are {age}"`
    /// Lexer emits: StringFragment("hello ") → StringInterp → [name tokens] → RBrace → StringFragment(", you are ") → StringInterp → [age tokens] → RBrace → StringFragment/StringEnd
    fn parse_string_lit(&mut self) -> Result<Expr, ParseError> {
        let mut segments = Vec::new();

        loop {
            match self.peek().clone() {
                TokenKind::StringFragment(text) => {
                    self.advance();
                    segments.push(StringSegment::Literal(text));
                }
                TokenKind::StringEnd => {
                    self.advance();
                    // Empty string or end after interpolation
                    break;
                }
                TokenKind::StringInterp => {
                    self.advance(); // consume the { marker
                    let expr = self.parse_expr()?;
                    segments.push(StringSegment::Interpolation(Box::new(expr)));
                    // Consume the closing } — lexer emits RBrace and decrements interp_depth,
                    // then next_token transitions back to string mode automatically.
                    if self.check(&TokenKind::RBrace) {
                        self.advance();
                    }
                }
                _ => break,
            }
        }

        if segments.is_empty() {
            segments.push(StringSegment::Literal(String::new()));
        }

        Ok(Expr::StringLit(StringLit { segments }))
    }

    fn parse_primary(&mut self) -> Result<Expr, ParseError> {
        match self.peek().clone() {
            TokenKind::Int(n) => {
                self.advance();
                Ok(Expr::IntLit(n))
            }
            TokenKind::Float(n) => {
                self.advance();
                Ok(Expr::FloatLit(n))
            }
            TokenKind::True => {
                self.advance();
                Ok(Expr::BoolLit(true))
            }
            TokenKind::False => {
                self.advance();
                Ok(Expr::BoolLit(false))
            }
            TokenKind::StringFragment(_) | TokenKind::StringEnd | TokenKind::StringInterp => {
                self.parse_string_lit()
            }

            // Containers
            TokenKind::Column => self.parse_container(ContainerKind::Column),
            TokenKind::Row => self.parse_container(ContainerKind::Row),
            TokenKind::Stack => self.parse_container(ContainerKind::Stack),
            TokenKind::Panel => self.parse_container_with_props(ContainerKind::Panel),
            TokenKind::Scene => self.parse_container_with_props(ContainerKind::Scene),

            // When conditional
            TokenKind::When => {
                self.advance();
                let cond = self.parse_comparison()?;
                self.expect(&TokenKind::Arrow)?;
                self.skip_newlines();
                let body = self.parse_expr()?;
                Ok(Expr::When(Box::new(cond), Box::new(body)))
            }

            // Match
            TokenKind::Match => {
                self.advance();
                let scrutinee = self.parse_primary()?;
                self.skip_newlines();
                let mut arms = Vec::new();
                while !self.is_at_end()
                    && !matches!(self.peek(), TokenKind::Let | TokenKind::View | TokenKind::On | TokenKind::Effect)
                {
                    self.skip_newlines();
                    if self.is_at_end() || matches!(self.peek(), TokenKind::Let | TokenKind::View | TokenKind::On | TokenKind::Effect) {
                        break;
                    }
                    let pattern = self.parse_pattern()?;
                    self.expect(&TokenKind::Arrow)?;
                    self.skip_newlines();
                    let body = self.parse_expr()?;
                    arms.push(MatchArm { pattern, body });
                    self.skip_newlines();
                }
                Ok(Expr::Match(Box::new(scrutinee), arms))
            }

            // If-then-else
            TokenKind::If => {
                self.advance();
                let cond = self.parse_expr()?;
                self.expect(&TokenKind::Then)?;
                let then_branch = self.parse_expr()?;
                self.expect(&TokenKind::Else)?;
                let else_branch = self.parse_expr()?;
                Ok(Expr::If(Box::new(cond), Box::new(then_branch), Box::new(else_branch)))
            }

            // Perform effect
            TokenKind::Perform => {
                self.advance();
                let name = self.expect_ident()?;
                let args = if self.check(&TokenKind::LParen) {
                    self.parse_call_args()?
                } else {
                    Vec::new()
                };
                Ok(Expr::Perform(name, args))
            }

            // For-in: `for item in items -> body` or `for item, idx in items -> body`
            TokenKind::For => {
                self.advance();
                let item = self.expect_ident()?;

                // Optional index: `for item, idx in ...`
                let index = if self.check(&TokenKind::Comma) {
                    self.advance();
                    Some(self.expect_ident()?)
                } else {
                    None
                };

                self.expect(&TokenKind::In)?;
                let iter_expr = self.parse_comparison()?;
                self.expect(&TokenKind::Arrow)?;
                self.skip_newlines();
                let body = self.parse_expr()?;

                Ok(Expr::ForIn {
                    item,
                    index,
                    iter: Box::new(iter_expr),
                    body: Box::new(body),
                })
            }

            // Navigate expression: `navigate "/path"`
            TokenKind::Navigate => {
                self.advance();
                let path_expr = self.parse_primary()?;
                Ok(Expr::Call("navigate".to_string(), vec![path_expr]))
            }

            // Spring expression: `spring(value, stiffness, damping)`
            TokenKind::Spring => {
                self.advance();
                if self.check(&TokenKind::LParen) {
                    let args = self.parse_call_args()?;
                    Ok(Expr::Call("spring".to_string(), args))
                } else {
                    // Just `spring` as an identifier
                    Ok(Expr::Ident("spring".to_string()))
                }
            }

            // Stream from
            TokenKind::Stream => {
                self.advance();
                self.expect(&TokenKind::From)?;
                let source = self.expect_ident()?;
                // Allow dotted source: `button.click`
                let mut full_source = source;
                while self.check(&TokenKind::Dot) {
                    self.advance();
                    let next = self.expect_ident()?;
                    full_source = format!("{full_source}.{next}");
                }
                Ok(Expr::StreamFrom { source: full_source, mode: None })
            }

            // Record: `{ key: value }`
            TokenKind::LBrace => {
                self.advance();
                self.skip_newlines();
                let mut fields = Vec::new();
                while !self.check(&TokenKind::RBrace) && !self.is_at_end() {
                    self.skip_newlines();
                    let key = self.expect_ident()?;
                    self.expect(&TokenKind::Colon)?;
                    self.skip_newlines();
                    let val = self.parse_expr()?;
                    fields.push((key, val));
                    self.skip_newlines();
                    if self.check(&TokenKind::Comma) {
                        self.advance();
                    }
                    self.skip_newlines();
                }
                self.expect(&TokenKind::RBrace)?;
                Ok(Expr::Record(fields))
            }

            // List: `[a, b, c]`
            TokenKind::LBracket => {
                self.advance();
                self.skip_newlines();
                let mut items = Vec::new();
                while !self.check(&TokenKind::RBracket) && !self.is_at_end() {
                    self.skip_newlines();
                    items.push(self.parse_expr()?);
                    self.skip_newlines();
                    if self.check(&TokenKind::Comma) {
                        self.advance();
                    }
                    self.skip_newlines();
                }
                self.expect(&TokenKind::RBracket)?;
                Ok(Expr::List(items))
            }

            // Parenthesized expression or lambda
            TokenKind::LParen => {
                self.advance();
                // Check for lambda: `(x -> x * 2)` or `(x, y -> x + y)`
                let expr = self.parse_expr()?;
                if self.check(&TokenKind::Arrow) {
                    // This is a lambda
                    let params = vec![self.expr_to_ident(&expr)?];
                    self.advance(); // consume ->
                    let body = self.parse_expr()?;
                    self.expect(&TokenKind::RParen)?;
                    Ok(Expr::Lambda(params, Box::new(body)))
                } else if self.check(&TokenKind::Comma) {
                    // Could be multi-param lambda or tuple — check ahead
                    let mut items = vec![expr];
                    while self.check(&TokenKind::Comma) {
                        self.advance();
                        items.push(self.parse_expr()?);
                    }
                    if self.check(&TokenKind::Arrow) {
                        // Multi-param lambda
                        let mut params = Vec::new();
                        for item in &items {
                            params.push(self.expr_to_ident(item)?);
                        }
                        self.advance(); // ->
                        let body = self.parse_expr()?;
                        self.expect(&TokenKind::RParen)?;
                        Ok(Expr::Lambda(params, Box::new(body)))
                    } else {
                        // Just a parenthesized expression (take first)
                        self.expect(&TokenKind::RParen)?;
                        Ok(items.into_iter().next().unwrap())
                    }
                } else {
                    self.expect(&TokenKind::RParen)?;
                    Ok(expr)
                }
            }

            // Animate modifier (used in pipe context)
            TokenKind::Animate => {
                self.advance();
                let name = self.expect_ident()?;
                // Parse optional duration: `200ms`
                let mut args = Vec::new();
                if let TokenKind::Int(n) = self.peek().clone() {
                    self.advance();
                    // Check for unit suffix (we treat `ms` as part of value for now)
                    args.push(Expr::IntLit(n));
                }
                Ok(Expr::Call(format!("animate_{name}"), args))
            }

            // Identifier — variable reference, element, or function call
            TokenKind::Ident(name) => {
                let name = name.clone();
                self.advance();

                // UI element with any arg (string, ident, parenthesized expr, or props-only)
                if is_ui_element(&name) {
                    let next = self.peek().clone();
                    let looks_like_element = matches!(
                        next,
                        TokenKind::StringFragment(_) | TokenKind::StringEnd | TokenKind::StringInterp
                        | TokenKind::LBrace | TokenKind::LParen
                    ) || matches!(next, TokenKind::Ident(ref n) if !is_declaration_keyword(n));

                    if looks_like_element {
                        let fallback = name.clone();
                        match self.parse_element(name)? {
                            Some(el) => Ok(el),
                            None => Ok(Expr::Ident(fallback)),
                        }
                    } else {
                        Ok(Expr::Ident(name))
                    }
                }
                // Function call: `name(args)` — only for non-element idents
                else if self.check(&TokenKind::LParen) {
                    let args = self.parse_call_args()?;
                    Ok(Expr::Call(name, args))
                }
                // Element with string arg: `text "hello"` — fallback for non-is_ui_element tags
                else if matches!(self.peek(), TokenKind::StringFragment(_)) {
                    let fallback = name.clone();
                    match self.parse_element(name)? {
                        Some(el) => Ok(el),
                        None => Ok(Expr::Ident(fallback)),
                    }
                }
                else {
                    Ok(Expr::Ident(name))
                }
            }

            _ => Err(self.error(format!("unexpected token: {:?}", self.peek()))),
        }
    }

    fn parse_element(&mut self, tag: String) -> Result<Option<Expr>, ParseError> {
        let mut args = Vec::new();
        let mut props = Vec::new();
        let mut modifiers = Vec::new();

        // Parse string, ident, or parenthesized expression args
        loop {
            match self.peek().clone() {
                TokenKind::StringFragment(_) | TokenKind::StringEnd | TokenKind::StringInterp => {
                    args.push(self.parse_string_lit()?);
                }
                TokenKind::Ident(name) if !is_declaration_keyword(&name) => {
                    // Only consume if it looks like an element argument
                    self.advance();
                    args.push(Expr::Ident(name));
                }
                TokenKind::LParen => {
                    // Parenthesized expression arg: button (if cond then "A" else "B")
                    self.advance(); // consume (
                    let expr = self.parse_expr()?;
                    self.expect(&TokenKind::RParen)?;
                    args.push(expr);
                }
                _ => break,
            }
        }

        // Parse props: `{ click: handler, class: "foo" }`
        if self.check(&TokenKind::LBrace) {
            self.advance();
            self.skip_newlines();
            while !self.check(&TokenKind::RBrace) && !self.is_at_end() {
                self.skip_newlines();
                let key = self.expect_ident()?;
                self.expect(&TokenKind::Colon)?;
                self.skip_newlines();
                let val = self.parse_expr()?;
                props.push((key, val));
                self.skip_newlines();
                if self.check(&TokenKind::Comma) {
                    self.advance();
                }
                self.skip_newlines();
            }
            self.expect(&TokenKind::RBrace)?;
        }

        // Parse modifiers: `| animate fade-in 200ms`
        while self.check(&TokenKind::Pipe) {
            self.advance();
            let name = self.expect_ident()?;
            let mut mod_args = Vec::new();
            while matches!(self.peek(), TokenKind::Ident(_) | TokenKind::Int(_) | TokenKind::Float(_)) {
                mod_args.push(self.parse_primary()?);
            }
            modifiers.push(Modifier { name, args: mod_args });
        }

        if args.is_empty() && props.is_empty() && modifiers.is_empty() {
            return Ok(None);
        }

        Ok(Some(Expr::Element(Element {
            tag,
            args,
            props,
            modifiers,
        })))
    }

    fn parse_container(&mut self, kind: ContainerKind) -> Result<Expr, ParseError> {
        self.advance(); // consume container keyword
        self.expect(&TokenKind::LBracket)?;
        self.skip_newlines();

        let mut children = Vec::new();
        while !self.check(&TokenKind::RBracket) && !self.is_at_end() {
            self.skip_newlines();
            if self.check(&TokenKind::RBracket) { break; }
            children.push(self.parse_expr()?);
            self.skip_newlines();
            if self.check(&TokenKind::Comma) {
                self.advance();
            }
            self.skip_newlines();
        }

        self.expect(&TokenKind::RBracket)?;
        Ok(Expr::Container(Container {
            kind,
            children,
            props: Vec::new(),
        }))
    }

    fn parse_container_with_props(&mut self, kind: ContainerKind) -> Result<Expr, ParseError> {
        self.advance(); // consume container keyword

        // Optional props: `panel { x: panel_x }`
        let mut props = Vec::new();
        if self.check(&TokenKind::LBrace) {
            self.advance();
            self.skip_newlines();
            while !self.check(&TokenKind::RBrace) && !self.is_at_end() {
                self.skip_newlines();
                let key = self.expect_ident()?;
                self.expect(&TokenKind::Colon)?;
                let val = self.parse_expr()?;
                props.push((key, val));
                self.skip_newlines();
                if self.check(&TokenKind::Comma) {
                    self.advance();
                }
                self.skip_newlines();
            }
            self.expect(&TokenKind::RBrace)?;
        }

        self.expect(&TokenKind::LBracket)?;
        self.skip_newlines();

        let mut children = Vec::new();
        while !self.check(&TokenKind::RBracket) && !self.is_at_end() {
            self.skip_newlines();
            if self.check(&TokenKind::RBracket) { break; }
            children.push(self.parse_expr()?);
            self.skip_newlines();
            if self.check(&TokenKind::Comma) {
                self.advance();
            }
            self.skip_newlines();
        }

        self.expect(&TokenKind::RBracket)?;
        Ok(Expr::Container(Container { kind, children, props }))
    }

    fn parse_call_args(&mut self) -> Result<Vec<Expr>, ParseError> {
        self.expect(&TokenKind::LParen)?;
        let mut args = Vec::new();
        while !self.check(&TokenKind::RParen) && !self.is_at_end() {
            args.push(self.parse_expr()?);
            if self.check(&TokenKind::Comma) {
                self.advance();
            }
        }
        self.expect(&TokenKind::RParen)?;
        Ok(args)
    }

    fn parse_pattern(&mut self) -> Result<Pattern, ParseError> {
        match self.peek().clone() {
            TokenKind::Ident(name) => {
                self.advance();
                if self.check(&TokenKind::LParen) {
                    // Constructor pattern: `Ok(value)`
                    self.advance();
                    let mut fields = Vec::new();
                    while !self.check(&TokenKind::RParen) && !self.is_at_end() {
                        fields.push(self.parse_pattern()?);
                        if self.check(&TokenKind::Comma) {
                            self.advance();
                        }
                    }
                    self.expect(&TokenKind::RParen)?;
                    Ok(Pattern::Constructor(name, fields))
                } else {
                    Ok(Pattern::Ident(name))
                }
            }
            TokenKind::Int(n) => {
                self.advance();
                Ok(Pattern::Literal(Expr::IntLit(n)))
            }
            TokenKind::StringFragment(s) => {
                self.advance();
                Ok(Pattern::Literal(Expr::StringLit(StringLit {
                    segments: vec![StringSegment::Literal(s)],
                })))
            }
            _ => Err(self.error(format!("expected pattern, got {:?}", self.peek()))),
        }
    }

    fn expr_to_ident(&self, expr: &Expr) -> Result<String, ParseError> {
        match expr {
            Expr::Ident(name) => Ok(name.clone()),
            _ => Err(self.error("expected identifier in lambda parameter".into())),
        }
    }
}

fn is_ui_element(name: &str) -> bool {
    matches!(
        name,
        "text" | "button" | "input" | "image" | "avatar" | "icon"
            | "link" | "label" | "badge" | "chip" | "card"
            | "header" | "footer" | "nav" | "section" | "div"
            | "spinner" | "skeleton"
            | "circle" | "rect" | "line"  // Physics scene elements
    )
}

fn is_declaration_keyword(name: &str) -> bool {
    matches!(name, "let" | "view" | "effect" | "on" | "handle")
}

#[derive(Debug)]
pub struct ParseError {
    pub message: String,
    pub line: usize,
    pub col: usize,
}

impl std::fmt::Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Parse error at line {}:{}: {}", self.line, self.col, self.message)
    }
}

impl std::error::Error for ParseError {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::Lexer;

    fn parse(src: &str) -> Program {
        let mut lexer = Lexer::new(src);
        let tokens = lexer.tokenize();
        let mut parser = Parser::new(tokens);
        parser.parse_program().expect("parse failed")
    }

    #[test]
    fn test_let_int() {
        let prog = parse("let count = 0");
        assert_eq!(prog.declarations.len(), 1);
        match &prog.declarations[0] {
            Declaration::Let(decl) => {
                assert_eq!(decl.name, "count");
                assert!(matches!(decl.value, Expr::IntLit(0)));
            }
            _ => panic!("expected let"),
        }
    }

    #[test]
    fn test_let_binop() {
        let prog = parse("let doubled = count * 2");
        match &prog.declarations[0] {
            Declaration::Let(decl) => {
                assert_eq!(decl.name, "doubled");
                match &decl.value {
                    Expr::BinOp(_, BinOp::Mul, _) => {}
                    other => panic!("expected BinOp(Mul), got {other:?}"),
                }
            }
            _ => panic!("expected let"),
        }
    }

    #[test]
    fn test_view_simple() {
        let prog = parse(
            r#"view counter =
  column [
    text "hello"
  ]"#
        );
        match &prog.declarations[0] {
            Declaration::View(v) => {
                assert_eq!(v.name, "counter");
                match &v.body {
                    Expr::Container(c) => {
                        assert!(matches!(c.kind, ContainerKind::Column));
                        assert_eq!(c.children.len(), 1);
                    }
                    other => panic!("expected Container, got {other:?}"),
                }
            }
            _ => panic!("expected view"),
        }
    }

    #[test]
    fn test_when_expr() {
        let prog = parse(
            r#"view test =
  column [
    when count > 10 ->
      text "big"
  ]"#
        );
        match &prog.declarations[0] {
            Declaration::View(v) => {
                match &v.body {
                    Expr::Container(c) => {
                        assert!(matches!(&c.children[0], Expr::When(_, _)));
                    }
                    other => panic!("expected Container, got {other:?}"),
                }
            }
            _ => panic!("expected view"),
        }
    }

    #[test]
    fn test_on_handler() {
        let prog = parse("on toggle_sidebar ->\n  count += 1");
        match &prog.declarations[0] {
            Declaration::OnHandler(h) => {
                assert_eq!(h.event, "toggle_sidebar");
                assert!(matches!(&h.body, Expr::Assign(_, AssignOp::AddAssign, _)));
            }
            _ => panic!("expected on handler"),
        }
    }

    #[test]
    fn test_full_counter() {
        let prog = parse(
            r#"let count = 0
let doubled = count * 2
let label = "hello"

view counter =
  column [
    text label
    button "+" { click: count += 1 }
    when count > 10 ->
      text "big"
  ]"#
        );
        assert_eq!(prog.declarations.len(), 4); // 3 lets + 1 view
    }

    #[test]
    fn test_stream_decl() {
        let prog = parse(r#"stream main on "ws://localhost:9100" { mode: signal }"#);
        match &prog.declarations[0] {
            Declaration::Stream(s) => {
                assert_eq!(s.view_name, "main");
                assert_eq!(s.mode, StreamMode::Signal);
            }
            other => panic!("expected Stream, got {other:?}"),
        }
    }

    #[test]
    fn test_stream_decl_pixel_mode() {
        let prog = parse(r#"stream main on "ws://localhost:9100" { mode: pixel }"#);
        match &prog.declarations[0] {
            Declaration::Stream(s) => {
                assert_eq!(s.mode, StreamMode::Pixel);
            }
            other => panic!("expected Stream, got {other:?}"),
        }
    }

    #[test]
    fn test_stream_decl_default_mode() {
        let prog = parse(r#"stream main on "ws://localhost:9100""#);
        match &prog.declarations[0] {
            Declaration::Stream(s) => {
                assert_eq!(s.mode, StreamMode::Signal);
            }
            other => panic!("expected Stream, got {other:?}"),
        }
    }
}