use crate::common::diagnostics::{Diagnostic, DiagnosticBundle, Severity};
use crate::common::spans::{FileId, Span};
use crate::frontends::pbs::ast::*;
use crate::frontends::pbs::lexer::Lexer;
use crate::frontends::pbs::token::{Token, TokenKind};
use prometeu_analysis::{NameId, NameInterner, NodeId};

pub struct Parser<'a> {
    tokens: Vec<Token>,
    pos: usize,
    file_id: FileId,
    errors: Vec<Diagnostic>,
    interner: &'a mut NameInterner,
    arena: AstArena,
    builtin_none: NameId,
    builtin_some: NameId,
    builtin_ok: NameId,
    builtin_err: NameId,
    builtin_true: NameId,
    builtin_false: NameId,
    builtin_void: NameId,
    builtin_array: NameId,
    builtin_optional: NameId,
    builtin_result: NameId,
    builtin_bounded: NameId,
}

impl<'a> Parser<'a> {
    pub fn new(source: &str, file_id: FileId, interner: &'a mut NameInterner) -> Self {
        let mut lexer = Lexer::new(source, file_id);
        let mut tokens = Vec::new();
        loop {
            let token = lexer.next_token();
            let is_eof = token.kind == TokenKind::Eof;
            tokens.push(token);
            if is_eof {
                break;
            }
        }

        let builtin_none = interner.intern("none");
        let builtin_some = interner.intern("some");
        let builtin_ok = interner.intern("ok");
        let builtin_err = interner.intern("err");
        let builtin_true = interner.intern("true");
        let builtin_false = interner.intern("false");
        let builtin_void = interner.intern("void");
        let builtin_array = interner.intern("array");
        let builtin_optional = interner.intern("optional");
        let builtin_result = interner.intern("result");
        let builtin_bounded = interner.intern("bounded");

        Self {
            tokens,
            pos: 0,
            file_id,
            errors: Vec::new(),
            interner,
            arena: AstArena::default(),
            builtin_none,
            builtin_some,
            builtin_ok,
            builtin_err,
            builtin_true,
            builtin_false,
            builtin_void,
            builtin_array,
            builtin_optional,
            builtin_result,
            builtin_bounded,
        }
    }

    pub fn parse_file(&mut self) -> Result<ParsedAst, DiagnosticBundle> {
        let start_span = self.peek().span.clone();
        let mut imports = Vec::new();
        let mut decls = Vec::new();

        while self.peek().kind != TokenKind::Eof {
            if self.peek().kind == TokenKind::Import {
                match self.parse_import() {
                    Ok(imp) => imports.push(imp),
                    Err(_) => self.recover_to_top_level(),
                }
            } else {
                match self.parse_top_level_decl() {
                    Ok(decl) => decls.push(decl),
                    Err(_) => self.recover_to_top_level(),
                }
            }
        }

        let end_span = self.peek().span.clone();

        if !self.errors.is_empty() {
            return Err(DiagnosticBundle {
                diagnostics: self.errors.clone(),
            });
        }

        let span = Span::new(self.file_id, start_span.start, end_span.end);
        let root = self.arena.push(NodeKind::File(FileNodeArena { imports, decls }), span);
        self.arena.roots.push(root);

        Ok(ParsedAst {
            arena: std::mem::take(&mut self.arena),
            root,
        })
    }

    fn recover_to_top_level(&mut self) {
        while self.peek().kind != TokenKind::Eof {
            match self.peek().kind {
                TokenKind::Import
                | TokenKind::Fn
                | TokenKind::Pub
                | TokenKind::Mod
                | TokenKind::Declare
                | TokenKind::Service => break,
                _ => self.advance(),
            };
        }
    }

    fn parse_import(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Import)?.span;
        let spec = self.parse_import_spec()?;
        self.consume(TokenKind::Identifier("from".to_string()))?;
        
        let path_tok = match self.peek().kind {
            TokenKind::StringLit(ref s) => {
                let s = s.clone();
                let span = self.advance().span;
                (s, span)
            }
            _ => return Err(self.error_with_code("Expected string literal after 'from'", Some("E_PARSE_EXPECTED_TOKEN"))),
        };

        if self.peek().kind == TokenKind::Semicolon {
            self.advance();
        }

