from transpiler.middleend import ir
from typing import Set, Any, Optional


# ---------------------------------------------------------------------------
# Purity / mutation analysis helpers
# ---------------------------------------------------------------------------

def _is_pure_expr(expr: ir.IRExpression) -> bool:
    """Return True if *expr* is free of side-effects and always evaluates to
    the same value (i.e. no function/method calls, no property mutations).

    Conservatively returns False for anything that isn't a literal, a simple
    variable reference, or a binary/unary operation composed of pure parts.
    """
    if isinstance(expr, ir.IRLiteral):
        return True
    if isinstance(expr, ir.IRVariable):
        return True
    if isinstance(expr, ir.IRBinaryOperation):
        return _is_pure_expr(expr.left) and _is_pure_expr(expr.right)
    if isinstance(expr, ir.IRUnaryOperation):
        return _is_pure_expr(expr.operand)
    if isinstance(expr, ir.IRIfElseExpression):
        return _is_pure_expr(expr.condition) and _is_pure_expr(expr.then_expr) and _is_pure_expr(expr.else_expr)
    # Anything else (calls, subscripts, method calls…) is not pure.
    return False


def _collect_var_mutations(stmts: list, var_name: str, result: Set[int]):
    """Recursively walk *stmts* and append the index of any statement that
    assigns to *var_name* into *result* (using the object identity of the
    stmt as key to avoid duplicates across nested blocks).
    """
    for stmt in stmts:
        if isinstance(stmt, ir.IRAssignment):
            if isinstance(stmt.target, ir.IRVariable) and stmt.target.name == var_name:
                result.add(id(stmt))
        # Recurse into nested blocks
        if hasattr(stmt, "body") and isinstance(stmt.body, list):
            _collect_var_mutations(stmt.body, var_name, result)
        if hasattr(stmt, "then_block") and isinstance(stmt.then_block, list):
            _collect_var_mutations(stmt.then_block, var_name, result)
        if hasattr(stmt, "else_block") and stmt.else_block:
            _collect_var_mutations(stmt.else_block, var_name, result)
        if hasattr(stmt, "elif_blocks"):
            for _cond, body in stmt.elif_blocks:
                _collect_var_mutations(body, var_name, result)


def _contains_call(expr: ir.IRExpression) -> bool:
    """Return True if *expr* contains any function or method call."""
    if isinstance(expr, (ir.IRFunctionCall, ir.IRMethodCall)):
        return True
    if isinstance(expr, ir.IRIfElseExpression):
        return _contains_call(expr.condition) or _contains_call(expr.then_expr) or _contains_call(expr.else_expr)
    if hasattr(expr, "__dict__"):
        for val in expr.__dict__.values():
            if isinstance(val, ir.IRExpression) and _contains_call(val):
                return True
            if isinstance(val, list):
                for item in val:
                    if isinstance(item, ir.IRExpression) and _contains_call(item):
                        return True
    return False


def _step_is_constant_one(inc_expr: ir.IRBinaryOperation, var_name: str) -> bool:
    """Check that *inc_expr* is exactly `var_name + 1`."""
    return (
        inc_expr.operator == "Add"
        and isinstance(inc_expr.left, ir.IRVariable)
        and inc_expr.left.name == var_name
        and isinstance(inc_expr.right, ir.IRLiteral)
        and inc_expr.right.value == 1
    )


def _is_safe_to_canonicalize(
    var_name: str,
    start_val: object,
    end_expr: ir.IRExpression,
    body: list,
    increment_stmt: ir.IRStatement,
) -> bool:
    """Return True iff the while-loop pattern can be safely rewritten as a
    numeric `for` loop.

    Safety conditions (all must hold):
    1. The loop bound expression (end_expr) is pure — no calls, no mutations.
    2. The step expression in the increment statement is `var + 1` with a
       constant 1 (not a computed step).
    3. The induction variable is only mutated *once* in the body — i.e. only
       by the final increment statement and nowhere else in the body.
    4. No call inside the body reads or writes the induction variable in a way
       that would affect loop progress (guarded by condition 3 above).
    """
    # --- Condition 1: pure bound ----------------------------------------
    if not _is_pure_expr(end_expr):
        return False

    # --- Condition 2: constant step of exactly 1 ------------------------
    if not (
        isinstance(increment_stmt, ir.IRAssignment)
        and isinstance(increment_stmt.value, ir.IRBinaryOperation)
        and _step_is_constant_one(increment_stmt.value, var_name)
    ):
        return False

    # --- Condition 3: induction variable mutated only by the increment ---
    all_mutations: Set[int] = set()
    _collect_var_mutations(body, var_name, all_mutations)
    # The increment stmt itself is one mutation — any more is unsafe.
    if len(all_mutations) > 1:
        return False
    # If the sole mutation is NOT our increment stmt, that's also wrong.
    if len(all_mutations) == 1 and id(increment_stmt) not in all_mutations:
        return False

    # --- Condition 4: no call in body reads back the induction var in a
    #     non-trivial way (already covered by condition 3 for assignments,
    #     but guard against i being passed into a call that mutates it via
    #     globals — this is conservative: if end_expr contains a call we
    #     already rejected it; here we just log and pass).
    # (Condition 3 is the primary safeguard; this is defence in depth.)

    return True


