RustAst/src/ast/parser.rs

use std::rc::Rc;
use crate::ast::lexer::{Lexer, Token, TokenKind};
use crate::ast::types::{Identity, NodeIdentity, Value, Keyword};
use crate::ast::nodes::{Node, UntypedKind, Context};

pub struct Parser<'a> {
    lexer: Lexer<'a>,
    current_token: Token,
}

impl<'a> Parser<'a> {
    pub fn new(input: &'a str) -> Result<Self, String> {
        let mut lexer = Lexer::new(input);
        let current_token = lexer.next_token()?;
        Ok(Self { lexer, current_token })
    }

    fn advance(&mut self) -> Result<Token, String> {
        let prev = std::mem::replace(&mut self.current_token, self.lexer.next_token()?);
        Ok(prev)
    }

    fn peek(&self) -> &TokenKind {
        &self.current_token.kind
    }

    pub fn parse_expression(&mut self) -> Result<Node<UntypedKind>, String> {
        match self.peek() {
            TokenKind::LeftParen => self.parse_list(),
            TokenKind::LeftBracket => self.parse_vector_literal(),
            TokenKind::LeftBrace => self.parse_map_literal(),
            TokenKind::Quote => {
                let identity = Rc::new(NodeIdentity { location: self.current_token.location });
                self.advance()?; // consume '
                let expr = self.parse_expression()?;
                Ok(Node {
                    identity: identity.clone(),
                    kind: UntypedKind::Call {
                        callee: Box::new(self.make_id_node("quote", identity)),
                        args: vec![expr],
                    },
                    ty: (),
                })
            }
            _ => self.parse_atom(),
        }
    }

    fn parse_atom(&mut self) -> Result<Node<UntypedKind>, String> {
        let token = self.advance()?;
        let identity = Rc::new(NodeIdentity { location: token.location });

        let kind = match token.kind {
            TokenKind::Integer(n) => UntypedKind::Constant(Value::Int(n)),
            TokenKind::Float(n) => UntypedKind::Constant(Value::Float(n)),
            TokenKind::String(s) => UntypedKind::Constant(Value::Text(s)),
            TokenKind::Keyword(k) => UntypedKind::Constant(Value::Keyword(Keyword::intern(&k))),
            TokenKind::Identifier(id) => {
                match id.as_ref() {
                    "..." => UntypedKind::Nop,
                    "true" => UntypedKind::Constant(Value::Bool(true)),
                    "false" => UntypedKind::Constant(Value::Bool(false)),
                    "void" => UntypedKind::Constant(Value::Void),
                    _ => UntypedKind::Identifier(id),
                }
            }
            _ => return Err(format!("Unexpected token in atom: {:?} at {:?}", token.kind, token.location)),
        };

        Ok(Node { identity, kind, ty: () })
    }

    fn parse_list(&mut self) -> Result<Node<UntypedKind>, String> {
        let start_loc = self.advance()?.location; // consume '('
        let identity = Rc::new(NodeIdentity { location: start_loc });
        
        if *self.peek() == TokenKind::RightParen {
            return Err(format!("Empty list () is not a valid expression at {:?}", start_loc));
        }

        let head = self.parse_expression()?;
        
        let result = if let UntypedKind::Identifier(name) = &head.kind {
            match name.as_ref() {
                "if" => self.parse_if(identity),
                "fn" => self.parse_fn(identity),
                "def" => self.parse_def(identity),
                "assign" => self.parse_assign(identity),
                "do" => self.parse_do(identity),
                _ => self.parse_call(head, identity),
            }
        } else {
            self.parse_call(head, identity)
        };

        self.expect(TokenKind::RightParen)?;
        result
    }

    fn parse_if(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let cond = Box::new(self.parse_expression()?);
        let then_br = Box::new(self.parse_expression()?);
        let mut else_br = None;

        if *self.peek() != TokenKind::RightParen {
            else_br = Some(Box::new(self.parse_expression()?));
        }

        Ok(Node {
            identity,
            kind: UntypedKind::If { cond, then_br, else_br },
            ty: (),
        })
    }

    fn parse_def(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let name_node = self.parse_expression()?;
        let name = match name_node.kind {
            UntypedKind::Identifier(s) => s,
            _ => return Err("Expected identifier for def name".to_string()),
        };

        let value = Box::new(self.parse_expression()?);

        Ok(Node {
            identity,
            kind: UntypedKind::Def { name, value },
            ty: (),
        })
    }

