use crate::common::diagnostics::{Diagnostic, DiagnosticBundle, DiagnosticLevel};
use crate::common::spans::Span;
use crate::frontends::pbs::ast::*;
use crate::frontends::pbs::lexer::Lexer;
use crate::frontends::pbs::token::{Token, TokenKind};

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

impl Parser {
    pub fn new(source: &str, file_id: usize) -> 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;
            }
        }

        Self {
            tokens,
            pos: 0,
            file_id,
            errors: Vec::new(),
        }
    }

    pub fn parse_file(&mut self) -> Result<FileNode, DiagnosticBundle> {
        let start_span = self.peek().span;
        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;

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

        Ok(FileNode {
            span: Span::new(self.file_id, start_span.start, end_span.end),
            imports,
            decls,
        })
    }

    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<Node, 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();
        }

        Ok(Node::Import(ImportNode {
            span: Span::new(self.file_id, start_span.start, path_tok.1.end),
            spec: Box::new(spec),
            from: path_tok.0,
        }))
    }

    fn parse_import_spec(&mut self) -> Result<Node, DiagnosticBundle> {
        let mut path = Vec::new();
        let start_span = self.peek().span;
        loop {
            if let TokenKind::Identifier(ref name) = self.peek().kind {
                path.push(name.clone());
                self.advance();
            } else {
                return Err(self.error("Expected identifier in import spec"));
            }

            if self.peek().kind == TokenKind::Dot {
                self.advance();
            } else {
                break;
            }
        }
        let end_span = self.tokens[self.pos-1].span;
        Ok(Node::ImportSpec(ImportSpecNode { 
            span: Span::new(self.file_id, start_span.start, end_span.end),
            path 
        }))
    }

    fn parse_top_level_decl(&mut self) -> Result<Node, DiagnosticBundle> {
        match self.peek().kind {
            TokenKind::Fn => self.parse_fn_decl("file".to_string()),
            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<Node, 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.unwrap_or_else(|| "pub".to_string())),
            TokenKind::Declare => self.parse_type_decl(vis),
            TokenKind::Fn => {
                let vis_str = vis.unwrap_or_else(|| "file".to_string());
                if vis_str == "pub" {
                    return Err(self.error_with_code(
                        "Functions cannot be public. They are always mod or file-private.",
                        Some("E_RESOLVE_VISIBILITY"),
                    ));
                }
                self.parse_fn_decl(vis_str)
            }
            _ => Err(self.error("Expected 'service', 'declare', or 'fn'")),
        }
    }

    fn parse_service_decl(&mut self, vis: String) -> Result<Node, 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;
        
        Ok(Node::ServiceDecl(ServiceDeclNode {
            span: Span::new(self.file_id, start_span.start, end_span.end),
            vis,
            name,
            extends,
            members,
        }))
    }

    fn parse_service_member(&mut self) -> Result<Node, 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();
            Box::new(self.parse_type_ref()?)
        } else {
            Box::new(Node::TypeName(TypeNameNode {
                span: Span::new(self.file_id, 0, 0), // Placeholder for void
                name: "void".to_string(),
            }))
        };

        Ok(Node::ServiceFnSig(ServiceFnSigNode {
            span: Span::new(self.file_id, start_span.start, ret.span().end),
            name,
            params,
            ret,
        }))
    }

    fn parse_type_decl(&mut self, vis: Option<String>) -> Result<Node, 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)?;
                constants.push(ConstantDeclNode {
                    span: Span::new(self.file_id, c_start, c_value.span().end),
                    name: c_name,
                    value: Box::new(c_value),
                });
                if self.peek().kind == TokenKind::Comma {
                    self.advance();
                }
            }
            self.consume(TokenKind::CloseDoubleBracket)?;
        }

        let mut body = None;
        if self.peek().kind == TokenKind::OpenBrace {
            body = Some(Box::new(self.parse_type_body()?));
        }
        
        let mut end_pos = start_span.end;
        if let Some(b) = &body {
            end_pos = b.span().end;
        } 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 = constants.last().unwrap().span.end;
        } else if !params.is_empty() {
            end_pos = params.last().unwrap().span.end;
        }

        Ok(Node::TypeDecl(TypeDeclNode {
            span: Span::new(self.file_id, start_span.start, end_pos),
            vis,
            type_kind,
            name,
            is_host,
            params,
            constructors,
            constants,
            body,
        }))
    }

    fn parse_type_body(&mut self) -> Result<Node, 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()?;
                if let Node::ServiceFnSig(sig) = sig_node {
                    methods.push(sig);
                }
                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 = ty.span().end;
                members.push(TypeMemberNode { 
                    span: Span::new(self.file_id, m_start, m_end),
                    name, 
                    ty: Box::new(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;
        Ok(Node::TypeBody(TypeBodyNode {
            span: Span::new(self.file_id, start_span.start, end_span.end),
            members,
            methods,
        }))
    }

    fn parse_fn_decl(&mut self, vis: String) -> Result<Node, 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(Box::new(self.parse_type_ref()?))
        } else {
            None
        };
        
        let mut else_fallback = None;
        if self.peek().kind == TokenKind::Else {
            self.advance();
            else_fallback = Some(Box::new(self.parse_block()?));
        }

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

    fn parse_param_list(&mut self) -> Result<Vec<ParamNode>, 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 = ty.span().end;
            params.push(ParamNode { 
                span: Span::new(self.file_id, p_start, p_end),
                name, 
                ty: Box::new(ty) 
            });
            if self.peek().kind == TokenKind::Comma {
                self.advance();
            } else {
                break;
            }
        }
        self.consume(TokenKind::CloseParen)?;
        Ok(params)
    }

    fn parse_type_ref(&mut self) -> Result<Node, DiagnosticBundle> {
        let id_tok = self.peek().clone();
        let name = match id_tok.kind {
            TokenKind::Identifier(ref s) => {
                self.advance();
                s.clone()
            }
            TokenKind::Optional => {
                self.advance();
                "optional".to_string()
            }
            TokenKind::Result => {
                self.advance();
                "result".to_string()
            }
            TokenKind::Bounded => {
                self.advance();
                "bounded".to_string()
            }
            TokenKind::None => {
                self.advance();
                "none".to_string()
            }
            TokenKind::Some => {
                self.advance();
                "some".to_string()
            }
            TokenKind::Ok => {
                self.advance();
                "ok".to_string()
            }
            TokenKind::Err => {
                self.advance();
                "err".to_string()
            }
            _ => 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)?;
            Node::TypeApp(TypeAppNode {
                span: Span::new(self.file_id, id_tok.span.start, end_tok.span.end),
                base: name,
                args,
            })
        } else {
            Node::TypeName(TypeNameNode {
                span: id_tok.span,
                name,
            })
        };

        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 span = Span::new(self.file_id, node.span().start, end_tok.span.end);
            
            // If it's array<T>[N], we want to represent it cleanly.
            // Currently TypeAppNode { base: name, args } was created.
            // If base was "array", it already has T in args.
            // We can just add N to args.
            match &mut node {
                Node::TypeApp(ta) if ta.base == "array" => {
                    ta.args.push(Node::IntLit(IntLitNode { span: size_tok.span, value: size as i64 }));
                    ta.span = span;
                }
                _ => {
                    // Fallback for T[N] if we want to support it, but spec says array<T>[N]
                    node = Node::TypeApp(TypeAppNode {
                        span,
                        base: "array".to_string(),
                        args: vec![node, Node::IntLit(IntLitNode { span: size_tok.span, value: size as i64 })],
                    });
                }
            }
        }

        Ok(node)
    }

    fn parse_block(&mut self) -> Result<Node, 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 = expr.span().start;
                    stmts.push(Node::ExprStmt(ExprStmtNode {
                        span: Span::new(self.file_id, expr_start, semi_span.end),
                        expr: Box::new(expr),
                    }));
                } else if self.peek().kind == TokenKind::CloseBrace {
                    tail = Some(Box::new(expr));
                } else {
                    // Treat as ExprStmt even without semicolon (e.g. for if/when used as statement)
                    let expr_span = expr.span();
                    stmts.push(Node::ExprStmt(ExprStmtNode {
                        span: expr_span,
                        expr: Box::new(expr),
                    }));
                }
            }
        }

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

    fn parse_let_stmt(&mut self) -> Result<Node, 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(Box::new(self.parse_type_ref()?))
        } else {
            None
        };
        self.consume(TokenKind::Assign)?;
        let init = self.parse_expr(0)?;
        let end_span = self.consume(TokenKind::Semicolon)?.span;
        
        Ok(Node::LetStmt(LetStmtNode {
            span: Span::new(self.file_id, start_span.start, end_span.end),
            name,
            is_mut,
            ty,
            init: Box::new(init),
        }))
    }

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

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

    fn parse_mutate_borrow_peek(&mut self, kind: TokenKind) -> Result<Node, DiagnosticBundle> {
        let start_span = self.consume(kind.clone())?.span;
        let target_expr = self.parse_expr(0)?;
        
        let (target, binding) = match target_expr {
            Node::Cast(cast) => {
                match *cast.ty {
                    Node::Ident(id) => (*cast.expr, id.name),
                    Node::TypeName(tn) => (*cast.expr, tn.name),
                    _ => return Err(self.error_with_code("Expected binding name after 'as'", Some("E_PARSE_EXPECTED_TOKEN"))),
                }
            }
            _ => {
                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, body.span().end);
        
        match kind {
            TokenKind::Mutate => Ok(Node::Mutate(MutateNode { span, target: Box::new(target), binding, body: Box::new(body) })),
            TokenKind::Borrow => Ok(Node::Borrow(BorrowNode { span, target: Box::new(target), binding, body: Box::new(body) })),
            TokenKind::Peek => Ok(Node::Peek(PeekNode { span, target: Box::new(target), binding, body: Box::new(body) })),
            _ => unreachable!(),
        }
    }

    fn parse_expr(&mut self, min_precedence: u8) -> Result<Node, 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, left.span().start, right.span().end);
            left = Node::Binary(BinaryNode {
                span,
                op,
                left: Box::new(left),
                right: Box::new(right),
            });
        }

        Ok(left)
    }

    fn parse_primary(&mut self) -> Result<Node, DiagnosticBundle> {
        let tok = self.peek().clone();
        match tok.kind {
            TokenKind::IntLit(v) => {
                self.advance();
                Ok(Node::IntLit(IntLitNode { span: tok.span, value: v }))
            }
            TokenKind::FloatLit(v) => {
                self.advance();
                Ok(Node::FloatLit(FloatLitNode { span: tok.span, value: v }))
            }
            TokenKind::BoundedLit(v) => {
                self.advance();
                Ok(Node::BoundedLit(BoundedLitNode { span: tok.span, value: v }))
            }
            TokenKind::StringLit(s) => {
                self.advance();
                Ok(Node::StringLit(StringLitNode { span: tok.span, value: s }))
            }
            TokenKind::Identifier(name) => {
                self.advance();
                let mut node = Node::Ident(IdentNode { span: tok.span, name });
                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 => "none",
                    TokenKind::Some => "some",
                    TokenKind::Ok => "ok",
                    TokenKind::Err => "err",
                    _ => unreachable!(),
                }.to_string();
                self.advance();
                let mut node = Node::Ident(IdentNode { span: tok.span, name });
                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)?;
                Ok(Node::Unary(UnaryNode {
                    span: Span::new(self.file_id, tok.span.start, expr.span().end),
                    op,
                    expr: Box::new(expr),
                }))
            }
            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: Node) -> Result<Node, DiagnosticBundle> {
        self.consume(TokenKind::Dot)?;
        let member = self.expect_identifier()?;
        Ok(Node::MemberAccess(MemberAccessNode {
            span: Span::new(self.file_id, object.span().start, self.tokens[self.pos-1].span.end),
            object: Box::new(object),
            member,
        }))
    }

    fn parse_call(&mut self, callee: Node) -> Result<Node, 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;
        Ok(Node::Call(CallNode {
            span: Span::new(self.file_id, callee.span().start, end_span.end),
            callee: Box::new(callee),
            args,
        }))
    }

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

    fn parse_if_expr(&mut self) -> Result<Node, 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(Box::new(self.parse_if_expr()?));
            } else {
                else_block = Some(Box::new(self.parse_block()?));
            }
        }
        
        let end_span = else_block.as_ref().map(|b| b.span().end).unwrap_or(then_block.span().end);

        Ok(Node::IfExpr(IfExprNode {
            span: Span::new(self.file_id, start_span.start, end_span),
            cond: Box::new(cond),
            then_block: Box::new(then_block),
            else_block,
        }))
    }

    fn parse_when_expr(&mut self) -> Result<Node, 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()?;
            arms.push(Node::WhenArm(WhenArmNode {
                span: Span::new(self.file_id, arm_start, body.span().end),
                cond: Box::new(cond),
                body: Box::new(body),
            }));
            if self.peek().kind == TokenKind::Comma {
                self.advance();
            }
        }
        let end_span = self.consume(TokenKind::CloseBrace)?.span;
        Ok(Node::WhenExpr(WhenExprNode {
            span: Span::new(self.file_id, start_span.start, end_span.end),
            arms,
        }))
    }

    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<String, DiagnosticBundle> {
        let peeked_kind = self.peek().kind.clone();
        match peeked_kind {
            TokenKind::Identifier(name) => {
                self.advance();
                Ok(name)
            }
            TokenKind::Optional => {
                self.advance();
                Ok("optional".to_string())
            }
            TokenKind::Result => {
                self.advance();
                Ok("result".to_string())
            }
            TokenKind::None => {
                self.advance();
                Ok("none".to_string())
            }
            TokenKind::Some => {
                self.advance();
                Ok("some".to_string())
            }
            TokenKind::Ok => {
                self.advance();
                Ok("ok".to_string())
            }
            TokenKind::Err => {
                self.advance();
                Ok("err".to_string())
            }
            TokenKind::Bounded => {
                self.advance();
                Ok("bounded".to_string())
            }
            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 {
            level: DiagnosticLevel::Error,
            code: code.map(|c| c.to_string()),
            message: message.to_string(),
            span: Some(self.peek().span),
        };
        self.errors.push(diag.clone());
        DiagnosticBundle::from(diag)
    }

    fn parse_constructor_list(&mut self) -> Result<Vec<ConstructorDeclNode>, 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;
            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()?;
            
            constructors.push(ConstructorDeclNode {
                span: Span::new(self.file_id, start_span.start, body.span().end),
                params,
                initializers,
                name,
                body: Box::new(body),
            });
        }
        self.consume(TokenKind::CloseBracket)?;
        Ok(constructors)
    }
}

