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

from typing import cast, List, Set, Tuple
from Lib.Errors import MiniCInternalError
from Lib import RiscV
from Lib.Graphes import DiGraph
from Lib.CFG import Block, BlockInstr, CFG
from Lib.Operands import (
    Register, DataLocation,
    Temporary, TemporaryPool)
from Lib.Statement import AbsoluteJump
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.
    """
    moves: List[BlockInstr] = []
    plabel = parent.get_label()
    for phi in phis:
        if(plabel in phi.used()):
            moves.append(RiscV.mv(phi.var,phi.used()[plabel]))
    return moves


def generate_moves_from_phis_smart(pool: TemporaryPool, phis: List[PhiNode], parent: Block) -> List[BlockInstr]:
    """
    `generate_moves_from_phis_smart(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 function is working with smart allocation, unlike its «dump» counterpart
    that makes allocations in no particular order.

    This is an helper function called during SSA exit.
    """
    moves: List[Tuple[DataLocation,DataLocation]] = []
    plabel = parent.get_label()
    for phi in phis:
        if(plabel in phi.used()):
            src = phi.var
            dest = phi.used()[plabel]
            if(isinstance(src,Temporary)):
                src = pool.get_alloced_loc(src)
            if(isinstance(dest,Temporary)):
                dest = pool.get_alloced_loc(dest)
            assert(isinstance(dest,DataLocation))
            moves.append((src,dest))
    realmoves = sequentialize_moves(set(moves))
    return realmoves


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._phis)  # Use cast for Pyright
        b._phis = []  # Remove all phi nodes in the block
        blabel = b.get_label()
        parents: List[Block] = b.get_in().copy()  # Copy as we modify it by adding blocks
        for p in parents:
            moves = generate_moves_from_phis_smart(cfg.fdata._pool,phis, p) if is_smart else generate_moves_from_phis(phis, p)
            if(len(moves)==0):
                continue
            # Creating the block
            ilabel = cfg.fdata.fresh_label(p.get_label().name+"_to_"+b.get_label().name)
            i = Block(ilabel,moves,AbsoluteJump(blabel))
            # Add the block
            cfg.add_block(i)
            # Changing the jumps
            ot = p.get_terminator()
            if(isinstance(ot,BranchingTerminator)):
                if(ot.label_then == blabel):
                    ot.label_then = ilabel
                if(ot.label_else == blabel):
                    ot.label_else = ilabel
            elif(isinstance(ot,AbsoluteJump)):
                assert(ot.label == blabel)
                ot = AbsoluteJump(ilabel)
            elif(isinstance(ot,Return)):
                raise MiniCInternalError("Malformed CFG, cannot have a return in a parent block")
            else:
                raise MiniCInternalError("I don't know of this terminator type:",type(ot))
            p.set_terminator(ot) # This instruction might be useless
            # Moving from p -> b to p -> i -> b
            cfg.remove_edge(p,b)
            cfg.add_edge(p,i)
            cfg.add_edge(i,b)