        let span = Span::new(self.file_id, start_span.start, path_tok.1.end);
        Ok(self.arena.push(
            NodeKind::Import(ImportNodeArena {
                spec,
                from: path_tok.0,
            }),
            span,
        ))
    }

    fn parse_import_spec(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let mut path = Vec::new();
        let start_span = self.peek().span.clone();

        if self.peek().kind == TokenKind::OpenBrace {
            self.advance(); // {
            loop {
                let name = match self.peek().kind.clone() {
                    TokenKind::Identifier(name) => name,
                    _ => return Err(self.error("Expected identifier in import spec")),
                };
                self.advance();
                path.push(self.interner.intern(&name));

                if self.peek().kind == TokenKind::Comma {
                    self.advance();
                } else if self.peek().kind == TokenKind::CloseBrace {
                    break;
                } else {
                    return Err(self.error("Expected ',' or '}' in import spec"));
                }
            }
            self.consume(TokenKind::CloseBrace)?;
        } else {
            loop {
                let name = match self.peek().kind.clone() {
                    TokenKind::Identifier(name) => name,
                    _ => return Err(self.error("Expected identifier in import spec")),
                };
                self.advance();
                path.push(self.interner.intern(&name));

                if self.peek().kind == TokenKind::Dot {
                    self.advance();
                } else {
                    break;
                }
            }
        }

        let end_span = self.tokens[self.pos - 1].span.clone();
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(NodeKind::ImportSpec(ImportSpecNodeArena { path }), span))
    }

    fn parse_top_level_decl(&mut self) -> Result<NodeId, DiagnosticBundle> {
        match self.peek().kind {
            TokenKind::Fn => self.parse_fn_decl(None),
            TokenKind::Pub | TokenKind::Mod | TokenKind::Declare | TokenKind::Service => self.parse_decl(),
            TokenKind::Invalid(ref msg) => {
                let code = if msg.contains("Unterminated string") {
                    "E_LEX_UNTERMINATED_STRING"
                } else {
                    "E_LEX_INVALID_CHAR"
                };
                let msg = msg.clone();
                Err(self.error_with_code(&msg, Some(code)))
            }
            _ => Err(self.error_with_code("Expected top-level declaration", Some("E_PARSE_UNEXPECTED_TOKEN"))),
        }
    }

    fn parse_decl(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let vis = if self.peek().kind == TokenKind::Pub {
            self.advance();
            Some("pub".to_string())
        } else if self.peek().kind == TokenKind::Mod {
            self.advance();
            Some("mod".to_string())
        } else {
            None
        };

        match self.peek().kind {
            TokenKind::Service => self.parse_service_decl(vis),
            TokenKind::Declare => self.parse_type_decl(vis),
            TokenKind::Fn => self.parse_fn_decl(vis),
            _ => Err(self.error("Expected 'service', 'declare', or 'fn'")),
        }
    }

    fn parse_service_decl(&mut self, vis: Option<String>) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Service)?.span;
        let name = self.expect_identifier()?;
        let mut extends = None;
        if self.peek().kind == TokenKind::Colon {
            self.advance();
            extends = Some(self.expect_identifier()?);
        }
        
        self.consume(TokenKind::OpenBrace)?;
        let mut members = Vec::new();
        while self.peek().kind != TokenKind::CloseBrace && self.peek().kind != TokenKind::Eof {
            members.push(self.parse_service_member()?);
            // Optional semicolon after signature
            if self.peek().kind == TokenKind::Semicolon {
                self.advance();
            }
        }
        let end_span = self.consume(TokenKind::CloseBrace)?.span;
        
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(
            NodeKind::ServiceDecl(ServiceDeclNodeArena {
                vis,
                name,
                extends,
                members,
            }),
            span,
        ))
    }

    fn parse_service_member(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Fn)?.span;
        let name = self.expect_identifier()?;
        let params = self.parse_param_list()?;
        let ret = if self.peek().kind == TokenKind::Colon {
            self.advance();
            self.parse_type_ref()?
        } else {
            self.arena.push(
                NodeKind::TypeName(TypeNameNodeArena {
                    name: self.builtin_void,
                }),
                Span::new(self.file_id, 0, 0),
            )
        };

        let span = Span::new(self.file_id, start_span.start, self.arena.span(ret).end);
        Ok(self.arena.push(
            NodeKind::ServiceFnSig(ServiceFnSigNodeArena { name, params, ret }),
            span,
        ))
    }

    fn parse_type_decl(&mut self, vis: Option<String>) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Declare)?.span;
        let type_kind = match self.peek().kind {
            TokenKind::Struct => { self.advance(); "struct".to_string() }
            TokenKind::Contract => { self.advance(); "contract".to_string() }
            TokenKind::Error => { self.advance(); "error".to_string() }
            _ => return Err(self.error_with_code("Expected 'struct', 'contract', or 'error'", Some("E_PARSE_EXPECTED_TOKEN"))),
        };
        let name = self.expect_identifier()?;

        let mut params = Vec::new();
        if self.peek().kind == TokenKind::OpenParen {
            params = self.parse_param_list()?;
        }

        let mut constructors = Vec::new();
        if self.peek().kind == TokenKind::OpenBracket {
            constructors = self.parse_constructor_list()?;
        }

        let mut is_host = false;
        if self.peek().kind == TokenKind::Host {
            self.advance();
            is_host = true;
        }

        let mut constants = Vec::new();
        if self.peek().kind == TokenKind::OpenDoubleBracket {
            self.advance();
            while self.peek().kind != TokenKind::CloseDoubleBracket && self.peek().kind != TokenKind::Eof {
                let c_start = self.peek().span.start;
                let c_name = self.expect_identifier()?;
                self.consume(TokenKind::Colon)?;
                let c_value = self.parse_expr(0)?;
                let c_span = Span::new(self.file_id, c_start, self.arena.span(c_value).end);
                constants.push(self.arena.push(
                    NodeKind::ConstantDecl(ConstantDeclNodeArena {
                        name: c_name,
                        value: c_value,
                    }),
                    c_span,
                ));
                if self.peek().kind == TokenKind::Comma {
                    self.advance();
                }
            }
            self.consume(TokenKind::CloseDoubleBracket)?;
        }

        let mut body = None;
        if self.peek().kind == TokenKind::OpenBrace {
            body = Some(self.parse_type_body()?);
        }
        
        let mut end_pos = start_span.end;
        if let Some(b) = body {
            end_pos = self.arena.span(b).end;
            body = Some(b);
        } else if !constants.is_empty() {
            // We should use the CloseDoubleBracket span here, but I don't have it easily
            // Let's just use the last constant's end
            end_pos = self.arena.span(*constants.last().unwrap()).end;
        } else if !params.is_empty() {
            end_pos = params.last().unwrap().span.end;
        }

        let span = Span::new(self.file_id, start_span.start, end_pos);
        Ok(self.arena.push(
            NodeKind::TypeDecl(TypeDeclNodeArena {
                vis,
                type_kind,
                name,
                is_host,
                params,
                constructors,
                constants,
                body,
            }),
            span,
        ))
    }

    fn parse_type_body(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::OpenBrace)?.span;
        let mut members = Vec::new();
        let mut methods = Vec::new();
        while self.peek().kind != TokenKind::CloseBrace && self.peek().kind != TokenKind::Eof {
            if self.peek().kind == TokenKind::Fn {
                let sig_node = self.parse_service_member()?;
                methods.push(sig_node);
                if self.peek().kind == TokenKind::Semicolon {
                    self.advance();
                }
            } else {
                let m_start = self.peek().span.start;
                let name = self.expect_identifier()?;
                self.consume(TokenKind::Colon)?;
                let ty = self.parse_type_ref()?;
                let m_end = self.arena.span(ty).end;
                members.push(TypeMemberNodeArena { 
                    span: Span::new(self.file_id, m_start, m_end),
                    name, 
                    ty,
                });
                if self.peek().kind == TokenKind::Comma {
                    self.advance();
                } else if self.peek().kind == TokenKind::Semicolon {
                    self.advance();
                }
            }
        }
        let end_span = self.consume(TokenKind::CloseBrace)?.span;
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(
            NodeKind::TypeBody(TypeBodyNodeArena { members, methods }),
            span,
        ))
    }

    fn parse_fn_decl(&mut self, vis: Option<String>) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Fn)?.span;
        let name = self.expect_identifier()?;
        let params = self.parse_param_list()?;
        let _ret = if self.peek().kind == TokenKind::Colon {
            self.advance();
            Some(self.parse_type_ref()?)
        } else {
            None
        };
        
        let mut else_fallback = None;
        if self.peek().kind == TokenKind::Else {
            self.advance();
            else_fallback = Some(self.parse_block()?);
        }

        let body = self.parse_block()?;
        let body_span = self.arena.span(body);
        
        let span = Span::new(self.file_id, start_span.start, body_span.end);
        Ok(self.arena.push(
            NodeKind::FnDecl(FnDeclNodeArena {
                vis,
                name,
                params,
                ret: _ret,
                else_fallback,
                body,
            }),
            span,
        ))
    }

    fn parse_param_list(&mut self) -> Result<Vec<ParamNodeArena>, DiagnosticBundle> {
        self.consume(TokenKind::OpenParen)?;
        let mut params = Vec::new();
        while self.peek().kind != TokenKind::CloseParen && self.peek().kind != TokenKind::Eof {
            let p_start = self.peek().span.start;
            let name = self.expect_identifier()?;
            self.consume(TokenKind::Colon)?;
            let ty = self.parse_type_ref()?;
            let p_end = self.arena.span(ty).end;
            params.push(ParamNodeArena { 
                span: Span::new(self.file_id, p_start, p_end),
                name, 
                ty,
            });
            if self.peek().kind == TokenKind::Comma {
                self.advance();
            } else {
                break;
            }
        }
        self.consume(TokenKind::CloseParen)?;
        Ok(params)
    }

    fn parse_type_ref(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let id_tok = self.peek().clone();
        let name = match id_tok.kind {
            TokenKind::Identifier(ref s) => {
                self.advance();
                self.interner.intern(s)
            }
            TokenKind::Optional => {
                self.advance();
                self.builtin_optional
            }
            TokenKind::Result => {
                self.advance();
                self.builtin_result
            }
            TokenKind::Bounded => {
                self.advance();
                self.builtin_bounded
            }
            TokenKind::None => {
                self.advance();
                self.builtin_none
            }
            TokenKind::Some => {
                self.advance();
                self.builtin_some
            }
            TokenKind::Ok => {
                self.advance();
                self.builtin_ok
            }
            TokenKind::Err => {
                self.advance();
                self.builtin_err
            }
            _ => return Err(self.error_with_code("Expected type name", Some("E_PARSE_EXPECTED_TOKEN"))),
        };
        let mut node = if self.peek().kind == TokenKind::Lt {
            self.advance(); // <
            let mut args = Vec::new();
            loop {
                args.push(self.parse_type_ref()?);
                if self.peek().kind == TokenKind::Comma {
                    self.advance();
                } else {
                    break;
                }
            }
            let end_tok = self.consume(TokenKind::Gt)?;
            self.arena.push(
                NodeKind::TypeApp(TypeAppNodeArena { base: name, args }),
                Span::new(self.file_id, id_tok.span.start, end_tok.span.end),
            )
        } else {
            self.arena.push(
                NodeKind::TypeName(TypeNameNodeArena { name }),
                id_tok.span,
            )
        };

        if self.peek().kind == TokenKind::OpenBracket {
            self.advance();
            let size_tok = self.peek().clone();
            let size = match size_tok.kind {
                TokenKind::IntLit(v) => {
                    self.advance();
                    v as u32
                }
                TokenKind::BoundedLit(v) => {
                    self.advance();
                    v
                }
                _ => {
                    return Err(self.error_with_code(
                        "integer or bounded literal for array size",
                        Some("E_PARSE_EXPECTED_TOKEN"),
                    ))
                }
            };
            let end_tok = self.consume(TokenKind::CloseBracket)?;
            let size_node = self.arena.push(
                NodeKind::IntLit(IntLitNodeArena { value: size as i64 }),
                size_tok.span,
            );
            let span = Span::new(self.file_id, self.arena.span(node).start, end_tok.span.end);

            let mut wrapped = true;
            let index = node.0 as usize;
            if let NodeKind::TypeApp(ta) = &mut self.arena.nodes[index] {
                if ta.base == self.builtin_array {
                    ta.args.push(size_node);
                    self.arena.spans[index] = span.clone();
                    wrapped = false;
                }
            }

            if wrapped {
                node = self.arena.push(
                    NodeKind::TypeApp(TypeAppNodeArena {
                        base: self.builtin_array,
                        args: vec![node, size_node],
                    }),
                    span,
                );
            }
        }

        Ok(node)
    }

    fn parse_block(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::OpenBrace)?.span;
        let mut stmts = Vec::new();
        let mut tail = None;

        while self.peek().kind != TokenKind::CloseBrace && self.peek().kind != TokenKind::Eof {
            if self.peek().kind == TokenKind::Let {
                stmts.push(self.parse_let_stmt()?);
            } else if self.peek().kind == TokenKind::Return {
                stmts.push(self.parse_return_stmt()?);
            } else {
                let expr = self.parse_expr(0)?;
                if self.peek().kind == TokenKind::Semicolon {
                    let semi_span = self.advance().span;
                    let expr_start = self.arena.span(expr).start;
                    let span = Span::new(self.file_id, expr_start, semi_span.end);
                    stmts.push(self.arena.push(
                        NodeKind::ExprStmt(ExprStmtNodeArena { expr }),
                        span,
                    ));
                } else if self.peek().kind == TokenKind::CloseBrace {
                    tail = Some(expr);
                } else {
                    // Treat as ExprStmt even without semicolon (e.g. for if/when used as statement)
                    let expr_span = self.arena.span(expr);
                    stmts.push(self.arena.push(
                        NodeKind::ExprStmt(ExprStmtNodeArena { expr }),
                        expr_span,
                    ));
                }
            }
        }

        let end_span = self.consume(TokenKind::CloseBrace)?.span;
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(NodeKind::Block(BlockNodeArena { stmts, tail }), span))
    }

    fn parse_let_stmt(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Let)?.span;
        let is_mut = if self.peek().kind == TokenKind::Mut {
            self.advance();
            true
        } else {
            false
        };
        let name = self.expect_identifier()?;
        let ty = if self.peek().kind == TokenKind::Colon {
            self.advance();
            Some(self.parse_type_ref()?)
        } else {
            None
        };
        self.consume(TokenKind::Assign)?;
        let init = self.parse_expr(0)?;
        let end_span = self.consume(TokenKind::Semicolon)?.span;

        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(
            NodeKind::LetStmt(LetStmtNodeArena {
                name,
                is_mut,
                ty,
                init,
            }),
            span,
        ))
    }

    fn parse_return_stmt(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Return)?.span;
        let mut expr = None;
        if self.peek().kind != TokenKind::Semicolon {
            expr = Some(self.parse_expr(0)?);
        }
        let end_span = self.consume(TokenKind::Semicolon)?.span;
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(NodeKind::ReturnStmt(ReturnStmtNodeArena { expr }), span))
    }

    fn parse_alloc(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::Alloc)?.span;
        let ty = self.parse_type_ref()?;
        let span = Span::new(self.file_id, start_span.start, self.arena.span(ty).end);
        Ok(self.arena.push(NodeKind::Alloc(AllocNodeArena { ty }), span))
    }

    fn parse_mutate_borrow_peek(&mut self, kind: TokenKind) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(kind.clone())?.span;
        let target_expr = self.parse_expr(0)?;
        
        let (target, binding) = if let NodeKind::Cast(cast) = self.arena.kind(target_expr) {
            let binding = match self.arena.kind(cast.ty) {
                NodeKind::Ident(id) => Some(id.name),
                NodeKind::TypeName(tn) => Some(tn.name),
                _ => None,
            };
            if let Some(binding) = binding {
                (cast.expr, binding)
            } else {
                return Err(self.error_with_code(
                    "Expected binding name after 'as'",
                    Some("E_PARSE_EXPECTED_TOKEN"),
                ));
            }
        } else {
            self.consume(TokenKind::As)?;
            let binding = self.expect_identifier()?;
            (target_expr, binding)
        };

        let body = self.parse_block()?;
        let span = Span::new(self.file_id, start_span.start, self.arena.span(body).end);
        
        match kind {
            TokenKind::Mutate => Ok(self.arena.push(
                NodeKind::Mutate(MutateNodeArena { target, binding, body }),
                span,
            )),
            TokenKind::Borrow => Ok(self.arena.push(
                NodeKind::Borrow(BorrowNodeArena { target, binding, body }),
                span,
            )),
            TokenKind::Peek => Ok(self.arena.push(
                NodeKind::Peek(PeekNodeArena { target, binding, body }),
                span,
            )),
            _ => unreachable!(),
        }
    }

    fn parse_expr(&mut self, min_precedence: u8) -> Result<NodeId, DiagnosticBundle> {
        let mut left = self.parse_primary()?;

        loop {
            let (op, precedence) = match self.get_binary_precedence() {
                Some((op, p)) if p >= min_precedence => (op, p),
                _ => break,
            };

            self.advance();
            let right = self.parse_expr(precedence + 1)?;
            let span = Span::new(
                self.file_id,
                self.arena.span(left).start,
                self.arena.span(right).end,
            );
            left = self.arena.push(
                NodeKind::Binary(BinaryNodeArena { op, left, right }),
                span,
            );
        }

        Ok(left)
    }

    fn parse_primary(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let tok = self.peek().clone();
        match tok.kind {
            TokenKind::IntLit(v) => {
                self.advance();
                Ok(self.arena.push(
                    NodeKind::IntLit(IntLitNodeArena { value: v }),
                    tok.span,
                ))
            }
            TokenKind::FloatLit(v) => {
                self.advance();
                Ok(self.arena.push(
                    NodeKind::FloatLit(FloatLitNodeArena { value: v }),
                    tok.span,
                ))
            }
            TokenKind::BoundedLit(v) => {
                self.advance();
                Ok(self.arena.push(
                    NodeKind::BoundedLit(BoundedLitNodeArena { value: v }),
                    tok.span,
                ))
            }
            TokenKind::StringLit(s) => {
                self.advance();
                Ok(self.arena.push(
                    NodeKind::StringLit(StringLitNodeArena { value: s }),
                    tok.span,
                ))
            }
            TokenKind::Identifier(name) => {
                self.advance();
                let name = self.interner.intern(&name);
                let mut node = self.arena.push(
                    NodeKind::Ident(IdentNodeArena { name }),
                    tok.span,
                );
                loop {
                    if self.peek().kind == TokenKind::OpenParen {
                        node = self.parse_call(node)?;
                    } else if self.peek().kind == TokenKind::As {
                        node = self.parse_cast(node)?;
                    } else if self.peek().kind == TokenKind::Dot {
                        node = self.parse_member_access(node)?;
                    } else {
                        break;
                    }
                }
                Ok(node)
            }
            TokenKind::None | TokenKind::Some | TokenKind::Ok | TokenKind::Err => {
                let name = match tok.kind {
                    TokenKind::None => self.builtin_none,
                    TokenKind::Some => self.builtin_some,
                    TokenKind::Ok => self.builtin_ok,
                    TokenKind::Err => self.builtin_err,
                    _ => unreachable!(),
                };
                self.advance();
                let mut node = self.arena.push(
                    NodeKind::Ident(IdentNodeArena { name }),
                    tok.span,
                );
                loop {
                    if self.peek().kind == TokenKind::OpenParen {
                        node = self.parse_call(node)?;
                    } else if self.peek().kind == TokenKind::As {
                        node = self.parse_cast(node)?;
                    } else if self.peek().kind == TokenKind::Dot {
                        node = self.parse_member_access(node)?;
                    } else {
                        break;
                    }
                }
                Ok(node)
            }
            TokenKind::OpenParen => {
                self.advance();
                let expr = self.parse_expr(0)?;
                self.consume(TokenKind::CloseParen)?;
                Ok(expr)
            }
            TokenKind::OpenBrace => self.parse_block(),
            TokenKind::If => self.parse_if_expr(),
            TokenKind::When => self.parse_when_expr(),
            TokenKind::Alloc => self.parse_alloc(),
            TokenKind::Mutate => self.parse_mutate_borrow_peek(TokenKind::Mutate),
            TokenKind::Borrow => self.parse_mutate_borrow_peek(TokenKind::Borrow),
            TokenKind::Peek => self.parse_mutate_borrow_peek(TokenKind::Peek),
            TokenKind::Minus | TokenKind::Not => {
                self.advance();
                let op = match tok.kind {
                    TokenKind::Minus => "-".to_string(),
                    TokenKind::Not => "!".to_string(),
                    _ => unreachable!(),
                };
                let expr = self.parse_expr(11)?;
                let span = Span::new(self.file_id, tok.span.start, self.arena.span(expr).end);
                Ok(self.arena.push(
                    NodeKind::Unary(UnaryNodeArena { op, expr }),
                    span,
                ))
            }
            TokenKind::Invalid(msg) => {
                let code = if msg.contains("Unterminated string") {
                    "E_LEX_UNTERMINATED_STRING"
                } else {
                    "E_LEX_INVALID_CHAR"
                };
                Err(self.error_with_code(&msg, Some(code)))
            }
            _ => Err(self.error_with_code("Expected expression", Some("E_PARSE_UNEXPECTED_TOKEN"))),
        }
    }

    fn parse_member_access(&mut self, object: NodeId) -> Result<NodeId, DiagnosticBundle> {
        self.consume(TokenKind::Dot)?;
        let member = self.expect_identifier()?;
        let span = Span::new(
            self.file_id,
            self.arena.span(object).start,
            self.tokens[self.pos - 1].span.end,
        );
        Ok(self.arena.push(
            NodeKind::MemberAccess(MemberAccessNodeArena { object, member }),
            span,
        ))
    }

    fn parse_call(&mut self, callee: NodeId) -> Result<NodeId, DiagnosticBundle> {
        self.consume(TokenKind::OpenParen)?;
        let mut args = Vec::new();
        while self.peek().kind != TokenKind::CloseParen && self.peek().kind != TokenKind::Eof {
            args.push(self.parse_expr(0)?);
            if self.peek().kind == TokenKind::Comma {
                self.advance();
            } else {
                break;
            }
        }
        let end_span = self.consume(TokenKind::CloseParen)?.span;
        let span = Span::new(self.file_id, self.arena.span(callee).start, end_span.end);
        Ok(self.arena.push(
            NodeKind::Call(CallNodeArena { callee, args }),
            span,
        ))
    }

    fn parse_cast(&mut self, expr: NodeId) -> Result<NodeId, DiagnosticBundle> {
        self.consume(TokenKind::As)?;
        let ty = self.parse_type_ref()?;
        let span = Span::new(self.file_id, self.arena.span(expr).start, self.arena.span(ty).end);
        Ok(self.arena.push(NodeKind::Cast(CastNodeArena { expr, ty }), span))
    }

    fn parse_if_expr(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::If)?.span;
        let cond = self.parse_expr(0)?;
        let then_block = self.parse_block()?;
        let mut else_block = None;
        if self.peek().kind == TokenKind::Else {
            self.advance();
            if self.peek().kind == TokenKind::If {
                else_block = Some(self.parse_if_expr()?);
            } else {
                else_block = Some(self.parse_block()?);
            }
        }

        let end_span = else_block
            .map(|b| self.arena.span(b).end)
            .unwrap_or(self.arena.span(then_block).end);

        let span = Span::new(self.file_id, start_span.start, end_span);
        Ok(self.arena.push(
            NodeKind::IfExpr(IfExprNodeArena {
                cond,
                then_block,
                else_block,
            }),
            span,
        ))
    }

    fn parse_when_expr(&mut self) -> Result<NodeId, DiagnosticBundle> {
        let start_span = self.consume(TokenKind::When)?.span;
        self.consume(TokenKind::OpenBrace)?;
        let mut arms = Vec::new();
        while self.peek().kind != TokenKind::CloseBrace && self.peek().kind != TokenKind::Eof {
            let arm_start = self.peek().span.start;
            let cond = self.parse_expr(0)?;
            self.consume(TokenKind::Arrow)?;
            let body = self.parse_block()?;
            let arm_span = Span::new(self.file_id, arm_start, self.arena.span(body).end);
            arms.push(self.arena.push(
                NodeKind::WhenArm(WhenArmNodeArena { cond, body }),
                arm_span,
            ));
            if self.peek().kind == TokenKind::Comma {
                self.advance();
            }
        }
        let end_span = self.consume(TokenKind::CloseBrace)?.span;
        let span = Span::new(self.file_id, start_span.start, end_span.end);
        Ok(self.arena.push(NodeKind::WhenExpr(WhenExprNodeArena { arms }), span))
    }

    fn get_binary_precedence(&self) -> Option<(String, u8)> {
        match self.peek().kind {
            TokenKind::Plus => Some(("+".to_string(), 5)),
            TokenKind::Minus => Some(("-".to_string(), 5)),
            TokenKind::Star => Some(("*".to_string(), 4)),
            TokenKind::Slash => Some(("/".to_string(), 4)),
            TokenKind::Percent => Some(("%".to_string(), 4)),
            TokenKind::Lt => Some(("<".to_string(), 7)),
            TokenKind::Lte => Some(("<=".to_string(), 7)),
            TokenKind::Gt => Some((">".to_string(), 7)),
            TokenKind::Gte => Some((">=".to_string(), 7)),
            TokenKind::Eq => Some(("==".to_string(), 8)),
            TokenKind::Neq => Some(("!=".to_string(), 8)),
            TokenKind::And => Some(("&&".to_string(), 9)),
            TokenKind::Or => Some(("||".to_string(), 10)),
            _ => None,
        }
    }

    fn peek(&self) -> &Token {
        &self.tokens[self.pos]
    }

    fn advance(&mut self) -> Token {
        let tok = self.tokens[self.pos].clone();
        if tok.kind != TokenKind::Eof {
            self.pos += 1;
        }
        tok
    }

    fn consume(&mut self, kind: TokenKind) -> Result<Token, DiagnosticBundle> {
        let peeked_kind = self.peek().kind.clone();
        if peeked_kind == kind {
            Ok(self.advance())
        } else {
            if let TokenKind::Invalid(ref msg) = peeked_kind {
                let code = if msg.contains("Unterminated string") {
                    "E_LEX_UNTERMINATED_STRING"
                } else {
                    "E_LEX_INVALID_CHAR"
                };
                let msg = msg.clone();
                Err(self.error_with_code(&msg, Some(code)))
            } else {
                Err(self.error_with_code(&format!("Expected {:?}, found {:?}", kind, peeked_kind), Some("E_PARSE_EXPECTED_TOKEN")))
            }
        }
    }

    fn expect_identifier(&mut self) -> Result<NameId, DiagnosticBundle> {
        let peeked_kind = self.peek().kind.clone();
        match peeked_kind {
            TokenKind::Identifier(name) => {
                self.advance();
                Ok(self.interner.intern(&name))
            }
            TokenKind::Optional => {
                self.advance();
                Ok(self.builtin_optional)
            }
            TokenKind::Result => {
                self.advance();
                Ok(self.builtin_result)
            }
            TokenKind::None => {
                self.advance();
                Ok(self.builtin_none)
            }
            TokenKind::Some => {
                self.advance();
                Ok(self.builtin_some)
            }
            TokenKind::Ok => {
                self.advance();
                Ok(self.builtin_ok)
            }
            TokenKind::Err => {
                self.advance();
                Ok(self.builtin_err)
            }
            TokenKind::Bounded => {
                self.advance();
                Ok(self.builtin_bounded)
            }
            TokenKind::Invalid(msg) => {
                 let code = if msg.contains("Unterminated string") {
                    "E_LEX_UNTERMINATED_STRING"
                } else {
                    "E_LEX_INVALID_CHAR"
                };
                Err(self.error_with_code(&msg, Some(code)))
            }
            _ => Err(self.error_with_code(&format!("Expected identifier, found {:?}", peeked_kind), Some("E_PARSE_EXPECTED_TOKEN"))),
        }
    }

    fn error(&mut self, message: &str) -> DiagnosticBundle {
        self.error_with_code(message, None)
    }

    fn error_with_code(&mut self, message: &str, code: Option<&str>) -> DiagnosticBundle {
        let diag = Diagnostic {
            severity: Severity::Error,
            code: code.unwrap_or("E_PARSE_ERROR").to_string(),
            message: message.to_string(),
            span: self.peek().span.clone(),
            related: Vec::new(),
        };
        self.errors.push(diag.clone());
        DiagnosticBundle::from(diag)
    }

    fn parse_constructor_list(&mut self) -> Result<Vec<NodeId>, DiagnosticBundle> {
        self.consume(TokenKind::OpenBracket)?;
        let mut constructors = Vec::new();
        while self.peek().kind != TokenKind::CloseBracket && self.peek().kind != TokenKind::Eof {
            let start_span = self.peek().span.clone();
            let params = self.parse_param_list()?;
            self.consume(TokenKind::Colon)?;
            
            let mut initializers = Vec::new();
            if self.peek().kind == TokenKind::OpenParen {
                self.advance();
                while self.peek().kind != TokenKind::CloseParen && self.peek().kind != TokenKind::Eof {
                    initializers.push(self.parse_expr(0)?);
                    if self.peek().kind == TokenKind::Comma {
                        self.advance();
                    }
                }
                self.consume(TokenKind::CloseParen)?;
            } else {
                initializers.push(self.parse_expr(0)?);
            }

            self.consume(TokenKind::As)?;
            let name = self.expect_identifier()?;
            
            let body = self.parse_block()?;
            
            let span = Span::new(self.file_id, start_span.start, self.arena.span(body).end);
            constructors.push(self.arena.push(
                NodeKind::ConstructorDecl(ConstructorDeclNodeArena {
                    params,
                    initializers,
                    name,
                    body,
                }),
                span,
            ));
        }
        self.consume(TokenKind::CloseBracket)?;
        Ok(constructors)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json;
    use crate::common::spans::FileId;

    #[test]
    fn test_parse_empty_file() {
        let mut interner = NameInterner::new();
        let mut parser = Parser::new("", FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.imports.len(), 0);
        assert_eq!(file.decls.len(), 0);
    }

    #[test]
    fn test_parse_imports() {
        let source = r#"
import std.io from "std";
import math from "./math.pbs";
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.imports.len(), 2);
        
        let imp = match parsed.arena.kind(file.imports[0]) {
            NodeKind::Import(imp) => imp,
            _ => panic!("Expected Import"),
        };
        assert_eq!(imp.from, "std");
        let spec = match parsed.arena.kind(imp.spec) {
            NodeKind::ImportSpec(spec) => spec,
            _ => panic!("Expected ImportSpec"),
        };
        let path: Vec<&str> = spec.path.iter().map(|id| interner.resolve(*id)).collect();
        assert_eq!(path, vec!["std", "io"]);
    }

    #[test]
    fn test_parse_fn_decl() {
        let source = r#"
fn add(a: int, b: int): int {
    return a + b;
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.decls.len(), 1);
        
        if let NodeKind::FnDecl(f) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(interner.resolve(f.name), "add");
            assert_eq!(f.params.len(), 2);
            assert_eq!(interner.resolve(f.params[0].name), "a");
            assert_eq!(interner.resolve(f.params[1].name), "b");
        } else {
            panic!("Expected FnDecl");
        }
    }

    #[test]
    fn test_parse_type_decl() {
        let source = r#"
pub declare struct Point {
    x: int,
    y: int
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.decls.len(), 1);
        
        if let NodeKind::TypeDecl(t) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(interner.resolve(t.name), "Point");
            assert_eq!(t.type_kind, "struct");
            assert_eq!(t.vis, Some("pub".to_string()));
        } else {
            panic!("Expected TypeDecl");
        }
    }

    #[test]
    fn test_parse_service_decl() {
        let source = r#"
pub service Audio {
    fn play(sound: Sound);
    fn stop(): bool;
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.decls.len(), 1);
        
        if let NodeKind::ServiceDecl(s) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(interner.resolve(s.name), "Audio");
            assert_eq!(s.members.len(), 2);
        } else {
            panic!("Expected ServiceDecl");
        }
    }

    #[test]
    fn test_parse_expressions() {
        let source = r#"
fn main() {
    let x = 10 + 20 * 30;
    let y = (x - 5) / 2;
    foo(x, y);
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.decls.len(), 1);
    }

    #[test]
    fn test_parse_if_when() {
        let source = r#"
fn main(x: int) {
    if x > 0 {
        print("positive");
    } else {
        print("non-positive");
    }

    let msg = when {
        x == 0 -> { return "zero"; },
        x == 1 -> { return "one"; }
    };
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        assert_eq!(file.decls.len(), 1);
    }

    #[test]
    fn test_parse_error_recovery() {
        let source = r#"
fn bad() { 
    let x = ; // Missing init
    let y = 10;
}

fn good() {}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let result = parser.parse_file();
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_mod_fn() {
        let source = "mod fn test() {}";
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().expect("mod fn should be allowed");
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        if let NodeKind::FnDecl(fn_decl) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(fn_decl.vis, Some("mod".to_string()));
        } else {
            panic!("Expected FnDecl");
        }
    }

    #[test]
    fn test_parse_pub_fn() {
        let source = "pub fn test() {}";
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().expect("pub fn should be allowed in parser");
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        if let NodeKind::FnDecl(fn_decl) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(fn_decl.vis, Some("pub".to_string()));
        } else {
            panic!("Expected FnDecl");
        }
    }

    #[test]
    fn test_ast_json_snapshot() {
        let source = r#"
fn main() {
    return 42;
}
"#;
        let mut interner = NameInterner::new();
        let mut parser = Parser::new(source, FileId(0), &mut interner);
        let parsed = parser.parse_file().unwrap();
        let file = match parsed.arena.kind(parsed.root) {
            NodeKind::File(file) => file,
            _ => panic!("Expected File"),
        };
        let json = serde_json::to_string_pretty(parsed.arena.kind(parsed.root)).unwrap();

        assert!(json.contains("\"kind\": \"File\""));
        if let NodeKind::FnDecl(fn_decl) = parsed.arena.kind(file.decls[0]) {
            assert_eq!(interner.resolve(fn_decl.name), "main");
        } else {
            panic!("Expected FnDecl");
        }
    }
}