"""
CAP, CodeGeneration, CFG linearization to a list of statements
"""

from typing import List, Set

from Lib.Statement import (
    Statement, AbsoluteJump, ConditionalJump
)
from Lib.Terminator import (Return, BranchingTerminator)
from Lib.CFG import Block


def ordered_blocks_list(cfg) -> List[Block]:
    """
    Compute a list of blocks with optimized ordering for linearization.
    """
    # TODO (Lab4b, Extension)
    return cfg.get_blocks()


def linearize(cfg) -> List[Statement]:
    """
    Linearize the given control flow graph as a list of instructions.
    """
    l: List[Statement] = []  # Linearized CFG
    blocks: List[Block] = ordered_blocks_list(cfg)
    
    labdict = {}
    # We make the label dictionary
    for j,block in enumerate(blocks):
        labdict[block.get_label()] = block
   
    while blocks != []:
        block = blocks[0]
        if(l != [] and isinstance(l[-1],AbsoluteJump)):
            ll : AbsoluteJump = l[-1]
            # We try to find the next jump according to that.
            try:
                # If we find the instruction
                b = labdict[ll.label]
                if b in blocks:
                    # We remove the jump instruction and we select the block to use
                    l.pop()
                    block = b
            except KeyError:
                pass
                
        # We remove the block we used
        blocks.remove(block)
            
        # 1. Add the label of the block to the linearization
        l.append(block.get_label())
        # 2. Add the body of the block to the linearization
        l.extend(block.get_body())
        # 3. Add the terminator of the block to the linearization
        match block.get_terminator():
            case BranchingTerminator() as j:
                l.append(ConditionalJump(j.cond, j.op1, j.op2, j.label_then))
                l.append(AbsoluteJump(j.label_else))
            case AbsoluteJump() as j:
                l.append(AbsoluteJump(j.label))
            case Return():
                l.append(AbsoluteJump(cfg.get_end()))
    return l