"""
Interpreter for rajini++ Language
Executes the Abstract Syntax Tree
"""

from typing import Any, Dict, Optional
from parser import (
    ASTNode, NumberNode, StringNode, BooleanNode, NullNode,
    IdentifierNode, BinaryOpNode, UnaryOpNode, AssignmentNode,
    BlockNode, IfNode, WhileNode, ForNode, FunctionDefNode,
    FunctionCallNode, ReturnNode, PrintNode, ProgramNode
)


class ReturnException(Exception):
    """Exception to handle return statements"""
    def __init__(self, value):
        self.value = value


class Function:
    """Represents a user-defined function"""
    def __init__(self, name: str, parameters: list, body: ASTNode, closure: Dict[str, Any]):
        self.name = name
        self.parameters = parameters
        self.body = body
        self.closure = closure


class Interpreter:
    def __init__(self):
        self.global_scope = {}
        self.current_scope = self.global_scope
        self.scope_stack = []
        
        # Built-in functions
        self.builtins = {
            'len': self.builtin_len,
            'str': self.builtin_str,
            'int': self.builtin_int,
            'float': self.builtin_float,
            'type': self.builtin_type,
        }
    
    # Built-in functions
    def builtin_len(self, value):
        """Get length of string or list"""
        if isinstance(value, (str, list)):
            return len(value)
        raise TypeError(f"len() not supported for {type(value).__name__}")
    
    def builtin_str(self, value):
        """Convert to string"""
        return str(value)
    
    def builtin_int(self, value):
        """Convert to integer"""
        return int(value)
    
    def builtin_float(self, value):
        """Convert to float"""
        return float(value)
    
    def builtin_type(self, value):
        """Get type of value"""
        return type(value).__name__
    
    def push_scope(self, new_scope: Optional[Dict[str, Any]] = None):
        """Push a new scope onto the stack"""
        self.scope_stack.append(self.current_scope)
        if new_scope is None:
            new_scope = {}
        self.current_scope = new_scope
    
    def pop_scope(self):
        """Pop the current scope"""
        if self.scope_stack:
            self.current_scope = self.scope_stack.pop()
    
    def get_variable(self, name: str) -> Any:
        """Get variable value from current or parent scopes"""
        # Check current scope
        if name in self.current_scope:
            return self.current_scope[name]
        
        # Check parent scopes
        for scope in reversed(self.scope_stack):
            if name in scope:
                return scope[name]
        
        # Check global scope
        if name in self.global_scope:
            return self.global_scope[name]
        
        # Check builtins
        if name in self.builtins:
            return self.builtins[name]
        
        raise NameError(f"Variable '{name}' is not defined")
    
    def set_variable(self, name: str, value: Any):
        """Set variable in current scope"""
        self.current_scope[name] = value
    
    def interpret(self, node: ASTNode) -> Any:
        """Interpret an AST node"""
        
        if isinstance(node, ProgramNode):
            result = None
            for statement in node.statements:
                result = self.interpret(statement)
            return result
        
        elif isinstance(node, NumberNode):
            return node.value
        
        elif isinstance(node, StringNode):
            return node.value
        
        elif isinstance(node, BooleanNode):
            return node.value
        
        elif isinstance(node, NullNode):
            return None
        
        elif isinstance(node, IdentifierNode):
            return self.get_variable(node.name)
        
        elif isinstance(node, BinaryOpNode):
            left = self.interpret(node.left)
            right = self.interpret(node.right)
            return self.eval_binary_op(left, node.operator, right)
        
        elif isinstance(node, UnaryOpNode):
            operand = self.interpret(node.operand)
            return self.eval_unary_op(node.operator, operand)
        
        elif isinstance(node, AssignmentNode):
            value = self.interpret(node.value)
            self.set_variable(node.name, value)
            return value
        
        elif isinstance(node, BlockNode):
            result = None
            for statement in node.statements:
                result = self.interpret(statement)
            return result
        
        elif isinstance(node, IfNode):
            condition = self.interpret(node.condition)
            if self.is_truthy(condition):
                return self.interpret(node.then_block)
            elif node.else_block:
                return self.interpret(node.else_block)
            return None
        
        elif isinstance(node, WhileNode):
            result = None
            while self.is_truthy(self.interpret(node.condition)):
                result = self.interpret(node.body)
            return result
        
        elif isinstance(node, ForNode):
            result = None
            
            # Initialize
            if node.init:
                self.interpret(node.init)
            
            # Loop
            while True:
                # Check condition
                if node.condition:
                    if not self.is_truthy(self.interpret(node.condition)):
                        break
                
                # Execute body
                result = self.interpret(node.body)
                
                # Increment
                if node.increment:
                    self.interpret(node.increment)
            
            return result
        
        elif isinstance(node, FunctionDefNode):
            # Create function object and store in current scope
            func = Function(
                node.name,
                node.parameters,
                node.body,
                dict(self.current_scope)  # Capture closure
            )
            self.set_variable(node.name, func)
            return func
        
        elif isinstance(node, FunctionCallNode):
            func = self.get_variable(node.name)
            
            # Built-in function
            if callable(func) and not isinstance(func, Function):
                args = [self.interpret(arg) for arg in node.arguments]
                return func(*args)
            
            # User-defined function
            if isinstance(func, Function):
                args = [self.interpret(arg) for arg in node.arguments]
                
                # Check parameter count
                if len(args) != len(func.parameters):
                    raise TypeError(
                        f"Function '{func.name}' expects {len(func.parameters)} "
                        f"arguments, got {len(args)}"
                    )
                
                # Create new scope for function
                func_scope = dict(func.closure)
                for param, arg in zip(func.parameters, args):
                    func_scope[param] = arg
                
                # Execute function
                self.push_scope(func_scope)
                try:
                    self.interpret(func.body)
                    result = None
                except ReturnException as ret:
                    result = ret.value
                finally:
                    self.pop_scope()
                
                return result
            
            raise TypeError(f"'{node.name}' is not callable")
        
        elif isinstance(node, ReturnNode):
            value = None
            if node.value:
                value = self.interpret(node.value)
            raise ReturnException(value)
        
        elif isinstance(node, PrintNode):
            value = self.interpret(node.value)
            output = self.value_to_string(value)
            print(output, flush=True)  # Force flush output
            return value
        
        else:
            raise RuntimeError(f"Unknown node type: {type(node).__name__}")
    
    def eval_binary_op(self, left: Any, operator: str, right: Any) -> Any:
        """Evaluate binary operation"""
        if operator == '+':
            return left + right
        elif operator == '-':
            return left - right
        elif operator == '*':
            return left * right
        elif operator == '/':
            if right == 0:
                raise ZeroDivisionError("Division by zero")
            return left / right
        elif operator == '%':
            return left % right
        elif operator == '==':
            return left == right
        elif operator == '!=':
            return left != right
        elif operator == '>':
            return left > right
        elif operator == '<':
            return left < right
        elif operator == '>=':
            return left >= right
        elif operator == '<=':
            return left <= right
        elif operator == '&&':
            return self.is_truthy(left) and self.is_truthy(right)
        elif operator == '||':
            return self.is_truthy(left) or self.is_truthy(right)
        else:
            raise RuntimeError(f"Unknown binary operator: {operator}")
    
    def eval_unary_op(self, operator: str, operand: Any) -> Any:
        """Evaluate unary operation"""
        if operator == '-':
            return -operand
        elif operator == '!':
            return not self.is_truthy(operand)
        else:
            raise RuntimeError(f"Unknown unary operator: {operator}")
    
    def is_truthy(self, value: Any) -> bool:
        """Determine if value is truthy"""
        if value is None:
            return False
        if isinstance(value, bool):
            return value
        if isinstance(value, (int, float)):
            return value != 0
        if isinstance(value, str):
            return len(value) > 0
        return True
    
    def value_to_string(self, value: Any) -> str:
        """Convert value to string for printing"""
        if value is None:
            return "onnum_illa"
        elif isinstance(value, bool):
            return "aama" if value else "illa"
        elif isinstance(value, str):
            return value
        elif isinstance(value, (int, float)):
            return str(value)
        elif isinstance(value, Function):
            return f"<function {value.name}>"
        else:
            return str(value)
    
    def run(self, program: ProgramNode):
        """Run the entire program"""
        try:
            # Interpret all statements (this defines functions)
            self.interpret(program)
            
            # If there's a __main__ function, call it
            if "__main__" in self.global_scope:
                main_func = self.global_scope["__main__"]
                if isinstance(main_func, Function):
                    # Create new scope for main
                    self.push_scope({})
                    try:
                        self.interpret(main_func.body)
                    except ReturnException:
                        pass  # Main returned, that's okay
                    finally:
                        self.pop_scope()
        except Exception as e:
            print(f"Runtime Error: {e}")
            raise
