from typing import Optional, Any, List, Set
from transpiler.backend import luau_ast as last

class PeepholeOptimizer:
    """
    Final pass on the Luau AST to apply idiom transformations (guard clauses, void strips, boolean simplifications).
    """
    def __init__(self):
        self._search_vars: Set[str] = set()

    def optimize(self, node: last.Node) -> last.Node:
        return self.visit(node)

    def visit(self, node: Any) -> Any:
        if node is None:
            return None
        if isinstance(node, list):
            res = []
            for x in node:
                v = self.visit(x)
                if v is not None:
                    res.append(v)
            return res
        
        if not isinstance(node, last.Node):
            return node

        method_name = f"visit_{type(node).__name__}"
        visitor = getattr(self, method_name, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node: last.Node) -> last.Node:
        for field_name, field_value in node.__dict__.items():
            if isinstance(field_value, last.Node):
                setattr(node, field_name, self.visit(field_value))
            elif isinstance(field_value, list):
                if field_name in ("body", "then_block", "else_block"):
                    # Use the block optimizer for statement lists
                    is_top_func = (field_name == "body" and isinstance(node, (last.FunctionDef, last.Lambda)))
                    is_top_loop = (field_name == "body" and isinstance(node, (last.WhileStatement, last.ForStatement, last.GenericForStatement)))
                    setattr(node, field_name, self._optimize_statements(field_value, is_function_top=is_top_func, is_loop_top=is_top_loop))
                else:
                    new_list = []
                    for item in field_value:
                        res = self.visit(item)
                        if res is not None:
                            new_list.append(res)
                    setattr(node, field_name, new_list)
        return node

    def _optimize_statements(self, statements: List[last.Statement], is_function_top: bool = False, is_loop_top: bool = False) -> List[last.Statement]:
        if not statements:
            return []
            
        new_body = []
        i = 0
        while i < len(statements):
            stmt = statements[i]
            opt_stmt = self.visit(stmt)
            if opt_stmt is None:
                i += 1
                continue

            # 1. Guard Clause Inversion (Loops)
            if isinstance(opt_stmt, last.IfStatement) and not opt_stmt.else_block and not opt_stmt.elif_blocks:
                is_loop_guard = False
                if len(opt_stmt.then_block) == 1 and isinstance(opt_stmt.then_block[0], (last.ContinueStatement, last.BreakStatement)):
                    is_loop_guard = True
                
                if is_loop_guard:
                    remaining = statements[i+1:]
                    if remaining:
                        inverted_cond = self._simplify_boolean(self._invert_condition(opt_stmt.condition))
                        wrapped_body = self._optimize_statements(remaining)
                        new_if = last.IfStatement(
                            condition=inverted_cond,
                            then_block=wrapped_body,
                            else_block=[],
                            lineno=opt_stmt.lineno
                        )
                        new_body.append(new_if)
                        return new_body
                    else:
                        if isinstance(opt_stmt.then_block[0], last.ContinueStatement):
                            i += 1
                            continue

            # 2. Guard Clause Hoisting (Functions)
            if is_function_top and isinstance(opt_stmt, last.IfStatement) and \
               not opt_stmt.else_block and not opt_stmt.elif_blocks:
                remaining_after_if = statements[i+1:]
                if not remaining_after_if and len(opt_stmt.then_block) >= 1:
                    inverted_cond = self._simplify_boolean(self._invert_condition(opt_stmt.condition))
                    new_body.append(last.IfStatement(
                        condition=inverted_cond,
                        then_block=[last.ReturnStatement(value=None)],
                        else_block=[],
                        lineno=opt_stmt.lineno
                    ))
                    new_body.extend(self._optimize_statements(opt_stmt.then_block, is_function_top=True))
                    return new_body

            # 3. Return/Else Lifting
            if isinstance(opt_stmt, last.IfStatement) and opt_stmt.else_block:
                if self._ends_with_return(opt_stmt.then_block):
                    lifted_else = opt_stmt.else_block
                    opt_stmt.else_block = []
                    new_body.append(opt_stmt)
                    new_body.extend(self._optimize_statements(lifted_else))
                    i += 1
                    continue
                
                if isinstance(opt_stmt.condition, last.UnaryOperation) and opt_stmt.condition.operator.strip() == "not":
                    inverted_cond = self._simplify_boolean(opt_stmt.condition.operand)
                    then_b = opt_stmt.then_block
                    else_b = opt_stmt.else_block
                    opt_stmt.condition = inverted_cond
                    opt_stmt.then_block = else_b
                    opt_stmt.else_block = then_b
            
            # 4. Empty if-then removal
            if isinstance(opt_stmt, last.IfStatement) and not opt_stmt.then_block and opt_stmt.else_block:
                 inverted_cond = self._simplify_boolean(self._invert_condition(opt_stmt.condition))
                 new_body.append(last.IfStatement(condition=inverted_cond, then_block=opt_stmt.else_block, else_block=[], lineno=opt_stmt.lineno))
                 i += 1
                 continue

            new_body.append(opt_stmt)
            i += 1
            
        # Elite-Tier Merge: Collapse adjacent guard clauses
        if not new_body: return []
        merged_body = []
        for stmt in new_body:
            # Robust merge check
            name = stmt.__class__.__name__
            is_prev_if = merged_body and merged_body[-1].__class__.__name__ == "IfStatement"
            
            if is_prev_if:
                prev = merged_body[-1]
                p_elif = bool(getattr(prev, "elif_blocks", []))
                p_else = bool(getattr(prev, "else_block", None))
                p_simple = self._is_simple_return(prev.then_block)
                
                if not p_else and not p_elif and p_simple:
                    if name == "IfStatement":
                        s_elif = bool(getattr(stmt, "elif_blocks", []))
                        s_else = bool(getattr(stmt, "else_block", None))
                        s_simple = self._is_simple_return(stmt.then_block)
                        
                        if not s_else and not s_elif and s_simple:
                           prev.condition = self._simplify_boolean(last.BinaryOperation(
                               left=prev.condition,
                               operator="or",
                               right=stmt.condition
                           ))
                           continue
            merged_body.append(stmt)
            
        # 2. Prune trailing continue in loops
        if is_loop_top and merged_body and merged_body[-1].__class__.__name__ == "ContinueStatement":
            merged_body.pop()
            
        return merged_body

    def _simplify_boolean(self, node: last.Expression) -> last.Expression:
        # Nil-Index Search Pattern: x ~= -1 -> x
        if isinstance(node, last.BinaryOperation):
            if node.operator == "~=" and isinstance(node.right, last.Literal) and node.right.value == -1:
                if isinstance(node.left, last.Variable) and node.left.name in self._search_vars:
                    return node.left
            if node.operator == "==" and isinstance(node.right, last.Literal) and node.right.value == -1:
                if isinstance(node.left, last.Variable) and node.left.name in self._search_vars:
                    return last.UnaryOperation(operator="not ", operand=node.left)

        if isinstance(node, last.BinaryOperation):
             # Elite Polish: Flatten or-nesting: A or (B or C) -> (A or B) or C
             # This ensures the printer emits A or B or C without parentheses.
             if node.operator == "or":
                 if isinstance(node.right, last.BinaryOperation) and node.right.operator == "or":
                     return self._simplify_boolean(last.BinaryOperation(
                         left=last.BinaryOperation(left=node.left, operator="or", right=node.right.left),
                         operator="or",
                         right=node.right.right
                     ))

             node.left = self._simplify_boolean(node.left)
             node.right = self._simplify_boolean(node.right)
             return node
             
        if isinstance(node, last.UnaryOperation) and node.operator.strip() == "not":
            operand = self._simplify_boolean(node.operand)
            if isinstance(operand, last.UnaryOperation) and operand.operator.strip() == "not":
                return operand.operand
            if isinstance(operand, last.BinaryOperation):
                if operand.operator == "or":
                    return last.BinaryOperation(
                        left=self._simplify_boolean(self._invert_condition(operand.left)),
                        operator="and",
                        right=self._simplify_boolean(self._invert_condition(operand.right))
                    )
                if operand.operator == "and":
                    inv_l = self._invert_condition(operand.left)
                    inv_r = self._invert_condition(operand.right)
                    if not (isinstance(inv_l, last.UnaryOperation) and inv_l.operator.strip() == "not") and \
                       not (isinstance(inv_r, last.UnaryOperation) and inv_r.operator.strip() == "not"):
                         return last.BinaryOperation(left=inv_l, operator="or", right=inv_r)
            return last.UnaryOperation(operator="not ", operand=operand)
            
        return node

    def _invert_condition(self, cond: last.Expression) -> last.Expression:
        # Nil-Index support for inversion: not (x) -> x == -1 -> x is nil
        if isinstance(cond, last.Variable) and cond.name in self._search_vars:
             return last.UnaryOperation(operator="not ", operand=cond)
             
        if isinstance(cond, last.UnaryOperation) and cond.operator.strip() == "not":
            return cond.operand
        if isinstance(cond, last.BinaryOperation):
            inv_map = {"==": "~=", "~=": "==", "<": ">=", ">": "<=", "<=": ">", ">=": "<"}
            if cond.operator in inv_map:
                return last.BinaryOperation(left=cond.left, operator=inv_map[cond.operator], right=cond.right)
            if cond.operator == "and":
                return last.BinaryOperation(left=self._invert_condition(cond.left), operator="or", right=self._invert_condition(cond.right))
            if cond.operator == "or":
                return last.BinaryOperation(left=self._invert_condition(cond.left), operator="and", right=self._invert_condition(cond.right))

        return last.UnaryOperation(operator="not ", operand=cond)

    def _ends_with_return(self, block: List[last.Statement]) -> bool:
        if not block: return False
        last_stmt = block[-1]
        return isinstance(last_stmt, (last.ReturnStatement, last.MultiReturnStatement, last.BreakStatement, last.ContinueStatement))

    def visit_Lambda(self, node: last.Lambda) -> last.Lambda:
        node.body = self._optimize_statements(node.body, is_function_top=True)
        return node

    def visit_FunctionDef(self, node: last.FunctionDef) -> last.FunctionDef:
        saved_search = self._search_vars.copy()
        node.body = self._optimize_statements(node.body, is_function_top=True)
        self._search_vars = saved_search
        return node
    
    def visit_Block(self, node: last.Block) -> last.Block:
        node.body = self._optimize_statements(node.body, is_loop_top=getattr(node, "is_loop_body", False))
        return node

    def visit_Assignment(self, node: last.Assignment) -> Optional[last.Node]:
        node.target = self.visit(node.target)
        node.value = self.visit(node.value)
        
        # 1. Nil-Index Pattern Detection: x = -1 -> local x = nil
        if isinstance(node.target, last.Variable) and isinstance(node.value, last.Literal) and node.value.value == -1:
             self._search_vars.add(node.target.name)
             return last.LocalAssign(name=node.target.name, value=None)

        # 2. Identity Pruning
        if isinstance(node.target, last.Variable) and isinstance(node.value, (last.IfElseExpression, last.IfExpression)):
            target_name = node.target.name
            cond = node.value.condition
            then_v = node.value.then_expr
            else_v = node.value.else_expr
            
            is_redundant_then = False
            if isinstance(then_v, last.Variable) and then_v.name == target_name:
                is_redundant_then = True
            elif isinstance(then_v, last.Literal) and then_v.value is None:
                if isinstance(cond, last.BinaryOperation) and cond.operator == "==":
                    if isinstance(cond.left, last.Variable) and cond.left.name == target_name and \
                       isinstance(cond.right, last.Literal) and cond.right.value is None:
                        is_redundant_then = True
            is_redundant_else = False
            if isinstance(else_v, last.Variable) and else_v.name == target_name:
                is_redundant_else = True
            if is_redundant_then and is_redundant_else:
                return None
        
        # 3. Compound Assignment
        if isinstance(node.value, last.BinaryOperation):
            if isinstance(node.target, last.Variable) and isinstance(node.value.left, last.Variable):
                if node.target.name == node.value.left.name:
                    op = node.value.operator
                    if op in ("+", "-", "*", "/", "%", "^", ".."):
                        return last.CompoundAssignment(target=node.target, operator=op, value=node.value.right, lineno=node.lineno)
        return node

    def visit_LocalAssign(self, node: last.LocalAssign) -> last.LocalAssign:
        node.value = self.visit(node.value)
        if isinstance(node.value, last.Literal) and node.value.value == -1:
            self._search_vars.add(node.name)
            node.value = None
        return node

    def _is_simple_return(self, block: List[last.Statement]) -> bool:
        if not block or len(block) != 1: return False
        stmt = block[0]
        name = stmt.__class__.__name__
        if name == "ReturnStatement":
            return getattr(stmt, "value", None) is None
        if name == "MultiReturnStatement":
            return not getattr(stmt, "values", [])
        return False

    def visit_IfStatement(self, node: last.IfStatement) -> last.IfStatement:
        # Collapse elif-guards BEFORE optimizing children
        if getattr(node, "elif_blocks", []) and self._is_simple_return(node.then_block):
            collapsible_count = 0
            for cond, body in node.elif_blocks:
                if self._is_simple_return(body):
                    collapsible_count += 1
                else:
                    break
            
            if collapsible_count > 0:
                # print(f"DEBUG_PEEPHOLE: Collapsing {collapsible_count} elif-guards: {node.condition}")
                new_cond = node.condition
                for i in range(collapsible_count):
                    new_cond = last.BinaryOperation(left=new_cond, operator="or", right=node.elif_blocks[i][0])
                node.condition = self._simplify_boolean(new_cond)
                node.elif_blocks = node.elif_blocks[collapsible_count:]

        node.condition = self._simplify_boolean(node.condition)
        node.then_block = self._optimize_statements(node.then_block)
        
        new_elif = []
        for cond, body in node.elif_blocks:
            new_elif.append((self._simplify_boolean(cond), self._optimize_statements(body)))
        node.elif_blocks = new_elif

        if node.else_block:
            node.else_block = self._optimize_statements(node.else_block)
                    
        return node

def optimize(node: last.Node) -> last.Node:
    return PeepholeOptimizer().optimize(node)
