from __future__ import annotations
import logging
from collections import defaultdict

import networkx as nx

from angr.ailment.block import Block
from angr.ailment.statement import ConditionalJump

from .errors import SAILRSemanticError
from .similarity import ail_similarity_to_orig_blocks
from .utils import (
    copy_graph_and_nodes,
    replace_node_in_graph,
    ail_block_from_stmts,
    correct_jump_targets,
    deepcopy_ail_anyjump,
)

_l = logging.getLogger(name=__name__)


class AILBlockSplit:
    """
    This class represents a block that has been split into three parts, which is best explained in the
    AILMergeGraph class example. See that class for more information.

    The up_split, is all statements above the matched Longest Common Sequence (LCS), the match_split is the LCS,
    and the down_split is all statements below the matched LCS. This class should only be used in the context of
    the AILMergeGraph class.
    """

    def __init__(self, original=None, up_split=None, match_split=None, down_split=None):
        """

        :param original:
        :param up_split:   The block split above the matched LCS
        :param match_split:   The block containing only the matched LCS
        :param down_split:  The block split below the matched LCS
        """
        self.original = original
        self.up_split = up_split
        self.match_split = match_split
        self.down_split = down_split

    @classmethod
    def from_block_lcs(cls, original_block: Block, idx, len_):
        pre, mid, post = cls.split_ail_block(original_block, idx, len_)
        return cls(
            original=original_block,
            up_split=pre,
            match_split=mid,
            down_split=post,
        )

    @staticmethod
    def split_ail_block(block, split_idx, split_len) -> tuple[Block | None, Block | None, Block | None]:
        if split_len == len(block.statements):
            return None, block, None

        up_split = ail_block_from_stmts(block.statements[:split_idx], block_addr=block.addr)
        match_split = ail_block_from_stmts(
            block.statements[split_idx : split_idx + split_len], block_addr=None if up_split else block.addr
        )
        down_split = ail_block_from_stmts(block.statements[split_idx + split_len :])

        return up_split, match_split, down_split

    def __str__(self):
        return (
            f"<AILBlockSplit: OG: {self.original.__repr__()} | Up: {self.up_split.__repr__()} | "
            f"Match: {self.match_split.__repr__()} | Down: {self.down_split.__repr__()}>"
        )

    def __repr__(self):
        return self.__str__()


