"""
transpiler/cfg_builder.py — Converts IR tree to Control Flow Graph.
"""

from __future__ import annotations
from typing import List, Optional, Dict
from transpiler.middleend import ir, cfg

class CFGBuilder:
    def __init__(self):
        self._block_count = 0
        self._blocks: Dict[int, cfg.BasicBlock] = {}
        self._current_block: Optional[cfg.BasicBlock] = None
        self._loop_stack: List[tuple[int, int]] = [] # (continue_block_id, break_block_id)

    def build(self, module: ir.IRModule) -> cfg.CFG:
        entry = self._create_block()
        exit_virtual = self._create_block() # Virtual sink for all exits
        
        self._current_block = entry
        self._exit_virtual = exit_virtual
        
        for stmt in module.body:
            self._process_statement(stmt)
            if self._current_block is None:
                self._current_block = self._create_block()
            
        # Any block still "open" or ending in return should point to exit_virtual
        for bid, block in list(self._blocks.items()):
            if bid == exit_virtual.id:
                continue
            if not block.terminator:
                if not block.successors:
                    # True leaf block: add default return and connect to sink
                    # (Note: simple return is enough for default)
                    ret = ir.IRReturnStatement(lineno=0, col_offset=0, value=None)
                    block.statements.append(ret)
                    block.add_successor(exit_virtual)
                elif len(block.successors) == 1:
                    # Exactly one successor: must be a fallthrough Jump
                    target = block.successors[0]
                    jump = ir.IRJump(target_block_id=target.id, lineno=0)
                    block.statements.append(jump)
                else:
                    # Multiple successors but no terminator. 
                    # This implies a Branch statement is missing.
                    # We'll skip adding a Jump to avoid "Jump block has N successors".
                    pass
        
        # Finally, ensure exit_virtual has no successors
        exit_virtual.successors = []
                
        print(f"DEBUG: CFG Built with {len(self._blocks)} blocks. Exit virtual: {exit_virtual.id}")
        for bid, b in self._blocks.items():
            print(f"DEBUG: Block {bid} successors: {[s.id for s in b.successors]}, predecessors: {[p.id for p in b.predecessors]}")
        return cfg.CFG(entry_block=entry, exit_blocks={exit_virtual}, blocks=self._blocks)

    def _create_block(self) -> cfg.BasicBlock:
        block = cfg.BasicBlock(id=self._block_count)
        self._blocks[self._block_count] = block
        self._block_count += 1
        return block

    def _add_statement(self, stmt: ir.IRStatement):
        curr = self._current_block
        if curr is not None:
            curr.statements.append(stmt)
            # Handle successors automatically for control flow nodes
            if isinstance(stmt, ir.IRJump):
                curr.add_successor(self._blocks[stmt.target_block_id])
            elif isinstance(stmt, ir.IRBranch):
                curr.add_successor(self._blocks[stmt.true_block_id])
                curr.add_successor(self._blocks[stmt.false_block_id])
            elif isinstance(stmt, (ir.IRReturnStatement, ir.IRMultiReturn)):
                curr.add_successor(self._exit_virtual)

            # If we just emitted a terminator, the block ends
            if isinstance(stmt, (ir.IRReturnStatement, ir.IRMultiReturn, ir.IRJump, ir.IRBranch)):
                self._current_block = None

    def _process_statement(self, stmt: ir.IRStatement):
        # If current block is None (e.g. after a return), we create a "dead" block 
        # for subsequent statements. DCE will remove it later if it stays unreachable.
        if self._current_block is None:
            self._current_block = self._create_block()

        if isinstance(stmt, ir.IRIfStatement):
            self._process_if(stmt)
        elif isinstance(stmt, ir.IRWhileStatement):
            self._process_while(stmt)
        elif isinstance(stmt, ir.IRForStatement):
            self._process_for(stmt)
        elif isinstance(stmt, ir.IRGenericForStatement):
            self._process_generic_for(stmt)
        elif isinstance(stmt, ir.IRBreakStatement):
            if self._loop_stack:
                _, break_id = self._loop_stack[-1]
                self._add_statement(ir.IRJump(target_block_id=break_id, lineno=stmt.lineno))
            self._current_block = None
        elif isinstance(stmt, ir.IRContinueStatement):
            if self._loop_stack:
                cont_id, _ = self._loop_stack[-1]
                self._add_statement(ir.IRJump(target_block_id=cont_id, lineno=stmt.lineno))
            self._current_block = None
        else:
            self._add_statement(stmt)

    def _process_if(self, node: ir.IRIfStatement):
        then_block = self._create_block()
        else_block = self._create_block()
        merge_block = self._create_block()
        
        # Branch to then/else
        branch = ir.IRBranch(
            condition=node.condition, 
            true_block_id=then_block.id, 
            false_block_id=else_block.id,
            lineno=node.lineno
        )
        
        curr = self._current_block
        if curr is not None:
            self._add_statement(branch)
            curr.add_successor(then_block)
            curr.add_successor(else_block)
            
        # Process then
        self._current_block = then_block
        for s in node.then_block:
            self._process_statement(s)
        target = self._current_block
        if target is not None:
            self._add_statement(ir.IRJump(target_block_id=merge_block.id, lineno=node.lineno))
            target.add_successor(merge_block)
            
        # Process else
        self._current_block = else_block
        if node.else_block:
            for s in node.else_block:
                self._process_statement(s)
        target = self._current_block
        if target is not None:
            self._add_statement(ir.IRJump(target_block_id=merge_block.id, lineno=node.lineno))
            target.add_successor(merge_block)
            
        self._current_block = merge_block

    def _process_while(self, stmt: ir.IRWhileStatement):
        cond_block = self._create_block()
        body_block = self._create_block()
        exit_block = self._create_block()
        
        prev_block = self._current_block
        if prev_block:
            self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
            prev_block.add_successor(cond_block)
            
        # Cond block
        self._current_block = cond_block
        branch = ir.IRBranch(
            condition=stmt.condition, 
            true_block_id=body_block.id, 
            false_block_id=exit_block.id,
            lineno=stmt.lineno
        )
        self._add_statement(branch)
        cond_block.add_successor(body_block)
        cond_block.add_successor(exit_block)
        
        # Body block
        self._loop_stack.append((cond_block.id, exit_block.id))
        self._current_block = body_block
        for s in stmt.body:
            self._process_statement(s)
        
        target = self._current_block
        if target:
            self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
            target.add_successor(cond_block)
        self._loop_stack.pop()
            
        self._current_block = exit_block

    def _process_for(self, stmt: ir.IRForStatement):
        """
        Lowers for i in range(start, end) to:
        i = start
        cond_block: if i < end goto body else goto exit
        body_block: ...; goto increment
        increment_block: i = i + 1; goto cond
        """
        # 1. Initialization: i = start
        self._process_statement(ir.IRAssignment(
            target=ir.IRVariable(name=stmt.var, lineno=stmt.lineno), 
            value=stmt.start,
            lineno=stmt.lineno
        ))
        
        cond_block = self._create_block()
        body_block = self._create_block()
        increment_block = self._create_block()
        exit_block = self._create_block()
        
        prev_block = self._current_block
        if prev_block:
            self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
            prev_block.add_successor(cond_block)
            
        # 2. Cond block: i < end
        self._current_block = cond_block
        cond_expr = ir.IRBinaryOperation(
            left=ir.IRVariable(name=stmt.var, lineno=stmt.lineno), 
            operator="Lt", 
            right=stmt.end, 
            lineno=stmt.lineno
        )
        branch = ir.IRBranch(
            condition=cond_expr, 
            true_block_id=body_block.id, 
            false_block_id=exit_block.id, 
            lineno=stmt.lineno
        )
        self._add_statement(branch)
        cond_block.add_successor(body_block)
        cond_block.add_successor(exit_block)
        
        # 3. Body block
        # continue jumps to increment_block, break jumps to exit_block
        self._loop_stack.append((increment_block.id, exit_block.id))
        self._current_block = body_block
        for s in stmt.body:
            self._process_statement(s)
            
        # 4. Increment block: i = i + 1
        curr = self._current_block
        if curr:
            self._add_statement(ir.IRJump(target_block_id=increment_block.id, lineno=stmt.lineno))
            curr.add_successor(increment_block)
            
        self._current_block = increment_block
        self._process_statement(ir.IRAssignment(
            target=ir.IRVariable(name=stmt.var, lineno=stmt.lineno),
            value=ir.IRBinaryOperation(
                left=ir.IRVariable(name=stmt.var, lineno=stmt.lineno), 
                operator="Add", 
                right=ir.IRLiteral(value=1, lineno=stmt.lineno),
                lineno=stmt.lineno
            ),
            lineno=stmt.lineno
        ))
        
        # Jump back to condition
        self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
        increment_block.add_successor(cond_block)
        
        self._loop_stack.pop()
        self._current_block = exit_block

    def _process_generic_for(self, stmt: ir.IRGenericForStatement):
        """
        Lowers for v1, v2 in iterator_expr to:
        cond_block: if iterator_active goto body else goto exit
        body_block: v1, v2 = ...; body_stmts; goto cond
        """
        # 1. Start with any setup (none required for generic for in Luau, 
        # but SSA needs the transition).
        cond_block = self._create_block()
        body_block = self._create_block()
        exit_block = self._create_block()
        
        prev_block = self._current_block
        if prev_block:
            self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
            prev_block.add_successor(cond_block)
            
        # 2. Cond block
        # Use IRGenericForTerminator to make the loop header explicit
        self._current_block = cond_block
        term = ir.IRGenericForTerminator(
            vars=stmt.vars,
            iterator=stmt.iterator,
            true_block_id=body_block.id,
            false_block_id=exit_block.id,
            lineno=stmt.lineno
        )
        self._add_statement(term)
        cond_block.add_successor(body_block)
        cond_block.add_successor(exit_block)
        
        # 3. Body block
        self._loop_stack.append((cond_block.id, exit_block.id))
        self._current_block = body_block
        
        # SSA needs to see the iteration variables being defined in the body
        self._add_statement(ir.IRMultiAssign(
            targets=[ir.IRVariable(name=v, lineno=stmt.lineno) for v in stmt.vars],
            rhs_values=[],
            rhs_call=stmt.iterator,
            is_local=True,
            is_loop_assignment=True,
            lineno=stmt.lineno
        ))
        
        for s in stmt.body:
            self._process_statement(s)
            
        # Jump back to condition
        curr = self._current_block
        if curr:
            self._add_statement(ir.IRJump(target_block_id=cond_block.id, lineno=stmt.lineno))
            curr.add_successor(cond_block)
            
        self._loop_stack.pop()
        self._current_block = exit_block

def build_cfg(module: ir.IRModule) -> cfg.CFG:
    builder = CFGBuilder()
    return builder.build(module)
