from typing import Optional, Any, List
from transpiler.middleend import ir, cfg, cfg_builder, ssa, purity
from transpiler.optimization import loop_detector, canonicalizer
from transpiler.backend import restructurer
from transpiler.infrastructure.flags import CompilerFlags
from transpiler.infrastructure import logger

def _optimize_block_list(stmts: List[ir.IRStatement], initial_vars: Optional[List[str]] = None, options: Optional[CompilerFlags] = None, scope_name: str = "unknown") -> List[ir.IRStatement]:
    """Helper to optimize a list of statements by building a temporary module/CFG."""
    tmp_module = ir.IRModule(body=stmts)
    
    # 1. Build CFG
    target_cfg = cfg_builder.build_cfg(tmp_module)
        
    # Validation (if enabled)
    if options and getattr(options, "verify_cfg", False):
        from transpiler.optimization.cfg_validator import CFGValidator
        validator = CFGValidator()
        validator.validate_cfg(target_cfg)
        
    # Optimization level 0: Only structural transformations required for Luau output
    opt_level = getattr(options, "opt", 2) if options else 2
    
    if opt_level == 0:
        # We still need DominatorAnalysis for Restructurer
        dom = ssa.DominatorAnalysis(target_cfg)
        dom.run()
        new_module = restructurer.restructure_cfg(target_cfg)
        canonicalizer.IRCanonicalizer().canonicalize(new_module)
        return new_module.body

    # 2. SSA Construction
    dom_analysis = ssa.DominatorAnalysis(target_cfg)
    dom_analysis.run()
    
    ssa_constructor = ssa.SSAConstructor(target_cfg, dom_analysis, initial_vars=initial_vars)
    ssa_constructor.run()
    
    # 3. Purity Analysis
    purity_analyzer = purity.PurityAnalyzer()
    for block in target_cfg.blocks.values():
        for stmt in block.statements:
            purity_analyzer.analyze_stmt(stmt)
            
    # 4. Global Optimization Fixpoint Loop
    pass_cp = ssa.SSAConstantPropagation(target_cfg)
    pass_dce = ssa.SSAMarkAndSweepDCE(target_cfg)
    pass_copy = ssa.SSACopyPropagation(target_cfg)
    pass_rte = ssa.SSARedundantTempElimination(target_cfg)
    
    # Level 1+: Basic Local Optimizations
    for i in range(5):
        changed = False
        # if pass_cp.run(): changed = True # Temporarily disabled to debug folding issues
        if pass_dce.run(initial_vars=initial_vars): changed = True
        if pass_copy.run(): changed = True
        if pass_rte.run(): changed = True
        if not changed: break
    # Level 2: Advanced Global Optimizations
    if opt_level >= 2:
        pass_cse = ssa.SSACSE(target_cfg)
        pass_cse.run()
        
        detector = loop_detector.LoopDetector(target_cfg, dom_analysis)
        detector.run()
        pass_licm = ssa.SSALICM(target_cfg, detector.loops)
        pass_licm.run()
        
        # Expression Inlining
        ssa.SSAExpressionInliner(target_cfg).run()

    # 7. SSA Destruction
    ssa.DeSSA(target_cfg).run()

    # 8. Restructuring
    dom = ssa.DominatorAnalysis(target_cfg)
    dom.run()
    new_module = restructurer.restructure_cfg(target_cfg)
        
    # 9. Canonicalization
    canonicalizer.IRCanonicalizer().canonicalize(new_module)
    
    return new_module.body

def optimize_ir(module: ir.IRModule, options: Optional[CompilerFlags] = None):
    """
    Orchestrates the RPy optimization pipeline recursively for modules and functions.
    """
    if options and options.explain:
        logger.set_explain_mode(True)
        
    # First, recursively optimize any nested functions
    def _optimize_functions(container: Any, parent_vars: List[str]):
        stmts = None
        if hasattr(container, "body") and isinstance(container.body, list): stmts = container.body
        elif isinstance(container, list): stmts = container
        
        if stmts:
            # Collect current scope variables for nested functions
            current_vars = parent_vars.copy()
            for stmt in stmts:
                if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.target, ir.IRVariable):
                    current_vars.append(stmt.target.name)

            for stmt in stmts:
                if isinstance(stmt, ir.IRFunctionDef):
                    # Combine parent vars (upvalues) and function args
                    all_initial = list(set(current_vars + stmt.args))
                    # Optimize the function body BEFORE we optimize the outer scope
                    stmt.body = _optimize_block_list(stmt.body, initial_vars=all_initial, options=options, scope_name=stmt.name)
                    # Recursively check for nested functions inside this one
                    _optimize_functions(stmt.body, all_initial)

    _optimize_functions(module.body, [])

    # Finally, optimize the top-level module code
    module.body = _optimize_block_list(module.body, options=options, scope_name="module")

    # Finally, optimize the top-level module code
    module.body = _optimize_block_list(module.body, options=options, scope_name="module")
