"""
Parser for rajini++ Language
Builds Abstract Syntax Tree from tokens
"""

from typing import List, Optional, Any
from lexer import Token, TokenType
from dataclasses import dataclass


# AST Node Classes
@dataclass
class ASTNode:
    """Base class for all AST nodes"""
    pass


@dataclass
class NumberNode(ASTNode):
    value: float


@dataclass
class StringNode(ASTNode):
    value: str


@dataclass
class BooleanNode(ASTNode):
    value: bool


@dataclass
class NullNode(ASTNode):
    pass


@dataclass
class IdentifierNode(ASTNode):
    name: str


@dataclass
class BinaryOpNode(ASTNode):
    left: ASTNode
    operator: str
    right: ASTNode


@dataclass
class UnaryOpNode(ASTNode):
    operator: str
    operand: ASTNode


@dataclass
class AssignmentNode(ASTNode):
    name: str
    value: ASTNode


@dataclass
class BlockNode(ASTNode):
    statements: List[ASTNode]


@dataclass
class IfNode(ASTNode):
    condition: ASTNode
    then_block: ASTNode
    else_block: Optional[ASTNode] = None


@dataclass
class WhileNode(ASTNode):
    condition: ASTNode
    body: ASTNode


@dataclass
class ForNode(ASTNode):
    init: Optional[ASTNode]
    condition: Optional[ASTNode]
    increment: Optional[ASTNode]
    body: ASTNode


@dataclass
class FunctionDefNode(ASTNode):
    name: str
    parameters: List[str]
    body: ASTNode


@dataclass
class FunctionCallNode(ASTNode):
    name: str
    arguments: List[ASTNode]


@dataclass
class ReturnNode(ASTNode):
    value: Optional[ASTNode]


@dataclass
class PrintNode(ASTNode):
    value: ASTNode


@dataclass
class ProgramNode(ASTNode):
    statements: List[ASTNode]