# ---------------------------------------------------------------------------
# Canonicalizer
# ---------------------------------------------------------------------------

class IRCanonicalizer:
    """
    Transforms restructured IR into more idiomatic forms.
    Phase 4: WhileStatement → ForStatement conversion.

    The conversion is guarded by `_is_safe_to_canonicalize` which enforces:
      - Pure (side-effect-free) loop bound
      - Constant step of exactly 1
      - Induction variable mutated only by the final increment statement
    Loops that fail any check are left as `while` loops.
    """
    def __init__(self):
        self.changed = False

    def canonicalize(self, module: ir.IRModule):
        # Run multiple passes until no more changes or limit reached
        for _ in range(3):
            self.changed = False
            module.body = self._canonicalize_block(module.body)
            if not self.changed:
                break
        return module

    def visit(self, node: Any) -> Any:
        if node is None: return None
        method = f"visit_{type(node).__name__}"
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node: Any) -> Any:
        return node

    def visit_IRIfElseExpression(self, node: ir.IRIfElseExpression) -> ir.IRIfElseExpression:
        node.condition = self.visit(node.condition)
        node.then_expr = self.visit(node.then_expr)
        node.else_expr = self.visit(node.else_expr)
        return node

    def _canonicalize_block(self, block: list) -> list:
        # 1. Inline trivial variable copies (e.g. _t1 = x)
        block = self._inline_variable_copies(block)
        
        # 2. Pattern-based block transformations (e.g. while -> for)
        new_block = []
        i = 0
        while i < len(block):
            stmt = block[i]

            if (
                isinstance(stmt, ir.IRAssignment)
                and isinstance(stmt.target, ir.IRVariable)
                and isinstance(stmt.value, ir.IRLiteral)
                and i + 1 < len(block)
            ):
                next_stmt = block[i + 1]
                if isinstance(next_stmt, ir.IRWhileStatement):
                    for_stmt = self._try_convert_to_for(stmt, next_stmt)
                    if for_stmt:
                        new_block.append(for_stmt)
                        i += 2
                        self.changed = True
                        continue

            # Recurse into nested blocks
            if isinstance(stmt, ir.IRWhileStatement):
                stmt.body = self._canonicalize_block(stmt.body)
            elif isinstance(stmt, ir.IRIfStatement):
                stmt.then_block = self._canonicalize_block(stmt.then_block)
                if stmt.else_block:
                    stmt.else_block = self._canonicalize_block(stmt.else_block)
                for _cond, body in stmt.elif_blocks:
                    if isinstance(body, list):
                        canonicalized = self._canonicalize_block(list(body))
                        body.clear()
                        body.extend(canonicalized)
            elif isinstance(stmt, ir.IRFunctionDef):
                stmt.body = self._canonicalize_block(stmt.body)
            elif isinstance(stmt, ir.IRForStatement):
                stmt.body = self._canonicalize_block(stmt.body)

            new_block.append(stmt)
            i += 1
        
        # 3. Final cleanup: dead assignment removal for temporaries
        new_block = self._remove_dead_assignments(new_block)
        
        return new_block

    def _inline_variable_copies(self, block: list) -> list:
        """Inline 'tmp = var' assignments when tmp is a temporary."""
        inlines = {} # name -> IRExpression
        result = []
        
        for stmt in block:
            # Check for candidate: _tmp = original_var
            if (isinstance(stmt, ir.IRAssignment) and 
                isinstance(stmt.target, ir.IRVariable) and 
                stmt.target.name.startswith("_") and
                isinstance(stmt.value, ir.IRVariable)):
                
                # Only inline if we haven't seen this tmp before in this block
                if stmt.target.name not in inlines:
                    inlines[stmt.target.name] = stmt.value
                    self.changed = True
                    continue # Drop the assignment
            
            # Apply existing inlines to this statement
            if inlines:
                self._apply_inlines(stmt, inlines)
            
            result.append(stmt)
        return result

    def _apply_inlines(self, node: Any, inlines: dict):
        """Recursively replace variables in expressions with their inlined values."""
        if hasattr(node, "__dict__"):
            for attr, val in node.__dict__.items():
                if isinstance(val, ir.IRVariable):
                    var_name = getattr(val, "name", "")
                    if var_name in inlines:
                        node.__dict__[attr] = inlines[var_name]
                elif isinstance(val, ir.IRExpression):
                    self._apply_inlines(val, inlines)
                elif isinstance(val, list):
                    for item in val:
                        self._apply_inlines(item, inlines)
                elif isinstance(val, ir.IRStatement):
                    # For nested blocks (if/while), we only inline into the conditions
                    # but not the body to avoid incorrect scoping issues across blocks
                    cond = getattr(val, "condition", None)
                    if cond is not None:
                        self._apply_inlines(cond, inlines)
                elif isinstance(val, ir.IRIfElseExpression):
                    self._apply_inlines(val.condition, inlines)
                    self._apply_inlines(val.then_expr, inlines)
                    self._apply_inlines(val.else_expr, inlines)

    def _remove_dead_assignments(self, block: list) -> list:
        """Remove assignments to temporaries if they are never used in the block."""
        used_vars = set()
        
        def collect_uses(node):
            if isinstance(node, ir.IRVariable):
                used_vars.add(node.name)
            elif hasattr(node, "__dict__"):
                for val in node.__dict__.values():
                    if isinstance(val, (ir.IRExpression, ir.IRStatement)):
                        collect_uses(val)
                    elif isinstance(val, list):
                        for item in val:
                            collect_uses(item)

        # Collect all used variables in the block
        for stmt in block:
            # Don't count the target of an assignment as a 'use'
            if isinstance(stmt, ir.IRAssignment):
                if not isinstance(stmt.target, ir.IRVariable):
                    collect_uses(stmt.target)
                collect_uses(stmt.value)
            else:
                collect_uses(stmt)
                
        # Filter out assignments to unused temporaries
        new_block = []
        for stmt in block:
            if (isinstance(stmt, ir.IRAssignment) and 
                isinstance(stmt.target, ir.IRVariable)):
                if (stmt.target.name.startswith("_") and 
                    stmt.target.name not in used_vars):
                    self.changed = True
                    continue
            new_block.append(stmt)
        return new_block

    def _try_convert_to_for(
        self,
        init_stmt: ir.IRAssignment,
        while_stmt: ir.IRWhileStatement,
    ) -> "ir.IRForStatement | None":
        var_name = init_stmt.target.name
        start_val = init_stmt.value.value

        # --- Structural pattern checks (fast path before safety analysis) ---

        # Condition: while body must be non-empty
        if not while_stmt.body:
            return None

        # Condition: condition must be a binary comparison `i < N` or `i <= N`
        cond = while_stmt.condition
        if not isinstance(cond, ir.IRBinaryOperation):
            return None
        if not (isinstance(cond.left, ir.IRVariable) and cond.left.name == var_name):
            return None
        if cond.operator not in ("Lt", "LtE"):
            return None

        end_expr = cond.right

        # Condition: last stmt in body must be the increment `i = i + 1`
        last_stmt = while_stmt.body[-1]
        if not (
            isinstance(last_stmt, ir.IRAssignment)
            and isinstance(last_stmt.target, ir.IRVariable)
            and last_stmt.target.name == var_name
        ):
            return None

        # --- Safety gate ------------------------------------------------
        if not _is_safe_to_canonicalize(
            var_name=var_name,
            start_val=start_val,
            end_expr=end_expr,
            body=while_stmt.body,
            increment_stmt=last_stmt,
        ):
            return None  # Leave as while — unsafe to canonicalize

        # --- Build ForStatement ----------------------------------------
        new_body = while_stmt.body[:-1]  # strip the trailing increment

        # Luau `for i = start, end` is inclusive on both ends.
        # `while i < N`  ↔  for i = start, N - 1  (use N directly via step)
        # `while i <= N` ↔  for i = start, N
        #
        # Actually Luau numeric for: for i = s, e, step → i ∈ {s, s+step, …≤e}
        # step defaults to 1.
        # while i < N   → for i = s, N-1, 1   when N is a literal
        #                  for i = s, N,   1   when we keep operator semantics
        # We use the simpler approach: keep `end_expr` as-is for Lt since
        # Luau for is inclusive but we already ensured step=1 and the body
        # uses the value before increment, so adjust end by operator.
        if cond.operator == "Lt":
            # while i < N → for i = start, N (exclusive mapped to inclusive - 1)
            # To keep generated code clean: if N is literal just subtract.
            if isinstance(end_expr, ir.IRLiteral) and isinstance(end_expr.value, int):
                luau_end: ir.IRExpression = ir.IRLiteral(
                    value=end_expr.value - 1,
                    lineno=end_expr.lineno,
                    col_offset=end_expr.col_offset,
                )
            else:
                luau_end = ir.IRBinaryOperation(
                    left=end_expr,
                    operator="Sub",
                    right=ir.IRLiteral(value=1),
                    lineno=end_expr.lineno,
                    col_offset=end_expr.col_offset,
                )
        else:  # LtE — already inclusive
            luau_end = end_expr

        return ir.IRForStatement(
            var=var_name,
            start=ir.IRLiteral(
                value=start_val,
                lineno=init_stmt.lineno,
                col_offset=init_stmt.col_offset,
            ),
            end=luau_end,
            body=new_body,
            lineno=while_stmt.lineno,
            col_offset=while_stmt.col_offset,
        )