class AILMergeGraph:
    """
    This class represents the merged results of two AIL graphs that have been found to be similar. We can reference
    these graphs as G1 and G2. The two graphs are of the following form, where any node other than D can be empty:

        A
       / \
      B   C
      \\  /
       D
      / \
     E   F

    The D node can be a subgraph, but in both G1 and G2, this D-subgraph are exact duplicates of each
    other, except their top and bottom statements. This class is the result of merging those two D subgraphs.

    To explain that last part about statements differing at the ends, see this example:

    D1:
    -----
    a = 10;
    puts(a);
    puts("bye");
    -----

    D2:
    -----
    a = 11;
    puts(a);
    puts("cya");
    -----

    In this case, the merged D would contain just `puts(a)`. The statements above it, referred to as up_split in
    the code, and the statements below it, referred to as down_split in the code, would be moved out of the block
    and bounded by the conditions that lead to those statements. This creates a graph even in the case of a single
    block being the original D.

    Lastly, since this class will deal a lot with splitting blocks into pieces, we keep a mapping of how the original
    blocks turned into the new ones and vice versa.
    """

    def __init__(
        self, graph=None, original_graph=None, conditional_block=None, original_blocks=None, original_split_blocks=None
    ):
        self.graph = graph or nx.DiGraph()
        self.original_graph = original_graph or nx.DiGraph()
        self.conditional_block = conditional_block

        # All the original blocks in both graphs that are going to be merged into a single
        # graph that is self.graph
        self.original_blocks: dict[Block, list[Block]] = original_blocks or {}

        # The mapping from an original_block in original_blocks to an AILBlockSplit,
        # which can have an up, match, and down split. This happens when only partial
        # blocks match from the original blocks
        self.original_split_blocks: dict[Block, list[AILBlockSplit]] = original_split_blocks or {}

        # The mapping from each block in the self.graph to the original blocks that
        # are contained in the larger CFG
        self.merge_blocks_to_originals: dict[Block, set[Block | AILBlockSplit]] = defaultdict(set)
        self.merge_end_to_start: dict[Block, Block] = {}
        self.starts = []
        self.original_ends = []

    def create_conditionless_graph(self, starting_blocks: list[Block], graph_lcs) -> dict[Block, Block] | None:
        # get all the original blocks (reverted from the LCS) and their split blocks.
        # split-blocks are blocks that need to be split at some stmt index to make the two blocks
        # equal across both graphs. At a highlevel, the first block in both matching graphs either need
        # to be a full match or a spilt block with a None up-split (since the up-split represent a starting
        # stmt that mismatches).
        self.starts = starting_blocks.copy()
        merge_base, other_base = self.starts[:]
        for block in self.starts:
            og_blks_bfs_order, split_blks = ail_similarity_to_orig_blocks(block, graph_lcs, self.original_graph)
            self.original_blocks[block] = og_blks_bfs_order
            self.original_split_blocks[block] = [
                AILBlockSplit.from_block_lcs(og, idx, len(lcs)) for og, (lcs, idx) in split_blks.items()
            ]

        # eliminate shared blocks that are the same in original blocks
        shared_blocks = set(self.original_blocks[merge_base]).intersection(set(self.original_blocks[other_base]))
        for shared_block in shared_blocks:
            for block in self.starts:
                if shared_block not in self.original_blocks[block]:
                    continue

                self.original_blocks[block].remove(shared_block)
                if shared_block in self.original_split_blocks[block]:
                    self.original_split_blocks[block].remove(shared_block)

        # we now know all the original blocks that will be merged, in order of BFS
        # so let's create the graph that will be the final output merged graph
        #
        # we start by finding all the blocks that are about to be split, i.e., blocks that
        # have at least one stmt mismatching from each other
        base_to_split = {
            split_block.original: split_block.match_split
            for split_block in self.original_split_blocks[merge_base]
            if split_block.match_split is not None
        }

        # we create a new graph, full of the original blocks of the base, with blocks
        # that should be split replaced.
        # this graph is only the initial merge_graph needed, where only the blocks
        subgraph = nx.subgraph(self.original_graph, self.original_blocks[merge_base])
        # ensure all base blocks are within the subgraph
        for block in base_to_split:
            if block not in subgraph:
                return None
        self.graph, update_blocks = self.clone_graph_replace_splits(subgraph, base_to_split)
        self._update_all_split_refs(update_blocks)
        for update_block, new_block in update_blocks.items():
            if update_block in starting_blocks:
                idx = self.starts.index(update_block)
                self.starts[idx] = new_block
        merge_base, other_base = self.starts[:]

        # Add all the blocks that start the graph that may have been split
        # (i.e., only blocks in the blocks list for this function)
        for block in self.starts:
            for split_block in self.original_split_blocks[block]:
                if split_block.up_split:
                    merge_base_split = self._find_split_block_by_original(merge_base) or merge_base
                    merge_base_split = (
                        merge_base_split.match_split
                        if isinstance(merge_base_split, AILBlockSplit)
                        else merge_base_split
                    )

                    self.graph.add_edge(split_block.up_split, merge_base_split)

        return update_blocks

    def add_edges_to_condition(self, conditional_block, true_target, merge_end_pairs):
        for match_node, merge_end_pair in merge_end_pairs.items():
            cond_copy = conditional_block.copy()
            cond_copy.statements = [deepcopy_ail_anyjump(cond_copy.statements[0])]
            cond_copy.idx += 1

            # fix the condition edges
            cond_jump_stmt: ConditionalJump = cond_copy.statements[-1]
            if not merge_end_pair:
                raise SAILRSemanticError("Encountered a conditional jump that has no successors! This can be bad!")
            if len(merge_end_pair) == 1:
                b0 = merge_end_pair[0]
                b0_og = next(iter(self.merge_blocks_to_originals[b0]))
                if isinstance(b0_og, AILBlockSplit):
                    b0_og = b0_og.original

                b1_og = self._find_block_pair_in_originals(b0_og)
                if b1_og is None:
                    raise SAILRSemanticError(
                        f"Encountered a conditional jump that has only 1 successor on {self.starts}!"
                    )

                b1_og_succs = list(self.original_graph.successors(b1_og))
                if len(b1_og_succs) != 1:
                    raise SAILRSemanticError(
                        "Encountered a merge-end pair which ends in a return, this should be skipped!"
                    )

                b1 = b1_og_succs[0]
            else:
                b0, b1 = merge_end_pair

            if true_target == self._find_og_start_by_merge_end(b0):
                cond_jump_stmt.true_target.value = b0.addr
                cond_jump_stmt.false_target.value = b1.addr
            else:
                cond_jump_stmt.false_target.value = b0.addr
                cond_jump_stmt.true_target.value = b1.addr

            self.graph.add_edge(match_node, cond_copy)
            self.graph.add_edge(cond_copy, b0)
            self.graph.add_edge(cond_copy, b1)

    def create_mapping_to_merge_graph(self, updated_blocks: dict[Block, Block], start_blocks: list[Block]):
        merge_base, other_base = self.starts[:]

        # collect all the normal blocks, which should be easy to pair
        for base, other in zip(self.original_blocks[merge_base], self.original_blocks[other_base]):
            if base in updated_blocks:
                base = updated_blocks[base]
            self.merge_blocks_to_originals[base] = {base, other}

        # collect all the split blocks, which have new addrs, and match them to merge_base
        for sblock in self.original_split_blocks[merge_base]:
            attrs = ("up_split", "match_split", "down_split")
            for attr in attrs:
                sblock_split = getattr(sblock, attr)
                if not sblock_split:
                    continue
                self.merge_blocks_to_originals[sblock_split].add(sblock)

        # finally, use the info from merge_base mapping to add to the mapping for other_base
        for sblock in self.original_split_blocks[other_base]:
            attrs = ("up_split", "match_split", "down_split")
            base_original = self._find_merge_block_by_original(sblock.original)
            for attr in attrs:
                sblock_split = getattr(sblock, attr)
                if not sblock_split:
                    continue

                if attr != "match_split":
                    self.merge_blocks_to_originals[sblock_split].add(sblock)
                    continue

                for base_split, base_sblocks in self.merge_blocks_to_originals.items():
                    # this split block must have only one pair right now
                    if len(base_sblocks) > 1:
                        continue

                    base_sblock = next(iter(base_sblocks))
                    # the original base must match
                    if base_sblock.original != base_original:
                        continue

                    # this split block must also be the right split (up, match, down)
                    if base_split != getattr(base_sblock, attr):
                        continue

                    base_sblocks.add(sblock)

        # remove any extra block that will not be in the graph because it was split up
        deletable_blocks = set()
        end_pair_map = {}
        end_pairs = set()
        merge_to_end_pair = {}
        for sblocks in self.merge_blocks_to_originals.values():
            for sblock in sblocks:
                if isinstance(sblock, AILBlockSplit) and sblock.original in self.merge_blocks_to_originals:
                    deletable_blocks.add(sblock.original)
        for del_block in deletable_blocks:
            b0, b1 = self.merge_blocks_to_originals[del_block]
            end_pair_map[b0] = b1
            end_pair_map[b1] = b0
            end_pairs.add((b0, b1))
            merge_to_end_pair[del_block] = (b0, b1)
            del self.merge_blocks_to_originals[del_block]

        # mappings are collected, now we need to connect the merge ends of the graph to the upper parts with
        # the conditions that blocked the original
        merge_ends = {
            block: split for block, split in self.merge_blocks_to_originals.items() if block not in self.graph
        }
        assert all(len(bs) == 1 for block, bs in self.merge_blocks_to_originals.items() if block in merge_ends)
        self.merge_end_to_start = {merge_end: self._find_og_start_by_merge_end(merge_end) for merge_end in merge_ends}
        merge_end_pairs = defaultdict(list)
        for merge_end, split_blocks in merge_ends.items():
            split_block = next(iter(split_blocks))
            if split_block.original in self.original_blocks[merge_base]:
                merge_end_pairs[split_block.match_split].append(merge_end)
                continue

            merge_original = end_pair_map.get(split_block.original, None) or split_block.original
            base_split = self._find_split_block_by_original(merge_original)
            merge_end_pairs[base_split.match_split].append(merge_end)

        for pair in end_pairs:
            for block in pair:
                self.original_ends.append(block)
        # the case of single block merges
        if not self.original_ends:
            self.original_ends = start_blocks

        # moved here
        for _, pair in merge_to_end_pair.items():
            for block in pair:
                other_block = pair[0] if pair[1] is block else pair[1]
                while True:
                    for merge, originals in self.merge_blocks_to_originals.items():
                        if self._block_in_originals(merge, other_block):
                            continue

                        if len(originals) == 1:
                            og = next(iter(originals))
                            if isinstance(og, Block) and og == block:
                                originals.add(other_block)
                                break

                        found = False
                        for og in originals:
                            if isinstance(og, AILBlockSplit) and og.original == block and merge == og.match_split:
                                originals.add(other_block)
                                found = True
                                break

                        if found:
                            break

                    else:
                        break

        return merge_end_pairs

    def merged_is_split_type(self, merge_block: Block, split_type: str):
        if split_type not in ["up_split", "down_split", "match_split"]:
            raise ValueError("Must use a supported split type (up_split, down_split, match_split)")

        for oblock in self.merge_blocks_to_originals[merge_block]:
            if not isinstance(oblock, AILBlockSplit):
                continue

            if getattr(oblock, split_type) == merge_block:
                return True

        return False

    #
    # Private Functions
    #

    def _find_block_pair_in_originals(self, block: Block):
        for _, originals in self.merge_blocks_to_originals.items():
            # need at least 2 for a pair
            if len(originals) < 2:
                continue

            # the block we are searching for a pair for should exist in the originals
            for og in originals:
                if isinstance(og, Block) and og == block:
                    break
                if isinstance(og, AILBlockSplit) and og.original == block:
                    break
            else:
                return None

            # we now know that our target block is in this originals set
            # now we just need to find any other block that is not itself
            for other_block in originals:
                if isinstance(other_block, Block) and other_block != block:
                    return other_block
                if isinstance(other_block, AILBlockSplit) and other_block.original != block:
                    return other_block.original

        return None

    def _block_in_originals(self, merge_block: Block, target_block: Block):
        for original in self.merge_blocks_to_originals[merge_block]:
            if (isinstance(original, Block) and original == target_block) or (
                isinstance(original, AILBlockSplit) and original.original == target_block
            ):
                return True

        return False

    def _update_all_split_refs(self, update_map: dict[Block, Block]):
        for original, updated in update_map.items():
            for k in list(self.original_split_blocks.keys()):
                if k == original:
                    self.original_split_blocks[updated] = self.original_split_blocks[k]
                    del self.original_split_blocks[k]

            for v in self.original_split_blocks.values():
                for sblock in v:
                    for attr in ["up_split", "match_split", "down_split"]:
                        if getattr(sblock, attr) == original:
                            setattr(sblock, attr, updated)

            for k in list(self.original_blocks.keys()):
                if k == original:
                    self.original_blocks[updated] = self.original_blocks[k]
                    del self.original_blocks[k]

    def _find_merge_block_by_original(self, block: Block):
        for merge_block, originals in self.merge_blocks_to_originals.items():
            for og in originals:
                if isinstance(og, AILBlockSplit) and og.original == block:
                    return merge_block

                if isinstance(og, Block) and og == block:
                    return merge_block

        _l.warning("Error in finding the merge block from the original block on %s", block)
        return None

    def _find_split_block_by_original(self, block: Block) -> AILBlockSplit | None:
        for _, split_blocks in self.original_split_blocks.items():
            for split_block in split_blocks:
                if split_block.original == block:
                    return split_block

        _l.warning("Error in finding split block by original on %s", block)
        return None

    def _find_og_start_by_merge_end(self, merge_end: Block) -> Block | None:
        original = next(iter(self.merge_blocks_to_originals[merge_end]))
        if isinstance(original, AILBlockSplit):
            original = original.original

        for og_start, og_blocks in self.original_blocks.items():
            if original in og_blocks:
                return og_start

        return None

    @staticmethod
    def clone_graph_replace_splits(graph: nx.DiGraph, split_map: dict[Block, Block]):
        # do a deepcopy, so we don't edit the original graph
        graph = copy_graph_and_nodes(graph)

        # replace every block that has been split
        updated_blocks = {}
        for original_node, new_node in split_map.items():
            graph.add_node(new_node)
            # correct every in_edge to this node, with new targets for jumps
            for pred in list(graph.predecessors(original_node)):
                new_pred = pred.copy()
                new_pred.statements[-1] = correct_jump_targets(
                    new_pred.statements[-1], {original_node.addr: new_node.addr}, new_stmt=True
                )
                updated_blocks[pred] = new_pred
                replace_node_in_graph(graph, pred, new_pred)
                graph.add_edge(new_pred, new_node)

            # re-add every out_edge
            blk = updated_blocks.get(original_node, original_node)
            for succ in graph.successors(blk):
                graph.add_edge(new_node, succ)

            # finally, kill the original
            if original_node in graph:
                graph.remove_node(original_node)

        return graph, updated_blocks