    fn parse_assign(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        // (assign target value)
        let target = Box::new(self.parse_expression()?);
        let value = Box::new(self.parse_expression()?);
        
        Ok(Node {
            identity,
            kind: UntypedKind::Assign { target, value },
            ty: (),
        })
    }

    fn parse_do(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let mut exprs = Vec::new();
        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            exprs.push(self.parse_expression()?);
        }
        Ok(Node {
            identity,
            kind: UntypedKind::Block { exprs },
            ty: (),
        })
    }

    fn parse_fn(&mut self, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let params = self.parse_param_vector()?;
        let body = self.parse_expression()?;
        
        Ok(Node {
            identity,
            kind: UntypedKind::Lambda {
                params,
                body: Rc::new(body),
            },
            ty: (),
        })
    }

    fn parse_param_vector(&mut self) -> Result<Vec<Rc<str>>, String> {
        if *self.peek() != TokenKind::LeftBracket {
             return Err(format!("Expected parameter vector [...] for fn, found {:?}", self.peek()));
        }
        self.advance()?;

        let mut params = Vec::new();
        while *self.peek() != TokenKind::RightBracket {
            let token = self.advance()?;
            match token.kind {
                TokenKind::Identifier(name) => params.push(name),
                _ => return Err(format!("Expected identifier in param vector, found {:?}", token.kind)),
            }
        }
        self.expect(TokenKind::RightBracket)?;
        Ok(params)
    }

    fn parse_call(&mut self, callee: Node<UntypedKind>, identity: Identity) -> Result<Node<UntypedKind>, String> {
        let mut args = Vec::new();
        
        while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF {
            args.push(self.parse_expression()?);
        }

        Ok(Node {
            identity,
            kind: UntypedKind::Call { callee: Box::new(callee), args },
            ty: (),
        })
    }

    fn parse_vector_literal(&mut self) -> Result<Node<UntypedKind>, String> {
        let token = self.advance()?;
        let mut elements = Vec::new();
        let mut temp_ctx = Context::new();

        while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF {
            let expr = self.parse_expression()?;
            match expr.eval(&mut temp_ctx) {
                Ok(val) => elements.push(val),
                Err(e) => return Err(format!("Vector literal error: {}", e)),
            }
        }
        
        self.expect(TokenKind::RightBracket)?;
        
        Ok(Node {
            identity: Rc::new(NodeIdentity { location: token.location }),
            kind: UntypedKind::Constant(Value::List(Rc::new(elements))),
            ty: (),
        })
    }

    fn parse_map_literal(&mut self) -> Result<Node<UntypedKind>, String> {
        let token = self.advance()?;
        let mut map = std::collections::HashMap::new();
        let mut temp_ctx = Context::new();

        while *self.peek() != TokenKind::RightBrace {
            if *self.peek() == TokenKind::EOF {
                return Err("Unexpected EOF in map literal".to_string());
            }

            let key_node = self.parse_expression()?;
            let key = match key_node.eval(&mut temp_ctx)? {
                Value::Keyword(k) => k,
                _ => return Err("Map keys must be keywords".to_string()),
            };

            if *self.peek() == TokenKind::RightBrace {
                return Err("Map literal must have even number of forms".to_string());
            }
            let val_node = self.parse_expression()?;
            let val = val_node.eval(&mut temp_ctx)?;

            map.insert(key, val);
        }
        self.expect(TokenKind::RightBrace)?;

        Ok(Node {
            identity: Rc::new(NodeIdentity { location: token.location }),
            kind: UntypedKind::Constant(Value::Record(Rc::new(map))),
            ty: (),
        })
    }

    fn expect(&mut self, kind: TokenKind) -> Result<(), String> {
        let token = self.advance()?;
        if token.kind == kind {
            Ok(())
        } else {
            Err(format!("Expected {:?}, but found {:?} at {:?}", kind, token.kind, token.location))
        }
    }

    fn make_id_node(&self, name: &str, identity: Identity) -> Node<UntypedKind> {
        Node {
            identity,
            kind: UntypedKind::Identifier(Rc::from(name)),
            ty: (),
        }
    }
}