lyanglyang 0.1.1

use crate::ast::{Statement, Condition, Value, StrSegment};
use crate::error::NepalError;
use crate::token::Token;

#[derive(Clone, Copy)]
enum MugOp {
    Add,
    Sub,
    Mul,
    Div,
}

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

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

    pub fn parse(&mut self) -> Result<Vec<Statement>, NepalError> {
        let mut statements = Vec::new();
        while self.position < self.tokens.len() {
            statements.push(self.parse_statement()?);
        }
        Ok(statements)
    }

    fn current(&self) -> Result<&Token, NepalError> {
        self.tokens
            .get(self.position)
            .ok_or(NepalError::ParseError("unexpected end of input"))
    }

    fn parse_statement(&mut self) -> Result<Statement, NepalError> {
        match self.current()? {
            Token::OiMug => {
                self.position += 1;
                if matches!(self.tokens.get(self.position), Some(Token::Bhan)) {
                    self.position += 1;
                    self.parse_input()
                } else {
                    self.parse_declaration()
                }
            },
            Token::BolMug => self.parse_print(),
            Token::Mug => {
                self.position += 1;  // Skip 'mug'
                match self.current()? {
                    Token::Jod => {
                        self.position += 1;
                        self.parse_standalone_arithmetic(MugOp::Add)
                    }
                    Token::Ghata => {
                        self.position += 1;
                        self.parse_standalone_arithmetic(MugOp::Sub)
                    }
                    Token::Guna => {
                        self.position += 1;
                        self.parse_standalone_arithmetic(MugOp::Mul)
                    }
                    Token::Bhag => {
                        self.position += 1;
                        self.parse_standalone_arithmetic(MugOp::Div)
                    }
                    _ => Err(NepalError::ParseError(
                        "Expected 'jod', 'ghata', 'guna', or 'bhag' after 'mug'",
                    )),
                }
            }
            Token::Yedi => self.parse_if_statement(),
            Token::Aile => {
                self.position += 1;
                if let Some(Token::Feri) = self.tokens.get(self.position) {
                    self.position += 1;
                    self.parse_if_statement()
                } else {
                    Err(NepalError::ParseError("Expected 'feri' after 'aile'"))
                }
            },
            _ => Err(NepalError::ParseError("Unexpected token")),
        }
    }

    fn parse_if_statement(&mut self) -> Result<Statement, NepalError> {
        if matches!(self.tokens.get(self.position), Some(Token::Yedi)) {
            self.position += 1;
        }
        
        let var1 = if let Token::Identifier(name) = self.current()? {
            let name = name.clone();
            self.position += 1;
            name
        } else {
            return Err(NepalError::ParseError("Expected variable name after yedi"));
        };

        let condition = match self.current()? {
            Token::Babaal => {
                self.position += 1;
                true
            }
            Token::Laamo => {
                self.position += 1;
                false
            }
            _ => return Err(NepalError::ParseError("Expected babaal or laamo")),
        };

        let var2 = if let Token::String(s) = self.current()? {
            let s = s.clone();
            self.position += 1;
            s
        } else {
            return Err(NepalError::ParseError("Expected string literal"));
        };

        if !matches!(self.tokens.get(self.position), Some(Token::Bhane)) {
            return Err(NepalError::ParseError("Expected 'bhane' after condition"));
        }
        self.position += 1;

        let mut statements = Vec::new();
        let mut else_branch: Option<Box<Statement>> = None;
        while self.position < self.tokens.len() {
            match &self.tokens[self.position] {
                Token::Sakiyo => {
                    self.position += 1;
                    break;
                }
                Token::Aile => {
                    // Else / else-if branch. parse_statement handles the
                    // `Aile` token: it consumes `aile feri` and parses the
                    // following if-statement, which recursively absorbs any
                    // further `aile feri ... sakiyo` chain (and its own
                    // trailing `sakiyo`). We attach the result as the else
                    // branch instead of leaving it orphaned at top level.
                    else_branch = Some(Box::new(self.parse_statement()?));
                    break;
                }
                _ => statements.push(self.parse_statement()?),
            }
        }

        Ok(Statement::If(
            if condition {
                Condition::Equals(var1, var2)
            } else {
                Condition::NotEquals(var1, var2)
            },
            statements,
            else_branch,
        ))
    }

    fn parse_declaration(&mut self) -> Result<Statement, NepalError> {
        let name = if let Token::Identifier(name) = self.current()? {
            let name = name.clone();
            self.position += 1;
            name
        } else {
            return Err(NepalError::ParseError("Expected identifier"));
        };

        if let Token::Equals = self.current()? {
            self.position += 1;
        } else {
            return Err(NepalError::ParseError("Expected '='"));
        }

        match self.current()?.clone() {
            Token::Number(value) => {
                self.position += 1;
                Ok(Statement::Declaration(name, Value::Number(value)))
            }
            Token::String(value) => {
                self.position += 1;
                if matches!(self.tokens.get(self.position), Some(Token::Plus)) {
                    self.position += 1;
                    let mut parts = vec![StrSegment::Literal(value.clone())];
                    
                    while self.position < self.tokens.len() {
                        match &self.tokens[self.position] {
                            Token::String(s) => {
                                parts.push(StrSegment::Literal(s.clone()));
                                self.position += 1;
                            }
                            Token::Identifier(id) => {
                                parts.push(StrSegment::Identifier(id.clone()));
                                self.position += 1;
                            }
                            Token::Plus => {
                                self.position += 1;
                                continue;
                            }
                            _ => break,
                        }
                    }
                    Ok(Statement::StringConcat(name, parts))
                } else {
                    Ok(Statement::Declaration(name, Value::String(value.clone())))
                }
            }
            Token::Identifier(value) => {
                self.position += 1;
                if matches!(self.tokens.get(self.position), Some(Token::Plus)) {
                    self.position += 1;
                    let mut parts = vec![StrSegment::Identifier(value.clone())];
                    
                    while self.position < self.tokens.len() {
                        match &self.tokens[self.position] {
                            Token::String(s) => {
                                parts.push(StrSegment::Literal(s.clone()));
                                self.position += 1;
                            }
                            Token::Identifier(id) => {
                                parts.push(StrSegment::Identifier(id.clone()));
                                self.position += 1;
                            }
                            Token::Plus => {
                                self.position += 1;
                                continue;
                            }
                            _ => break,
                        }
                    }
                    Ok(Statement::StringConcat(name, parts))
                } else {
                    Ok(Statement::Declaration(name, Value::String(value.clone())))
                }
            }
            Token::Jod | Token::Ghata | Token::Guna | Token::Bhag => {
                self.parse_arithmetic_operation(name)
            }
            _ => Err(NepalError::ParseError("Expected value")),
        }
    }

    fn parse_arithmetic_operation(&mut self, target: String) -> Result<Statement, NepalError> {
        let op = self.tokens[self.position].clone();
        self.position += 1;

        let mut sources = Vec::new();
        while self.position < self.tokens.len() {
            match &self.tokens[self.position] {
                Token::Identifier(id) => {
                    sources.push(id.clone());
                    self.position += 1;
                    if matches!(self.tokens.get(self.position), Some(Token::Comma)) {
                        self.position += 1;
                    } else {
                        break;
                    }
                }
                _ => break,
            }
        }

        if sources.is_empty() {
            return Err(NepalError::ParseError(
                "arithmetic operation requires at least one operand",
            ));
        }

        match op {
            Token::Jod => Ok(Statement::Addition(target, sources)),
            Token::Ghata => Ok(Statement::Subtraction(target, sources)),
            Token::Guna => Ok(Statement::Multiplication(target, sources)),
            Token::Bhag => Ok(Statement::Division(target, sources)),
            _ => unreachable!(),
        }
    }

    fn parse_standalone_arithmetic(&mut self, op: MugOp) -> Result<Statement, NepalError> {
        let mut sources = Vec::new();
        
        while self.position < self.tokens.len() {
            match &self.tokens[self.position] {
                Token::Identifier(name) => {
                    sources.push(name.clone());
                    self.position += 1;
                    
                    match self.tokens.get(self.position) {
                        Some(Token::Comma) => {
                            self.position += 1;
                            continue;
                        }
                        Some(Token::Lai) => {
                            self.position += 1;
                            break;
                        }
                        _ => return Err(NepalError::ParseError("Expected comma or 'lai'")),
                    }
                }
                _ => return Err(NepalError::ParseError("Expected identifier")),
            }
        }

        let target = if let Token::Identifier(name) = self.current()? {
            let name = name.clone();
            self.position += 1;
            name
        } else {
            return Err(NepalError::ParseError("Expected target identifier after 'lai'"));
        };

        if sources.is_empty() {
            return Err(NepalError::ParseError(
                "arithmetic operation requires at least one operand",
            ));
        }

        match op {
            MugOp::Add => Ok(Statement::Addition(target, sources)),
            MugOp::Sub => Ok(Statement::Subtraction(target, sources)),
            MugOp::Mul => Ok(Statement::Multiplication(target, sources)),
            MugOp::Div => Ok(Statement::Division(target, sources)),
        }
    }

    fn parse_print(&mut self) -> Result<Statement, NepalError> {
        self.position += 1; // Skip 'bol mug'
        if let Token::String(_) = self.current()? {
            let mut parts = Vec::new();
            while self.position < self.tokens.len() {
                match &self.tokens[self.position] {
                    Token::String(s) => {
                        parts.push(s.clone());
                        self.position += 1;
                    }
                    Token::Plus => {
                        self.position += 1;
                    }
                    Token::Identifier(name) => {
                        parts.push(format!("{{{}}}", name));
                        self.position += 1;
                    }
                    _ => break,
                }
            }
            Ok(Statement::PrintString(parts))
        } else {
            let name = if let Token::Identifier(name) = self.current()? {
                let name = name.clone();
                self.position += 1;
                name
            } else {
                return Err(NepalError::ParseError("Expected identifier or string"));
            };
            Ok(Statement::Print(name))
        }
    }

    fn parse_input(&mut self) -> Result<Statement, NepalError> {
        let name = if let Token::Identifier(name) = self.current()? {
            let name = name.clone();
            self.position += 1;
            name
        } else {
            return Err(NepalError::ParseError("Expected identifier after bhan"));
        };
        Ok(Statement::Input(name))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::StrSegment;
    use crate::token::Token;

    #[test]
    fn string_concat_preserves_literal_segments() {
        let tokens = vec![
            Token::OiMug,
            Token::Identifier("fullName".into()),
            Token::Equals,
            Token::Identifier("firstName".into()),
            Token::Plus,
            Token::String(" ".into()),
            Token::Plus,
            Token::Identifier("lastName".into()),
        ];
        let mut p = Parser::new(tokens);
        let stmts = p.parse().unwrap();
        match &stmts[0] {
            Statement::StringConcat(name, parts) => {
                assert_eq!(name, "fullName");
                assert_eq!(parts.len(), 3);
                assert!(matches!(&parts[0], StrSegment::Identifier(s) if s == "firstName"));
                assert!(matches!(&parts[1], StrSegment::Literal(s) if s == " "));
                assert!(matches!(&parts[2], StrSegment::Identifier(s) if s == "lastName"));
            }
            _ => panic!("expected StringConcat"),
        }
    }

    #[test]
    fn parses_simple_if() {
        let tokens = vec![
            Token::Yedi,
            Token::Identifier("age".into()),
            Token::Babaal,
            Token::String("18".into()),
            Token::Bhane,
            Token::BolMug,
            Token::String("ok".into()),
            Token::Sakiyo,
        ];
        let mut p = Parser::new(tokens);
        let stmts = p.parse().unwrap();
        assert_eq!(stmts.len(), 1);
        assert!(matches!(stmts[0], Statement::If(_, _, None)));
    }
}