"""
transpiler/purity.py — Purity Analysis and Effect Lattice.
"""

from transpiler.middleend import ir
from typing import Dict, Set

class PurityAnalyzer:
    def __init__(self):
        # Whitelist of pure functions / built-ins
        self.pure_builtins = {
            "math.abs", "math.acos", "math.asin", "math.atan", "math.atan2",
            "math.ceil", "math.cos", "math.cosh", "math.deg", "math.exp",
            "math.floor", "math.fmod", "math.frexp", "math.ldexp", "math.log",
            "math.log10", "math.max", "math.min", "math.modf", "math.pow",
            "math.rad", "math.sin", "math.sinh", "math.sqrt", "math.tan",
            "math.tanh", "type", "typeof", "len", "py_len", "str", "py_str",
            "int", "py_int", "float", "py_float"
        }
        
    def analyze_module(self, module: ir.IRModule):
        """Annotates all expressions in the module with their effect level."""
        for stmt in module.body:
            self.analyze_stmt(stmt)
            
    def analyze_stmt(self, stmt: ir.IRStatement):
        if isinstance(stmt, ir.IRAssignment):
            self.analyze_expr(stmt.target)
            self.analyze_expr(stmt.value)
        elif isinstance(stmt, ir.IRIfStatement):
            self.analyze_expr(stmt.condition)
            for s in stmt.then_block: self.analyze_stmt(s)
            for cond, body in stmt.elif_blocks:
                self.analyze_expr(cond)
                for s in body: self.analyze_stmt(s)
            if stmt.else_block:
                for s in stmt.else_block: self.analyze_stmt(s)
        elif isinstance(stmt, (ir.IRWhileStatement, ir.IRForStatement)):
            if hasattr(stmt, "condition"): self.analyze_expr(stmt.condition)
            if hasattr(stmt, "start"): self.analyze_expr(stmt.start)
            if hasattr(stmt, "end"): self.analyze_expr(stmt.end)
            if hasattr(stmt, "step") and stmt.step: self.analyze_expr(stmt.step)
            for s in stmt.body: self.analyze_stmt(s)
        elif isinstance(stmt, ir.IRFunctionDef):
            # We don't do interprocedural purity yet, assume unknown/has side effects
            # but we still analyze the body to annotate internal expressions
            for s in stmt.body: self.analyze_stmt(s)
        elif isinstance(stmt, ir.IRReturnStatement) and stmt.value:
            self.analyze_expr(stmt.value)
        elif isinstance(stmt, ir.IRModule):
            for s in stmt.body: self.analyze_stmt(s)
        elif hasattr(stmt, "__dict__"):
            # Generic traversal for other statement types (Break, Continue, etc)
            pass

    def analyze_expr(self, expr: ir.IRExpression) -> ir.EffectLevel:
        if isinstance(expr, ir.IRLiteral):
            expr.effect_level = ir.EffectLevel.PURE
        elif isinstance(expr, ir.IRVariable):
            # Reading a variable is pure (SSA handles definitions/uses)
            expr.effect_level = ir.EffectLevel.PURE
        elif isinstance(expr, ir.IRBinaryOperation):
            l = self.analyze_expr(expr.left)
            r = self.analyze_expr(expr.right)
            expr.effect_level = self._combine_effects(l, r)
        elif isinstance(expr, ir.IRUnaryOperation):
            o = self.analyze_expr(expr.operand)
            expr.effect_level = o
        elif isinstance(expr, ir.IRFunctionCall):
            # 1. Analyze arguments
            max_arg_effect = ir.EffectLevel.PURE
            for arg in expr.args:
                max_arg_effect = self._combine_effects(max_arg_effect, self.analyze_expr(arg))
            
            # 2. Check function purity
            fn_name = self._get_fn_name(expr.func)
            if fn_name in self.pure_builtins:
                # Built-in pure functions are PURE if their arguments are PURE.
                # If args READ_MEMORY, then the call READS_MEMORY.
                expr.effect_level = self._combine_effects(ir.EffectLevel.PURE, max_arg_effect)
            else:
                # Unknown calls are assumed to have side effects
                expr.effect_level = ir.EffectLevel.HAS_SIDE_EFFECT
                
        elif isinstance(expr, ir.IRMethodCall):
            # Recurse
            self.analyze_expr(expr.receiver)
            for arg in expr.args: self.analyze_expr(arg)
            # Methods can trigger metamethods or mutate state
            expr.effect_level = ir.EffectLevel.HAS_SIDE_EFFECT
            
        elif isinstance(expr, ir.IRPropertyAccess):
            self.analyze_expr(expr.receiver)
            # Property access can trigger __index metamethod
            expr.effect_level = ir.EffectLevel.READS_MEMORY
            
        elif isinstance(expr, ir.IRIndexOperation):
            self.analyze_expr(expr.receiver)
            self.analyze_expr(expr.index)
            # Indexing can trigger __index
            expr.effect_level = ir.EffectLevel.READS_MEMORY
            
        elif isinstance(expr, ir.IRTableLiteral):
            max_eff = ir.EffectLevel.PURE
            for el in expr.elements:
                max_eff = self._combine_effects(max_eff, self.analyze_expr(el))
            for k, v in expr.fields:
                max_eff = self._combine_effects(max_eff, self.analyze_expr(v))
            # Table allocation is deterministic but introduces new identity.
            # For LICM/CSE we can treat it as PURE-ish if elements are pure,
            # but to be safe with identity we'll call it PURE + allocation overhead.
            expr.effect_level = max_eff
            
        elif isinstance(expr, ir.IRPhi):
            max_phi_eff = ir.EffectLevel.PURE
            for v in expr.options.values():
                max_phi_eff = self._combine_effects(max_phi_eff, self.analyze_expr(v))
            expr.effect_level = max_phi_eff
            
        elif isinstance(expr, ir.IRIfElseExpression):
            cond_eff = self.analyze_expr(expr.condition)
            then_eff = self.analyze_expr(expr.then_expr)
            else_eff = self.analyze_expr(expr.else_expr)
            expr.effect_level = self._combine_effects(cond_eff, self._combine_effects(then_eff, else_eff))
            
        else:
            # Fallback for complex/unknown expressions: gather effects from all children
            max_fallback_eff = ir.EffectLevel.PURE
            if hasattr(expr, "get_children"):
                for child in expr.get_children():
                    if isinstance(child, ir.IRExpression):
                        max_fallback_eff = self._combine_effects(max_fallback_eff, self.analyze_expr(child))
            expr.effect_level = self._combine_effects(max_fallback_eff, ir.EffectLevel.HAS_SIDE_EFFECT)
            
        return expr.effect_level

    def _combine_effects(self, e1: ir.EffectLevel, e2: ir.EffectLevel) -> ir.EffectLevel:
        # Lattice: PURE < READS < WRITES < HAS_SIDE_EFFECT < UNKNOWN
        return ir.EffectLevel(max(e1.value, e2.value))

    def _get_fn_name(self, expr: ir.IRExpression) -> str:
        if isinstance(expr, ir.IRVariable):
            return expr.name
        if isinstance(expr, ir.IRPropertyAccess):
            recv = self._get_fn_name(expr.receiver)
            if recv:
                return f"{recv}.{expr.property}"
        return ""

def analyze_purity(module: ir.IRModule):
    analyzer = PurityAnalyzer()
    analyzer.analyze_module(module)
