sharp 0.1.0

// Parser for Sharp language

use crate::lexer::Lexer;
use crate::lexer::Token;
use crate::ast::*;
use std::rc::Rc;

/// The main parser struct.
pub struct Parser<'a> {
    lexer: Lexer<'a>,
    /// Lookahead token (the current token).
    current: Option<(Token, std::ops::Range<usize>)>,
}

impl<'a> Parser<'a> {
    /// Create a new parser from source code.
    pub fn new(source: &'a str) -> Self {
        let mut lexer = Lexer::new(source);
        let first = lexer.next();
        Self { lexer, current: first }
    }

    /// Advance to the next token.
    fn advance(&mut self) {
        self.current = self.lexer.next();
    }

    /// Peek the current token.
    fn peek(&self) -> Option<&Token> {
        self.current.as_ref().map(|(tok, _)| tok)
    }

    /// Consume the current token if it matches the expected one.
    fn expect(&mut self, expected: Token) -> Result<std::ops::Range<usize>, String> {
        match &self.current {
            Some((tok, span)) if std::mem::discriminant(tok) == std::mem::discriminant(&expected) => {
                let span = span.clone();
                self.advance();
                Ok(span)
            }
            Some((tok, _)) => Err(format!("Expected {:?}, found {:?}", expected, tok)),
            None => Err("Unexpected end of input".to_string()),
        }
    }

    /// Skip newlines, indents, and dedents
    fn skip_newlines(&mut self) {
        while matches!(self.peek(), Some(Token::Newline | Token::Indent | Token::Dedent)) {
            self.advance();
        }
    }

    /// Parse the entire program.
    pub fn parse_program(&mut self) -> Result<Program, String> {
        let mut items = Vec::new();
        self.skip_newlines();
        while self.peek().is_some() {
            match self.peek() {
                Some(Token::Def) => items.push(TopLevel::Func(self.parse_func()?)),
                Some(Token::Struct) => items.push(TopLevel::Struct(self.parse_struct()?)),
                Some(Token::Class) => items.push(TopLevel::Class(self.parse_class()?)),
                Some(Token::Const) => items.push(TopLevel::Const(self.parse_const()?)),
                Some(Token::Import) => items.push(TopLevel::Import(self.parse_import()?)),
                Some(Token::Impl) => items.push(TopLevel::Impl(self.parse_impl()?)),
                Some(Token::Enum) => items.push(TopLevel::Enum(self.parse_enum()?)),
                None => break,
                _ => return Err(format!("Unexpected top-level token: {:?}", self.peek())),
            }
            self.skip_newlines();
        }
        Ok(Program { items })
    }

    /// Parse import declaration
    fn parse_import(&mut self) -> Result<ImportDecl, String> {
        self.expect(Token::Import)?;
        let mut path = vec![self.parse_ident()?];
        while matches!(self.peek(), Some(Token::Dot)) {
            self.advance();
            path.push(self.parse_ident()?);
        }
        let alias = if matches!(self.peek(), Some(Token::As)) {
            self.advance();
            Some(self.parse_ident()?)
        } else {
            None
        };
        Ok(ImportDecl { module_path: path, alias })
    }

    /// Parse const declaration
    fn parse_const(&mut self) -> Result<ConstDecl, String> {
        self.expect(Token::Const)?;
        let name = self.parse_ident()?;
        let ty = if matches!(self.peek(), Some(Token::Colon)) {
            self.advance();
            Some(self.parse_type()?)
        } else {
            None
        };
        self.expect(Token::Eq)?;
        let value = self.parse_expr()?;
        Ok(ConstDecl { name, ty, value })
    }

    /// Parse struct declaration
    fn parse_struct(&mut self) -> Result<StructDecl, String> {
        self.expect(Token::Struct)?;
        let name = self.parse_ident()?;
        let type_params = self.parse_type_params()?;
        let fields = self.parse_field_list()?;
        Ok(StructDecl { name, type_params, fields })
    }

