# This file is part of Androguard.
#
# Copyright (c) 2012 Geoffroy Gueguen <geoffroy.gueguen@gmail.com>
# All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from collections import defaultdict
from androguard.decompiler.dad.basic_blocks import (build_node_from_block,
                                                    StatementBlock, CondBlock)
from androguard.decompiler.dad.instruction import Variable

logger = logging.getLogger('dad.graph')


class Graph(object):
    def __init__(self):
        self.entry = None
        self.exit = None
        self.nodes = list()
        self.rpo = []
        self.edges = defaultdict(list)
        self.catch_edges = defaultdict(list)
        self.reverse_edges = defaultdict(list)
        self.reverse_catch_edges = defaultdict(list)
        self.loc_to_ins = None
        self.loc_to_node = None

    def sucs(self, node):
        return self.edges.get(node, [])

    def all_sucs(self, node):
        return self.edges.get(node, []) + self.catch_edges.get(node, [])

    def preds(self, node):
        return [n for n in self.reverse_edges.get(node, [])
                if not n.in_catch]

    def all_preds(self, node):
        return (self.reverse_edges.get(node, []) +
                self.reverse_catch_edges.get(node, []))

    def add_node(self, node):
        self.nodes.append(node)

    def add_edge(self, e1, e2):
        lsucs = self.edges[e1]
        if e2 not in lsucs:
            lsucs.append(e2)
        lpreds = self.reverse_edges[e2]
        if e1 not in lpreds:
            lpreds.append(e1)

    def add_catch_edge(self, e1, e2):
        lsucs = self.catch_edges[e1]
        if e2 not in lsucs:
            lsucs.append(e2)
        lpreds = self.reverse_catch_edges[e2]
        if e1 not in lpreds:
            lpreds.append(e1)

    def remove_node(self, node):
        preds = self.reverse_edges.get(node, [])
        for pred in preds:
            self.edges[pred].remove(node)

        succs = self.edges.get(node, [])
        for suc in succs:
            self.reverse_edges[suc].remove(node)

        exc_preds = self.reverse_catch_edges.pop(node, [])
        for pred in exc_preds:
            self.catch_edges[pred].remove(node)

        exc_succs = self.catch_edges.pop(node, [])
        for suc in exc_succs:
            self.reverse_catch_edges[suc].remove(node)

        self.nodes.remove(node)
        if node in self.rpo:
            self.rpo.remove(node)
        del node

    def number_ins(self):
        self.loc_to_ins = {}
        self.loc_to_node = {}
        num = 0
        for node in self.rpo:
            start_node = num
            num = node.number_ins(num)
            end_node = num - 1
            self.loc_to_ins.update(node.get_loc_with_ins())
            self.loc_to_node[(start_node, end_node)] = node

    def get_ins_from_loc(self, loc):
        return self.loc_to_ins.get(loc)

    def get_node_from_loc(self, loc):
        for (start, end), node in self.loc_to_node.iteritems():
            if start <= loc <= end:
                return node

    def remove_ins(self, loc):
        ins = self.get_ins_from_loc(loc)
        self.get_node_from_loc(loc).remove_ins(loc, ins)
        self.loc_to_ins.pop(loc)

    def split_if_nodes(self):
        '''
        Split IfNodes in two nodes, the first node is the header node, the
        second one is only composed of the jump condition.
        '''
        node_map = {n: n for n in self.nodes}
        to_update = set()
        for node in self.nodes[:]:
            if node.type.is_cond:
                if len(node.get_ins()) > 1:
                    pre_ins = node.get_ins()[:-1]
                    last_ins = node.get_ins()[-1]
                    pre_node = StatementBlock('%s-pre' % node.name, pre_ins)
                    cond_node = CondBlock('%s-cond' % node.name, [last_ins])
                    node_map[node] = pre_node
                    node_map[pre_node] = pre_node
                    node_map[cond_node] = cond_node

                    pre_node.copy_from(node)
                    cond_node.copy_from(node)
                    for var in node.var_to_declare:
                        pre_node.add_variable_declaration(var)
                    pre_node.type.is_stmt = True
                    cond_node.true = node.true
                    cond_node.false = node.false

                    for pred in self.all_preds(node):
                        pred_node = node_map[pred]
                        # Verify that the link is not an exception link
                        if node not in self.sucs(pred):
                            self.add_catch_edge(pred_node, pre_node)
                            continue
                        if pred is node:
                            pred_node = cond_nod
e
                        if pred.type.is_cond:  # and not (pred is node):
                            if pred.true is node:
                                pred_node.true = pre_node
                            if pred.false is node:
                            
    pred_node.false = pre_node
                        self.add_edge(pred_node, pre_node)
                    for suc in self.sucs(node):
                        self.add_edge(cond_node, node_map[suc])

                    # We link all the exceptions to the pre node instead of the
                    # condition node, which should not trigger any of them.
                    for suc in self.catch_edges.get(node, []):
                        self.add_catch_edge(pre_node, node_map[suc])

                    if node is self.entry:
                        self.entry = pre_node

                    self.add_node(pre_node)
                    self.add_node(cond_node)
                    self.add_edge(pre_node, cond_node)
                    pre_node.update_attribute_with(node_map)
                    cond_node.update_attribute_with(node_map)
                    self.remove_node(node)
            else:
                to_update.add(node)
        for node in to_update:
            node.update_attribute_with(node_map)

    def simplify(self):
        '''
        Simplify the CFG by merging/deleting statement nodes when possible:
        If statement B follows statement A and if B has no other predecessor
        besides A, then we can merge A and B into a new statement node.
        We also remove nodes which do nothing except redirecting the control
        flow (nodes which only contains a goto).
        '''
        redo = True
        while redo:
            redo = False
            node_map = {}
            to_update = set()
            for node in self.nodes[:]:
                if node.type.is_stmt and node in self.nodes:
                    sucs = self.all_sucs(node)
                    if len(sucs) != 1:
                        continue
                    suc = sucs[0]
                    if len(node.get_ins()) == 0:
                        if any(pred.type.is_switch
                               for pred in self.all_preds(node)):
                            continue
                        if node is suc:
                            continue
                        node_map[node] = suc

                        for pred in self.all_preds(node):
                            pred.update_attribute_with(node_map)
                            if node not in self.sucs(pred):
                                self.add_catch_edge(pred, suc)
                                continue
                            self.add_edge(pred, suc)
                        redo = True
                        if node is self.entry:
                            self.entry = suc
                     