"""
CAP, SSA Intro, Elimination and Optimisations
Functions to convert a CFG out of SSA Form.
"""

from typing import cast, List
from Lib import RiscV
from Lib.CFG import Block, BlockInstr, CFG
from Lib.Operands import Temporary
from Lib.Statement import AbsoluteJump, Label
from Lib.Terminator import BranchingTerminator, Return
from Lib.PhiNode import PhiNode
from TP05.SequentializeMoves import sequentialize_moves


def generate_moves_from_phis(phis: List[PhiNode], parent: Block) -> List[BlockInstr]:
    """
    `generate_moves_from_phis(phis, parent)` builds a list of move instructions
    to be inserted in a new block between `parent` and the block with phi nodes
    `phis`.

    This is an helper function called during SSA exit.
    """
    # (Lab 5a, Exercise 6)
    return [
        RiscV.mv(
            phi.defined()[0],
            phi.get_srcs()[parent.get_label()]
        ) for phi in phis if parent.get_label() in phi.get_srcs()
    ]


def exit_ssa(cfg: CFG, is_smart: bool) -> None:
    """
    `exit_ssa(cfg)` replaces phi nodes with move instructions to exit SSA form.

    `is_smart` is set to true when smart register allocation is enabled (Lab 5b).
    """
    for b in cfg.get_blocks():
        phis = cast(List[PhiNode], b.get_phis())  # Use cast for Pyright
        b.remove_all_phis()  # Remove all phi nodes in the block
        parents: List[Block] = b.get_in().copy()  # Copy as we modify it by adding blocks
        for parent in parents:
            # Add the block containing 'moves' to 'cfg'
            moves_label = cfg.fdata.fresh_label("merge_phi")
            block_moves = Block(moves_label, generate_moves_from_phis(phis, parent), AbsoluteJump(b.get_label()))
            cfg.add_block(block_moves)

            # Update parent terminator
            j = parent.get_terminator()
            match j:
                case AbsoluteJump():
                    new_terminator = AbsoluteJump(moves_label)
                case BranchingTerminator():
                    lbl_else : Label = moves_label if (j.label_else == b.get_label()) else j.label_else
                    lbl_then : Label = moves_label if (j.label_then == b.get_label()) else j.label_then
                    new_terminator = BranchingTerminator(j.cond, j.op1, j.op2, lbl_else, lbl_then)
                case Return():
                    new_terminator = AbsoluteJump(moves_label)
                    block_moves.set_terminator(Return())

            parent.set_terminator(new_terminator)