package ch.usi.inf.sp.cfg;

import ch.usi.inf.sp.bytecode.Disassembler;
import ch.usi.inf.sp.cfg.builder.JumpSource;
import ch.usi.inf.sp.cfg.builder.LookupSwitchInstructionNodeInfo;
import ch.usi.inf.sp.cfg.builder.MultiMap;
import ch.usi.inf.sp.cfg.builder.SwitchInstructionNodeInfo;
import ch.usi.inf.sp.cfg.builder.TableSwitchInstructionNodeInfo;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.util.*;

public final class ControlFlowGraphBuilder {
    private static final String DEFAULT_LABEL = "default";
    private static final String TRUE_LABEL = "T";

    private final MultiMap<LabelNode, JumpSource> labelToIncomingEdges = new MultiMap<>();
    private final List<LabelNode> labels = new ArrayList<>();
    private final Map<AbstractInsnNode, BasicBlock> insnToBlock = new HashMap<>();
    private final ControlFlowGraph graph = new ControlFlowGraph();
    private final MethodNode method;
    private BasicBlock currentBasicBlock;

    private ControlFlowGraphBuilder(MethodNode method) {
        this.method = method;
        setNewBasicBlock(0);
        graph.addEntryEdge(currentBasicBlock);
    }

    public static ControlFlowGraph createControlFlowGraph(final MethodNode method) {
        return new ControlFlowGraphBuilder(method).create();
    }

    private void setNewBasicBlock(int id) {
        currentBasicBlock = new BasicBlock(id);
        graph.addNode(currentBasicBlock);
    }

    public ControlFlowGraph create() {
        // get the list of all instructions in that method
        final InsnList instructions = method.instructions;
        for (int i = 0; i < instructions.size(); i++) {
            final AbstractInsnNode instruction = instructions.get(i);

            switch (instruction.getType()) {
                case AbstractInsnNode.JUMP_INSN:
                        labelToIncomingEdges.put(((JumpInsnNode) instruction).label, new JumpSource(instruction, TRUE_LABEL));
                        break;
                case AbstractInsnNode.LOOKUPSWITCH_INSN:
                        captureSwitchEdges(new LookupSwitchInstructionNodeInfo((LookupSwitchInsnNode) instruction));
                        break;
                case AbstractInsnNode.TABLESWITCH_INSN:
                        captureSwitchEdges(new TableSwitchInstructionNodeInfo((TableSwitchInsnNode) instruction));
            }
        }

        for (int i = 0; i < instructions.size(); i++) {
            final AbstractInsnNode instruction = instructions.get(i);

            currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions));

            if (isReturnInstruction(instruction)) {
                graph.addExitEdge(currentBasicBlock);
            }

            insnToBlock.put(instruction, currentBasicBlock);

            if (isInsnSignificantLabel(instruction)) {
                labels.add((LabelNode) instruction);
            }

            if (isEndOfBlock(instruction)) {
                final BasicBlock previousBasicBlock = currentBasicBlock;
                setNewBasicBlock(i + 1);

                // GOTO and SWITCH instructions do not have a fallthrough edge, otherwise add one
                if (instruction.getOpcode() != Opcodes.GOTO
                        && instruction.getType() != AbstractInsnNode.LOOKUPSWITCH_INSN
                        && instruction.getType() != AbstractInsnNode.TABLESWITCH_INSN) {
                    graph.addFallthroughEdge(previousBasicBlock, currentBasicBlock);
                }
            }
        }

        for (final LabelNode label : labels) {
            final BasicBlock toBB = Objects.requireNonNull(insnToBlock.get(label));
            for (final JumpSource jumpSource : labelToIncomingEdges.getAll(label)) {
                final BasicBlock fromBB = jumpSource.block(insnToBlock);
                final ControlFlowEdge edge = jumpSource.edge();

                graph.addEdge(edge);
                graph.connect(fromBB, edge, toBB);
            }
        }

        return graph;
    }

    private void captureSwitchEdges(final SwitchInstructionNodeInfo info) {
        for (int j = 0; j < info.getCaseCount(); j++) {
            final LabelNode label = info.getLabelForCase(j);
            final String value = info.getKeyForCase(j);
            labelToIncomingEdges.put(label, new JumpSource(info.getNode(), value));
        }
        if (info.getDefaultCase() != null) {
            labelToIncomingEdges.put(info.getDefaultCase(), new JumpSource(info.getNode(), DEFAULT_LABEL));
        }
    }

    private boolean isInsnSignificantLabel(final AbstractInsnNode node) {
        return node instanceof LabelNode && labelToIncomingEdges.containsKey((LabelNode) node);
    }

    private boolean isReturnInstruction(final AbstractInsnNode instruction) {
        final int opcode = instruction.getOpcode();
        return opcode == Opcodes.RETURN ||
                opcode == Opcodes.ARETURN ||
                opcode == Opcodes.LRETURN ||
                opcode == Opcodes.IRETURN ||
                opcode == Opcodes.FRETURN;
    }

    private boolean isEndOfBlock(final AbstractInsnNode instruction) {
        final AbstractInsnNode nextInsn = instruction.getNext();

        if (nextInsn == null) {
            return false; // cannot start another bb at the end of the method with 0 instructions
        }

        final int type = instruction.getType();
        if (type == AbstractInsnNode.JUMP_INSN ||
                type == AbstractInsnNode.LOOKUPSWITCH_INSN ||
                type == AbstractInsnNode.TABLESWITCH_INSN) {
            return true; // if we're branching or jumping after this instruction, we NEED to cut the current bb short
        }

        // if the next instruction is a label some other bb may jump into then cut, otherwise continue
        return isInsnSignificantLabel(nextInsn);
    }
}