class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0
    
    def current_token(self) -> Token:
        """Get current token"""
        if self.pos >= len(self.tokens):
            return self.tokens[-1]  # Return EOF
        return self.tokens[self.pos]
    
    def peek_token(self, offset: int = 1) -> Token:
        """Peek ahead at token"""
        pos = self.pos + offset
        if pos >= len(self.tokens):
            return self.tokens[-1]
        return self.tokens[pos]
    
    def advance(self) -> Token:
        """Move to next token"""
        token = self.current_token()
        if self.pos < len(self.tokens) - 1:
            self.pos += 1
        return token
    
    def expect(self, token_type: TokenType) -> Token:
        """Expect a specific token type"""
        token = self.current_token()
        if token.type != token_type:
            raise SyntaxError(
                f"Expected {token_type.name}, got {token.type.name} "
                f"at line {token.line}, column {token.column}"
            )
        return self.advance()
    
    def parse(self) -> ProgramNode:
        """Parse the entire program"""
        statements = []
        while self.current_token().type != TokenType.EOF:
            stmt = self.parse_statement()
            if stmt:
                statements.append(stmt)
        return ProgramNode(statements)
    
    def parse_statement(self) -> Optional[ASTNode]:
        """Parse a single statement"""
        token = self.current_token()
        
        if token.type == TokenType.VAR:
            return self.parse_assignment()
        elif token.type == TokenType.FUNCTION:
            return self.parse_function_def()
        elif token.type == TokenType.START:
            return self.parse_start()
        elif token.type == TokenType.IF:
            return self.parse_if()
        elif token.type == TokenType.WHILE:
            return self.parse_while()
        elif token.type == TokenType.FOR:
            return self.parse_for()
        elif token.type == TokenType.RETURN:
            return self.parse_return()
        elif token.type == TokenType.PRINT:
            return self.parse_print()
        elif token.type == TokenType.LBRACE:
            return self.parse_block()
        elif token.type == TokenType.SEMICOLON:
            self.advance()
            return None
        else:
            # Expression statement
            expr = self.parse_expression()
            self.expect(TokenType.SEMICOLON)
            return expr
    
    def parse_assignment(self) -> AssignmentNode:
        """Parse variable assignment: vachuko x = value;"""
        self.expect(TokenType.VAR)
        name_token = self.expect(TokenType.IDENTIFIER)
        self.expect(TokenType.ASSIGN)
        value = self.parse_expression()
        self.expect(TokenType.SEMICOLON)
        return AssignmentNode(name_token.value, value)
    
    def parse_function_def(self) -> FunctionDefNode:
        """Parse function definition: sollu name(params) { body }"""
        self.expect(TokenType.FUNCTION)
        name_token = self.expect(TokenType.IDENTIFIER)
        
        self.expect(TokenType.LPAREN)
        parameters = []
        if self.current_token().type != TokenType.RPAREN:
            parameters.append(self.expect(TokenType.IDENTIFIER).value)
            while self.current_token().type == TokenType.COMMA:
                self.advance()
                parameters.append(self.expect(TokenType.IDENTIFIER).value)
        self.expect(TokenType.RPAREN)
        
        body = self.parse_block()
        return FunctionDefNode(name_token.value, parameters, body)
    
    def parse_start(self) -> FunctionDefNode:
        """Parse main entry point: aarambam { body }"""
        self.expect(TokenType.START)
        body = self.parse_block()
        return FunctionDefNode("__main__", [], body)
    
    def parse_if(self) -> IfNode:
        """Parse if statement: na_sonna (condition) { then } illa_na_sonna { else if } illana { else }"""
        self.expect(TokenType.IF)
        
        self.expect(TokenType.LPAREN)
        condition = self.parse_expression()
        self.expect(TokenType.RPAREN)
        
        then_block = self.parse_block()
        else_block = None
        
        # Check for else-if (illa_na_sonna) or else (illana)
        if self.current_token().type == TokenType.ELSEIF:
            self.advance()
            # Parse the else-if as a nested if statement
            self.expect(TokenType.LPAREN)
            elif_condition = self.parse_expression()
            self.expect(TokenType.RPAREN)
            elif_then = self.parse_block()
            
            # Check for more else-if or final else
            elif_else = None
            if self.current_token().type in (TokenType.ELSEIF, TokenType.ELSE):
                elif_else = self.parse_if_continuation()
            
            else_block = IfNode(elif_condition, elif_then, elif_else)
        elif self.current_token().type == TokenType.ELSE:
            self.advance()
            else_block = self.parse_block()
        
        return IfNode(condition, then_block, else_block)
    
    def parse_if_continuation(self) -> Optional[ASTNode]:
        """Parse else-if chain continuation"""
        if self.current_token().type == TokenType.ELSEIF:
            self.advance()
            self.expect(TokenType.LPAREN)
            condition = self.parse_expression()
            self.expect(TokenType.RPAREN)
            then_block = self.parse_block()
            
            else_block = None
            if self.current_token().type in (TokenType.ELSEIF, TokenType.ELSE):
                else_block = self.parse_if_continuation()
            
            return IfNode(condition, then_block, else_block)
        elif self.current_token().type == TokenType.ELSE:
            self.advance()
            return self.parse_block()
        return None
    
    def parse_while(self) -> WhileNode:
        """Parse while loop: varaikkum (condition) { body }"""
        self.expect(TokenType.WHILE)
        self.expect(TokenType.LPAREN)
        condition = self.parse_expression()
        self.expect(TokenType.RPAREN)
        body = self.parse_block()
        return WhileNode(condition, body)
    
    def parse_for(self) -> ForNode:
        """Parse for loop: ellathukum (init; condition; increment) { body }"""
        self.expect(TokenType.FOR)
        self.expect(TokenType.LPAREN)
        
        # Parse init
        init = None
        if self.current_token().type == TokenType.IDENTIFIER:
            name_token = self.advance()
            self.expect(TokenType.ASSIGN)
            value = self.parse_expression()
            init = AssignmentNode(name_token.value, value)
        
        self.expect(TokenType.SEMICOLON)
        
        # Parse condition
        condition = None
        if self.current_token().type != TokenType.SEMICOLON:
            condition = self.parse_expression()
        
        self.expect(TokenType.SEMICOLON)
        
        # Parse increment
        increment = None
        if self.current_token().type != TokenType.RPAREN:
            name_token = self.expect(TokenType.IDENTIFIER)
            self.expect(TokenType.ASSIGN)
            value = self.parse_expression()
            increment = AssignmentNode(name_token.value, value)
        
        self.expect(TokenType.RPAREN)
        body = self.parse_block()
        
        return ForNode(init, condition, increment, body)
    
    def parse_return(self) -> ReturnNode:
        """Parse return statement: thirupi_kudu value;"""
        self.expect(TokenType.RETURN)
        value = None
        if self.current_token().type != TokenType.SEMICOLON:
            value = self.parse_expression()
        self.expect(TokenType.SEMICOLON)
        return ReturnNode(value)
    
    def parse_print(self) -> PrintNode:
        """Parse print statement: sollu_da(value);"""
        self.expect(TokenType.PRINT)
        self.expect(TokenType.LPAREN)
        value = self.parse_expression()
        self.expect(TokenType.RPAREN)
        self.expect(TokenType.SEMICOLON)
        return PrintNode(value)
    
    def parse_block(self) -> BlockNode:
        """Parse code block: { statements }"""
        self.expect(TokenType.LBRACE)
        statements = []
        while self.current_token().type != TokenType.RBRACE:
            stmt = self.parse_statement()
            if stmt:
                statements.append(stmt)
        self.expect(TokenType.RBRACE)
        return BlockNode(statements)
    
    def parse_expression(self) -> ASTNode:
        """Parse expression"""
        return self.parse_logical_or()
    
    def parse_logical_or(self) -> ASTNode:
        """Parse logical OR: expr || expr"""
        left = self.parse_logical_and()
        
        while self.current_token().type == TokenType.OR:
            op_token = self.advance()
            right = self.parse_logical_and()
            left = BinaryOpNode(left, '||', right)
        
        return left
    
    def parse_logical_and(self) -> ASTNode:
        """Parse logical AND: expr && expr"""
        left = self.parse_equality()
        
        while self.current_token().type == TokenType.AND:
            op_token = self.advance()
            right = self.parse_equality()
            left = BinaryOpNode(left, '&&', right)
        
        return left
    
    def parse_equality(self) -> ASTNode:
        """Parse equality: expr == expr, expr != expr"""
        left = self.parse_comparison()
        
        while self.current_token().type in (TokenType.EQUAL, TokenType.NOT_EQUAL):
            op_token = self.advance()
            right = self.parse_comparison()
            left = BinaryOpNode(left, op_token.value, right)
        
        return left
    
    def parse_comparison(self) -> ASTNode:
        """Parse comparison: expr > expr, expr < expr, etc."""
        left = self.parse_term()
        
        while self.current_token().type in (
            TokenType.GREATER, TokenType.LESS,
            TokenType.GREATER_EQUAL, TokenType.LESS_EQUAL
        ):
            op_token = self.advance()
            right = self.parse_term()
            left = BinaryOpNode(left, op_token.value, right)
        
        return left
    
    def parse_term(self) -> ASTNode:
        """Parse addition/subtraction: expr + expr, expr - expr"""
        left = self.parse_factor()
        
        while self.current_token().type in (TokenType.PLUS, TokenType.MINUS):
            op_token = self.advance()
            right = self.parse_factor()
            left = BinaryOpNode(left, op_token.value, right)
        
        return left
    
    def parse_factor(self) -> ASTNode:
        """Parse multiplication/division: expr * expr, expr / expr, expr % expr"""
        left = self.parse_unary()
        
        while self.current_token().type in (TokenType.MULTIPLY, TokenType.DIVIDE, TokenType.MODULO):
            op_token = self.advance()
            right = self.parse_unary()
            left = BinaryOpNode(left, op_token.value, right)
        
        return left
    
    def parse_unary(self) -> ASTNode:
        """Parse unary operations: -expr, !expr"""
        if self.current_token().type in (TokenType.MINUS, TokenType.NOT):
            op_token = self.advance()
            operand = self.parse_unary()
            return UnaryOpNode(op_token.value, operand)
        
        return self.parse_primary()
    
    def parse_primary(self) -> ASTNode:
        """Parse primary expressions: literals, identifiers, function calls, parentheses"""
        token = self.current_token()
        
        # Numbers
        if token.type == TokenType.NUMBER:
            self.advance()
            return NumberNode(token.value)
        
        # Strings
        elif token.type == TokenType.STRING:
            self.advance()
            return StringNode(token.value)
        
        # Booleans
        elif token.type == TokenType.TRUE:
            self.advance()
            return BooleanNode(True)
        
        elif token.type == TokenType.FALSE:
            self.advance()
            return BooleanNode(False)
        
        # Null
        elif token.type == TokenType.NULL:
            self.advance()
            return NullNode()
        
        # Identifier or function call
        elif token.type == TokenType.IDENTIFIER:
            name = token.value
            self.advance()
            
            # Function call
            if self.current_token().type == TokenType.LPAREN:
                self.advance()
                arguments = []
                if self.current_token().type != TokenType.RPAREN:
                    arguments.append(self.parse_expression())
                    while self.current_token().type == TokenType.COMMA:
                        self.advance()
                        arguments.append(self.parse_expression())
                self.expect(TokenType.RPAREN)
                return FunctionCallNode(name, arguments)
            else:
                return IdentifierNode(name)
        
        # Parenthesized expression
        elif token.type == TokenType.LPAREN:
            self.advance()
            expr = self.parse_expression()
            self.expect(TokenType.RPAREN)
            return expr
        
        else:
            raise SyntaxError(
                f"Unexpected token {token.type.name} at line {token.line}, column {token.column}"
            )
