"""
CAP, SSA Intro, Elimination and Optimisations
Helper functions to convert a CFG out of SSA Form
for the Smart Allocator.
"""

from typing import List, Set, Tuple
from Lib import RiscV
from Lib.Graphes import DiGraph
from Lib.CFG import BlockInstr
from Lib.Operands import Register, Offset, DataLocation, S


def generate_smart_move(dest: DataLocation, src: DataLocation) -> List[BlockInstr]:
    """
    Generate a list of move, store and load instructions, depending on
    whether the operands are registers or memory locations.
    This is an helper function for `sequentialize_moves`.
    """
    instr: List[BlockInstr] = []
    tmp: Register = S[1]

    match dest, src:
        case Register(), Register():
            instr.append(RiscV.mv(dest, src))
        case Register(), Offset():
            instr.append(RiscV.ld(dest, src))
        case Offset(), Register():
            instr.append(RiscV.sd(src, dest))
        case Offset(), Offset():
            instr.append(RiscV.ld(tmp, src))
            instr.append(RiscV.sd(tmp, dest))
        case _, _:
            raise ValueError(f"Unsupported operands {type(dest)} <- {type(src)}")

    return instr


def swap_registers(r1: Register, r2: Register) -> List[BlockInstr]:
    """Swap two registers"""
    return [
        RiscV.xor(r1, r1, r2),
        RiscV.xor(r2, r1, r2),
        RiscV.xor(r1, r1, r2)
    ]


def sequentialize_moves(parallel_moves: Set[Tuple[DataLocation, DataLocation]]
                        ) -> List[BlockInstr]:
    """
    Take a set of parallel moves represented as (destination, source) pairs,
    and return a list of sequential moves which respect the cycles.
    Use the register `tmp` S2 for the cycles.
    Return a corresponding list of RiscV instructions.
    This is an helper function called during SSA exit.
    """
    tmp: Register = S[2]  # S2 is not a general purpose register
    # Build the graph of the moves
    move_graph: DiGraph = DiGraph()
    for dest, src in parallel_moves:
        move_graph.add_edge((src, dest))
    # List for the sequentialized moves to do
    # Convention: in moves we put (dest, src) for each move
    moves: List[Tuple[DataLocation, DataLocation] | BlockInstr] = []
    # First iteratively remove all the vertices without successors
    vars_without_successor = {src
                              for src, dests in move_graph.neighbourhoods()
                              if len(dests) == 0}
    while vars_without_successor:
        var = vars_without_successor.pop()
        for src in move_graph.pred(var):
            moves.append((var, src))
            dests = [dests for sr, dests in move_graph.neighbourhoods() if sr == src][0]
            if len(dests) == 1:
                vars_without_successor.add(src)

        move_graph.delete_vertex(var)
    # Then handle the cycles
    cycles: List = move_graph.connected_components()
    for cycle in cycles:
        if len(cycle) == 1:
            continue
        if len(cycle) == 2 and isinstance(cycle[0], Register) and isinstance(cycle[1], Register):
            moves += swap_registers(*cycle)
            continue
        previous = tmp
        for var in reversed(cycle):
            moves.append((previous, var))
            previous = var
        moves.append((previous, tmp))
    # Transform the moves to do in actual RiscV instructions
    moves_instr: List[BlockInstr] = []
    for val in moves:
        if isinstance(val, BlockInstr):
            moves_instr.append(val)
            continue
        instrs = generate_smart_move(*val)
        moves_instr.extend(instrs)
    return moves_instr