impl Node {
    pub fn span(&self) -> Span {
        match self {
            Node::File(n) => n.span,
            Node::Import(n) => n.span,
            Node::ImportSpec(n) => n.span,
            Node::ServiceDecl(n) => n.span,
            Node::ServiceFnSig(n) => n.span,
            Node::FnDecl(n) => n.span,
            Node::TypeDecl(n) => n.span,
            Node::TypeBody(n) => n.span,
            Node::Block(n) => n.span,
            Node::LetStmt(n) => n.span,
            Node::ExprStmt(n) => n.span,
            Node::ReturnStmt(n) => n.span,
            Node::IntLit(n) => n.span,
            Node::FloatLit(n) => n.span,
            Node::BoundedLit(n) => n.span,
            Node::StringLit(n) => n.span,
            Node::Ident(n) => n.span,
            Node::Call(n) => n.span,
            Node::Unary(n) => n.span,
            Node::Binary(n) => n.span,
            Node::Cast(n) => n.span,
            Node::IfExpr(n) => n.span,
            Node::WhenExpr(n) => n.span,
            Node::WhenArm(n) => n.span,
            Node::TypeName(n) => n.span,
            Node::TypeApp(n) => n.span,
            Node::ConstructorDecl(n) => n.span,
            Node::ConstantDecl(n) => n.span,
            Node::Alloc(n) => n.span,
            Node::Mutate(n) => n.span,
            Node::Borrow(n) => n.span,
            Node::Peek(n) => n.span,
            Node::MemberAccess(n) => n.span,
        }
    }
}

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

    #[test]
    fn test_parse_empty_file() {
        let mut parser = Parser::new("", 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.imports.len(), 0);
        assert_eq!(result.decls.len(), 0);
    }

    #[test]
    fn test_parse_imports() {
        let source = r#"
import std.io from "std";
import math from "./math.pbs";
"#;
        let mut parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.imports.len(), 2);
        
        if let Node::Import(ref imp) = result.imports[0] {
            assert_eq!(imp.from, "std");
            if let Node::ImportSpec(ref spec) = *imp.spec {
                assert_eq!(spec.path, vec!["std", "io"]);
            } else { panic!("Expected ImportSpec"); }
        } else { panic!("Expected Import"); }
    }

    #[test]
    fn test_parse_fn_decl() {
        let source = r#"
fn add(a: int, b: int): int {
    return a + b;
}
"#;
        let mut parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.decls.len(), 1);
        
        if let Node::FnDecl(ref f) = result.decls[0] {
            assert_eq!(f.name, "add");
            assert_eq!(f.params.len(), 2);
            assert_eq!(f.params[0].name, "a");
            assert_eq!(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 parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.decls.len(), 1);
        
        if let Node::TypeDecl(ref t) = result.decls[0] {
            assert_eq!(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 parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.decls.len(), 1);
        
        if let Node::ServiceDecl(ref s) = result.decls[0] {
            assert_eq!(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 parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.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 parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        assert_eq!(result.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 parser = Parser::new(source, 0);
        let result = parser.parse_file();
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_mod_fn() {
        let source = "mod fn test() {}";
        let mut parser = Parser::new(source, 0);
        let result = parser.parse_file().expect("mod fn should be allowed");
        if let Node::FnDecl(fn_decl) = &result.decls[0] {
            assert_eq!(fn_decl.vis, "mod");
        } else {
            panic!("Expected FnDecl");
        }
    }

    #[test]
    fn test_parse_pub_fn() {
        let source = "pub fn test() {}";
        let mut parser = Parser::new(source, 0);
        let result = parser.parse_file();
        assert!(result.is_err(), "pub fn should be disallowed");
        let err = result.unwrap_err();
        assert!(err.diagnostics[0].message.contains("Functions cannot be public"));
    }

    #[test]
    fn test_ast_json_snapshot() {
        let source = r#"
fn main() {
    return 42;
}
"#;
        let mut parser = Parser::new(source, 0);
        let result = parser.parse_file().unwrap();
        let json = serde_json::to_string_pretty(&Node::File(result)).unwrap();
        
        assert!(json.contains("\"kind\": \"File\""));
        assert!(json.contains("\"kind\": \"FnDecl\""));
        assert!(json.contains("\"name\": \"main\""));
    }
}