    /// Parse class declaration
    fn parse_class(&mut self) -> Result<ClassDecl, String> {
        self.expect(Token::Class)?;
        let name = self.parse_ident()?;
        let base = if matches!(self.peek(), Some(Token::Colon)) {
            self.advance();
            Some(self.parse_ident()?)
        } else {
            None
        };
        let body = self.parse_block()?;
        Ok(ClassDecl { name, base, body })
    }

    /// Parse field list for struct
    fn parse_field_list(&mut self) -> Result<Vec<Field>, String> {
        self.expect(Token::LBrace)?;
        self.skip_newlines();
        let mut fields = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            let name = self.parse_ident()?;
            self.expect(Token::Colon)?;
            let ty = self.parse_type()?;
            let default = if matches!(self.peek(), Some(Token::Eq)) {
                self.advance();
                Some(self.parse_expr()?)
            } else {
                None
            };
            fields.push(Field { name, ty, default });
            if matches!(self.peek(), Some(Token::Comma)) {
                self.advance();
            }
            self.skip_newlines();
        }
        self.expect(Token::RBrace)?;
        Ok(fields)
    }

    /// Parse type parameters
    fn parse_type_params(&mut self) -> Result<Vec<String>, String> {
        if !matches!(self.peek(), Some(Token::Lt)) {
            return Ok(Vec::new());
        }
        self.advance();
        let mut params = vec![self.parse_ident()?];
        while matches!(self.peek(), Some(Token::Comma)) {
            self.advance();
            params.push(self.parse_ident()?);
        }
        self.expect(Token::Gt)?;
        Ok(params)
    }

    /// Parse impl block
    fn parse_impl(&mut self) -> Result<ImplBlock, String> {
        self.expect(Token::Impl)?;
        let ty = self.parse_ident()?;
        self.expect(Token::LBrace)?;
        self.skip_newlines();
        
        let mut methods = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            if matches!(self.peek(), Some(Token::Def)) {
                methods.push(self.parse_func()?);
            } else {
                return Err("Expected function definition in impl block".to_string());
            }
            self.skip_newlines();
        }
        
        self.expect(Token::RBrace)?;
        Ok(ImplBlock { ty, methods })
    }

    /// Parse enum declaration
    fn parse_enum(&mut self) -> Result<EnumDecl, String> {
        self.expect(Token::Enum)?;
        let name = self.parse_ident()?;
        self.expect(Token::LBrace)?;
        self.skip_newlines();
        
        let mut variants = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            let variant_name = self.parse_ident()?;
            let fields = if matches!(self.peek(), Some(Token::LParen)) {
                self.advance();
                let mut types = Vec::new();
                while !matches!(self.peek(), Some(Token::RParen)) {
                    types.push(self.parse_type()?);
                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                    }
                }
                self.expect(Token::RParen)?;
                types
            } else {
                Vec::new()
            };
            
            variants.push(EnumVariant { name: variant_name, fields });
            
            if matches!(self.peek(), Some(Token::Comma)) {
                self.advance();
            }
            self.skip_newlines();
        }
        
        self.expect(Token::RBrace)?;
        Ok(EnumDecl { name, variants })
    }

    /// Parse a function declaration.
    fn parse_func(&mut self) -> Result<FuncDecl, String> {
        self.expect(Token::Def)?;
        let name = self.parse_ident()?;
        let type_params = self.parse_type_params()?;
        self.expect(Token::LParen)?;
        let params = self.parse_param_list()?;
        self.expect(Token::RParen)?;
        
        let ret_type = if matches!(self.peek(), Some(Token::Arrow)) {
            self.advance();
            Some(self.parse_type()?)
        } else {
            None
        };
        
        let body = self.parse_block()?;
        Ok(FuncDecl { name, type_params, params, ret_type, body })
    }

    /// Parse parameter list
    fn parse_param_list(&mut self) -> Result<Vec<Param>, String> {
        let mut params = Vec::new();
        if matches!(self.peek(), Some(Token::RParen)) {
            return Ok(params);
        }
        
        loop {
            let mutable = if matches!(self.peek(), Some(Token::Mut)) {
                self.advance();
                true
            } else {
                false
            };
            
            let name = self.parse_ident()?;
            let ty = if matches!(self.peek(), Some(Token::Colon)) {
                self.advance();
                Some(self.parse_type()?)
            } else {
                None
            };
            let default = if matches!(self.peek(), Some(Token::Eq)) {
                self.advance();
                Some(self.parse_expr()?)
            } else {
                None
            };
            params.push(Param { name, ty, mutable, default });
            
            if !matches!(self.peek(), Some(Token::Comma)) {
                break;
            }
            self.advance();
        }
        Ok(params)
    }

    /// Parse a block (expects ':' INDENT ... DEDENT or '{' ... '}').
    fn parse_block(&mut self) -> Result<Block, String> {
        if matches!(self.peek(), Some(Token::Colon)) {
            self.advance();
            self.skip_newlines();
            self.expect(Token::Indent)?;
            let mut stmts = Vec::new();
            while !matches!(self.peek(), Some(Token::Dedent)) {
                stmts.push(self.parse_stmt()?);
                self.skip_newlines();
            }
            self.expect(Token::Dedent)?;
            Ok(Block { statements: stmts })
        } else {
            self.expect(Token::LBrace)?;
            self.skip_newlines();
            let mut stmts = Vec::new();
            while !matches!(self.peek(), Some(Token::RBrace)) {
                stmts.push(self.parse_stmt()?);
                self.skip_newlines();
            }
            self.expect(Token::RBrace)?;
            Ok(Block { statements: stmts })
        }
    }

    /// Parse a statement
    fn parse_stmt(&mut self) -> Result<Stmt, String> {
        match self.peek() {
            Some(Token::Return) => {
                self.advance();
                let expr = if matches!(self.peek(), Some(Token::Newline | Token::Semi | Token::RBrace)) {
                    None
                } else {
                    Some(self.parse_expr()?)
                };
                Ok(Stmt::Return(expr))
            }
            Some(Token::If) => Ok(Stmt::If(self.parse_if()?)),
            Some(Token::While) => Ok(Stmt::While(self.parse_while()?)),
            Some(Token::For) => Ok(Stmt::For(self.parse_for()?)),
            Some(Token::Match) => Ok(Stmt::Match(self.parse_match()?)),
            Some(Token::Ident(_)) => {
                // Could be var decl or expression
                let checkpoint = self.current.clone();
                let name = self.parse_ident()?;
                
                if matches!(self.peek(), Some(Token::Colon)) {
                    // Variable declaration: name: type = expr
                    self.advance();
                    let ty = Some(self.parse_type()?);
                    self.expect(Token::Eq)?;
                    let init = self.parse_expr()?;
                    Ok(Stmt::VarDecl(VarDecl { name, ty, init }))
                } else if matches!(self.peek(), Some(Token::Eq)) {
                    // Variable declaration without type: name = expr
                    self.advance();
                    let init = self.parse_expr()?;
                    Ok(Stmt::VarDecl(VarDecl { name, ty: None, init }))
                } else {
                    // Expression statement - restore and re-parse
                    self.current = checkpoint;
                    Ok(Stmt::Expr(self.parse_expr()?))
                }
            }
            _ => Ok(Stmt::Expr(self.parse_expr()?)),
        }
    }

    /// Parse if statement
    fn parse_if(&mut self) -> Result<IfStmt, String> {
        self.expect(Token::If)?;
        let condition = self.parse_expr()?;
        let then_branch = self.parse_block()?;
        
        let mut elif_branches = Vec::new();
        while matches!(self.peek(), Some(Token::Elif)) {
            self.advance();
            let elif_cond = self.parse_expr()?;
            let elif_block = self.parse_block()?;
            elif_branches.push((elif_cond, elif_block));
        }
        
        let else_branch = if matches!(self.peek(), Some(Token::Else)) {
            self.advance();
            Some(self.parse_block()?)
        } else {
            None
        };
        
        Ok(IfStmt { condition, then_branch, elif_branches, else_branch })
    }

    /// Parse while statement
    fn parse_while(&mut self) -> Result<WhileStmt, String> {
        self.expect(Token::While)?;
        let condition = self.parse_expr()?;
        let body = self.parse_block()?;
        Ok(WhileStmt { condition, body })
    }

    /// Parse for statement
    fn parse_for(&mut self) -> Result<ForStmt, String> {
        self.expect(Token::For)?;
        // Support two forms:
        // 1) C-style: for (init; cond; update) { body }
        // 2) Range sugar: for i in start .. end { body }
        if matches!(self.peek(), Some(Token::LParen)) {
            self.advance();
            let init = if matches!(self.peek(), Some(Token::Semi)) { None } else { Some(Box::new(self.parse_stmt()?)) };
            self.expect(Token::Semi)?;
            let condition = self.parse_expr()?;
            self.expect(Token::Semi)?;
            let update = if matches!(self.peek(), Some(Token::RParen)) { None } else { Some(Box::new(self.parse_stmt()?)) };
            self.expect(Token::RParen)?;
            let body = self.parse_block()?;
            Ok(ForStmt { init, condition, update, body })
        } else {
            // Range sugar: for ident in start .. end { body }
            let var_name = self.parse_ident()?;
            self.expect(Token::In)?;
            let start_expr = self.parse_expr()?;
            self.expect(Token::DotDot)?;
            let end_expr = self.parse_expr()?;
            // Synthesize: init = var_name = start_expr;
            let init_stmt = Stmt::VarDecl(VarDecl { name: var_name.clone(), ty: None, init: start_expr });
            // cond = var_name < end_expr;
            let cond_expr = Expr::Binary(Box::new(Expr::Ident(var_name.clone())), BinOp::Lt, Box::new(end_expr));
            // update = var_name = var_name + 1;
            let one = Expr::Literal(Literal::Int(1));
            let add = Expr::Binary(Box::new(Expr::Ident(var_name.clone())), BinOp::Add, Box::new(one));
            let update_stmt = Stmt::VarDecl(VarDecl { name: var_name.clone(), ty: None, init: add });
            let body = self.parse_block()?;
            Ok(ForStmt { init: Some(Box::new(init_stmt)), condition: cond_expr, update: Some(Box::new(update_stmt)), body })
        }
    }

    /// Parse match statement
    fn parse_match(&mut self) -> Result<MatchStmt, String> {
        self.expect(Token::Match)?;
        let expr = self.parse_expr()?;
        self.expect(Token::LBrace)?;
        self.skip_newlines();
        
        let mut arms = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            let pattern = self.parse_pattern()?;
            self.expect(Token::Arrow)?;
            let block = self.parse_block()?;
            arms.push((pattern, block));
            self.skip_newlines();
        }
        self.expect(Token::RBrace)?;
        Ok(MatchStmt { expr, arms })
    }

    /// Parse pattern
    fn parse_pattern(&mut self) -> Result<Pattern, String> {
        match self.peek() {
            Some(Token::IntLit(n)) => {
                let n = *n;
                self.advance();
                Ok(Pattern::Lit(Literal::Int(n)))
            }
            Some(Token::Ident(name)) if name == "_" => {
                self.advance();
                Ok(Pattern::Wildcard)
            }
            Some(Token::Ident(_)) => {
                Ok(Pattern::Ident(self.parse_ident()?))
            }
            _ => Err("Expected pattern".to_string()),
        }
    }

    /// Parse an expression (with precedence climbing).
    fn parse_expr(&mut self) -> Result<Expr, String> {
        self.parse_assignment()
    }

    /// Parse assignment (right-associative)
    fn parse_assignment(&mut self) -> Result<Expr, String> {
        let lhs = self.parse_or()?;
        
        match self.peek() {
            Some(Token::Eq) => {
                self.advance();
                let rhs = self.parse_assignment()?;
                Ok(Expr::Assign(Box::new(lhs), Box::new(rhs)))
            }
            Some(Token::PlusEq) => {
                self.advance();
                let rhs = self.parse_assignment()?;
                Ok(Expr::CompoundAssign(Box::new(lhs), BinOp::Add, Box::new(rhs)))
            }
            Some(Token::MinusEq) => {
                self.advance();
                let rhs = self.parse_assignment()?;
                Ok(Expr::CompoundAssign(Box::new(lhs), BinOp::Sub, Box::new(rhs)))
            }
            Some(Token::StarEq) => {
                self.advance();
                let rhs = self.parse_assignment()?;
                Ok(Expr::CompoundAssign(Box::new(lhs), BinOp::Mul, Box::new(rhs)))
            }
            Some(Token::SlashEq) => {
                self.advance();
                let rhs = self.parse_assignment()?;
                Ok(Expr::CompoundAssign(Box::new(lhs), BinOp::Div, Box::new(rhs)))
            }
            _ => Ok(lhs),
        }
    }

    /// Parse logical or
    fn parse_or(&mut self) -> Result<Expr, String> {
        let mut lhs = self.parse_and()?;
        while matches!(self.peek(), Some(Token::Or)) {
            self.advance();
            let rhs = self.parse_and()?;
            lhs = Expr::Binary(Box::new(lhs), BinOp::Or, Box::new(rhs));
        }
        Ok(lhs)
    }

    /// Parse logical and
    fn parse_and(&mut self) -> Result<Expr, String> {
        let mut lhs = self.parse_comparison()?;
        while matches!(self.peek(), Some(Token::And)) {
            self.advance();
            let rhs = self.parse_comparison()?;
            lhs = Expr::Binary(Box::new(lhs), BinOp::And, Box::new(rhs));
        }
        Ok(lhs)
    }

    /// Parse comparison operators
    fn parse_comparison(&mut self) -> Result<Expr, String> {
        let mut lhs = self.parse_additive()?;
        
        loop {
            let op = match self.peek() {
                Some(Token::EqEq) => BinOp::Eq,
                Some(Token::NotEq) => BinOp::NotEq,
                Some(Token::Lt) => BinOp::Lt,
                Some(Token::Gt) => BinOp::Gt,
                Some(Token::LtEq) => BinOp::LtEq,
                Some(Token::GtEq) => BinOp::GtEq,
                _ => break,
            };
            self.advance();
            let rhs = self.parse_additive()?;
            lhs = Expr::Binary(Box::new(lhs), op, Box::new(rhs));
        }
        Ok(lhs)
    }

    /// Parse additive operators
    fn parse_additive(&mut self) -> Result<Expr, String> {
        let mut lhs = self.parse_multiplicative()?;
        
        loop {
            let op = match self.peek() {
                Some(Token::Plus) => BinOp::Add,
                Some(Token::Minus) => BinOp::Sub,
                _ => break,
            };
            self.advance();
            let rhs = self.parse_multiplicative()?;
            lhs = Expr::Binary(Box::new(lhs), op, Box::new(rhs));
        }
        Ok(lhs)
    }

    /// Parse multiplicative operators
    fn parse_multiplicative(&mut self) -> Result<Expr, String> {
        let mut lhs = self.parse_unary()?;
        
        loop {
            let op = match self.peek() {
                Some(Token::Star) => BinOp::Mul,
                Some(Token::Slash) => BinOp::Div,
                Some(Token::Percent) => BinOp::Mod,
                _ => break,
            };
            self.advance();
            let rhs = self.parse_unary()?;
            lhs = Expr::Binary(Box::new(lhs), op, Box::new(rhs));
        }
        Ok(lhs)
    }

    /// Parse unary operators
    fn parse_unary(&mut self) -> Result<Expr, String> {
        match self.peek() {
            Some(Token::Minus) => {
                self.advance();
                let expr = self.parse_unary()?;
                Ok(Expr::Unary(UnOp::Neg, Box::new(expr)))
            }
            Some(Token::Not) => {
                self.advance();
                let expr = self.parse_unary()?;
                Ok(Expr::Unary(UnOp::Not, Box::new(expr)))
            }
            _ => self.parse_postfix(),
        }
    }

    /// Parse postfix expressions (call, index, field access)
    fn parse_postfix(&mut self) -> Result<Expr, String> {
        let mut expr = self.parse_primary()?;
        
        loop {
            match self.peek() {
                Some(Token::LParen) => {
                    self.advance();
                    let args = self.parse_arg_list()?;
                    self.expect(Token::RParen)?;
                    expr = Expr::Call(Box::new(expr), args);
                }
                Some(Token::LBracket) => {
                    self.advance();
                    let index = self.parse_expr()?;
                    self.expect(Token::RBracket)?;
                    expr = Expr::Index(Box::new(expr), Box::new(index));
                }
                Some(Token::Dot) => {
                    self.advance();
                    let field = self.parse_ident()?;
                    
                    // Check if this is a method call
                    if matches!(self.peek(), Some(Token::LParen)) {
                        self.advance();
                        let args = self.parse_arg_list()?;
                        self.expect(Token::RParen)?;
                        expr = Expr::MethodCall(Box::new(expr), field, args);
                    } else {
                        expr = Expr::FieldAccess(Box::new(expr), field);
                    }
                }
                Some(Token::LBrace) => {
                    // This might be struct initialization
                    if let Expr::Ident(struct_name) = expr {
                        self.advance();
                        let mut fields = Vec::new();
                        
                        while !matches!(self.peek(), Some(Token::RBrace)) {
                            let field_name = self.parse_ident()?;
                            self.expect(Token::Colon)?;
                            let field_value = self.parse_expr()?;
                            fields.push((field_name, field_value));
                            
                            if matches!(self.peek(), Some(Token::Comma)) {
                                self.advance();
                            } else {
                                break;
                            }
                        }
                        
                        self.expect(Token::RBrace)?;
                        expr = Expr::StructInit(struct_name, fields);
                    } else {
                        break;
                    }
                }
                _ => break,
            }
        }
        Ok(expr)
    }

    /// Parse argument list
    fn parse_arg_list(&mut self) -> Result<Vec<Expr>, String> {
        let mut args = Vec::new();
        if matches!(self.peek(), Some(Token::RParen)) {
            return Ok(args);
        }
        
        loop {
            args.push(self.parse_expr()?);
            if !matches!(self.peek(), Some(Token::Comma)) {
                break;
            }
            self.advance();
        }
        Ok(args)
    }

    /// Parse primary expressions
    fn parse_primary(&mut self) -> Result<Expr, String> {
        match self.peek() {
            Some(Token::True) => { self.advance(); Ok(Expr::Literal(Literal::Bool(true))) }
            Some(Token::False) => { self.advance(); Ok(Expr::Literal(Literal::Bool(false))) }
            Some(Token::IntLit(n)) => {
                let n = *n;
                self.advance();
                Ok(Expr::Literal(Literal::Int(n)))
            }
            Some(Token::FloatLit(f)) => {
                let f = *f;
                self.advance();
                Ok(Expr::Literal(Literal::Float(f)))
            }
            Some(Token::StringLit(s)) => {
                let s = s.clone();
                self.advance();
                Ok(Expr::Literal(Literal::String(Rc::from(s.as_str()))))
            }
            Some(Token::Ident(_)) => {
                Ok(Expr::Ident(self.parse_ident()?))
            }
            Some(Token::LParen) => {
                self.advance();
                let expr = self.parse_expr()?;
                self.expect(Token::RParen)?;
                Ok(expr)
            }
            Some(Token::LBracket) => {
                self.advance();
                let mut elements = Vec::new();
                while !matches!(self.peek(), Some(Token::RBracket)) {
                    elements.push(self.parse_expr()?);
                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                    } else {
                        break;
                    }
                }
                self.expect(Token::RBracket)?;
                Ok(Expr::List(elements))
            }
            Some(Token::Lambda) => self.parse_lambda(),
            Some(Token::LBrace) => {
                let block = self.parse_block()?;
                Ok(Expr::Block(block))
            }
            _ => Err(format!("Unexpected token in expression: {:?}", self.peek())),
        }
    }

    /// Parse lambda expression
    fn parse_lambda(&mut self) -> Result<Expr, String> {
        self.expect(Token::Lambda)?;
        let params = if matches!(self.peek(), Some(Token::Colon)) {
            Vec::new()
        } else {
            self.parse_param_list()?
        };
        self.expect(Token::Colon)?;
        let body = self.parse_expr()?;
        Ok(Expr::Lambda(params, Box::new(body)))
    }

    /// Parse a type.
    fn parse_type(&mut self) -> Result<Type, String> {
        let base_type = match self.peek() {
            Some(Token::TypeInt) => { self.advance(); Type::Int }
            Some(Token::TypeFloat) => { self.advance(); Type::Float }
            Some(Token::TypeBool) => { self.advance(); Type::Bool }
            Some(Token::TypeString) => { self.advance(); Type::String }
            Some(Token::TypeVoid) => { self.advance(); Type::Void }
            Some(Token::TypeAuto) => { self.advance(); Type::Auto }
            Some(Token::LBracket) => {
                self.advance();
                let elem_type = self.parse_type()?;
                let size = if matches!(self.peek(), Some(Token::Semi)) {
                    self.advance();
                    match self.peek() {
                        Some(Token::IntLit(n)) => {
                            let n = *n as usize;
                            self.advance();
                            Some(n)
                        }
                        _ => return Err("Expected array size".to_string()),
                    }
                } else {
                    None
                };
                self.expect(Token::RBracket)?;
                Type::Array(Box::new(elem_type), size)
            }
            Some(Token::Ident(_)) => {
                let name = self.parse_ident()?;
                let type_args = if matches!(self.peek(), Some(Token::Lt)) {
                    self.advance();
                    let mut args = vec![self.parse_type()?];
                    while matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                        args.push(self.parse_type()?);
                    }
                    self.expect(Token::Gt)?;
                    args
                } else {
                    Vec::new()
                };
                Type::Named(name, type_args)
            }
            _ => return Err(format!("Unexpected token in type: {:?}", self.peek())),
        };
        
        Ok(base_type)
    }

    /// Parse an identifier
    fn parse_ident(&mut self) -> Result<String, String> {
        match self.current.clone() {
            Some((Token::Ident(name), _)) => {
                self.advance();
                Ok(name)
            }
            _ => Err(format!("Expected identifier, found {:?}", self.peek())),
        }
    }
}

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

    #[test]
    fn test_parse_simple_function() {
        let source = "def foo() -> int { return 0 }";
        let mut parser = Parser::new(source);
        let program = parser.parse_program().expect("Failed to parse program");
        assert_eq!(program.items.len(), 1);
    }

    #[test]
    fn test_parse_function_with_params() {
        let source = "def add(a: int, b: int) -> int { return a + b }";
        let mut parser = Parser::new(source);
        let program = parser.parse_program().expect("Failed to parse");
        assert_eq!(program.items.len(), 1);
        if let TopLevel::Func(func) = &program.items[0] {
            assert_eq!(func.params.len(), 2);
        }
    }

    #[test]
    fn test_parse_if_stmt() {
        let source = "def test() { if x > 0 { return 1 } else { return 0 } }";
        let mut parser = Parser::new(source);
        let _program = parser.parse_program().expect("Failed to parse");
    }
}