import re
import inspect
import networkx
from networkx.drawing.nx_pydot import to_pydot

NOT_DECORATORS = ['property']

def _get_jump_edges(obj):
    dest_nodes = set()
    for source_line in inspect.getsourcelines(obj)[0]:
        match = re.match(r"\s*return '(.*)'\s*$", source_line)
        if match:
            dest_nodes.update([match.group(1)])
    return dest_nodes

def _is_comment_or_blank(source_line):
    return re.match(r"\s*#", source_line) or re.match(r"^\s*$", source_line)

def find_jump_edges(guardstate, module):
    """Find all jump edges for a GuardState object."""

    dest_nodes = _get_jump_edges(guardstate)

    source_lines = inspect.getsourcelines(guardstate)[0]
    i = 0
    while i < len(source_lines):
        # look for jump edges in object decorator source code
        # decorators can be stacked on top of each other, so we need to
        # be sure to get all of them
        decorator_matches = set()
        j = i
        while True:
            match = re.match(r"\s*@([^\s\(]*)(\(.*\))?\s*$", source_lines[j])
            if _is_comment_or_blank(source_lines[j]):
                pass
            elif match:
                decorator_matches.update([match.group(1)])
            else:
                break
            j += 1

        # we know these are decorators if they are followed by a function definition
        if decorator_matches:
            while True:
                if _is_comment_or_blank(source_lines[j]):
                    pass
                elif re.match(r".*\bdef\b.*", source_lines[j]):
                    # look for jump edges in decorator code
                    for decorator_match in decorator_matches:
                        if decorator_match in NOT_DECORATORS:
                            continue
                        dest_nodes.update(_get_jump_edges(getattr(module, decorator_match)))
                        break
                else:
                    break
                j += 1

        i = j+1

    return dest_nodes

def fullgraph(system, find_jumps=False):
    """Return MultiDiGraph copy of system graph with JUMP edges added."""
    G = networkx.MultiDiGraph(system.graph)
    for node, code in G.nodes(data='code'):
        obj = code()
        if find_jumps:
            for jump in find_jump_edges(code, system.module):
                G.add_edge(node, jump, jump=True, goto=False, weight=1)
    return G

def sys2dot(system,
            path=None,
            show_index=False,
            show_gotos=False,
            show_jumps=False,
            edge_constraints=True):
    """Create graphviz dot code for system graph."""

    # if states is a tuple it specifies a start and stop state, so
    # calc path
    fpath = []
    path_edges = []
    if path:
        assert isinstance(path, tuple), "'path' should be tuple"
        if len(path) > 1:
            fpath = system.shortest_path(path[0], path[1])
        else:
            fpath = path

        for ind in range(len(fpath)-1):
            path_edges.append((fpath[ind], fpath[ind+1]))

    G = fullgraph(system, find_jumps=show_jumps)

    G.graph['graph'] = {
        'label': '<%s Guardian: %s>' % (system.ifo, system.name),
        'labelloc': 't',
        'rankdir': 'BT',
    }

    # http://www.graphviz.org/content/attrs
    for node, data in G.nodes(data=True):
        # node attribute defaults
        label = '%s' % (node)
        shape = 'oval'
        style = 'filled'
        colorscheme = 'set19'
        color = '2'
        outline = 'black'
        fontcolor = 'white'
        penwidth = 1

        # state index
        if show_index:
            label = '%s (%d)' % (node, data['index'])

        # if not requestable use pastel shade
        if not data['request']:
            colorscheme = 'pastel19'
            fontcolor = 'grey30'

        # if not redirectable use red outline
        if not data['redirect']:
            outline = 'red'
            #style = 'dashed, filled'
            penwidth = 2

        # if node is goto use green
        if data['goto']:
            color = '3'

        # INIT state
        if node == 'INIT':
            shape = 'doublecircle'
            color = '4'

        # if node in path
        if fpath and node in fpath:
            penwidth = 8
            colorscheme = 'set19'
            if node == fpath[0]:
                color = 3
            elif node == fpath[-1]:
                color = 1
            else:
                color = 5

        # set attributes
        G.nodes[node]['label'] = label
        G.nodes[node]['shape'] = shape
        G.nodes[node]['style'] = style
        G.nodes[node]['fillcolor'] = '/%s/%s' % (colorscheme, color)
        G.nodes[node]['penwidth'] = penwidth
        G.nodes[node]['color'] = outline
        G.nodes[node]['fontcolor'] = fontcolor

    # set edge attributes
    remove_edges = []
    for edge in G.edges(data=True, keys=True):
        source, sink, key, data = edge

        # label edges with their weight
        if data['weight'] != 1:
            data['label'] = data['weight']

        # graphviz handles edge weights opposite of networkx.  edge
        # weights add to path length in networkx whereas large weights
        # represent "tighter" couplings between nodes in graphviz.
        # therefore flip the sign of the weight for graphviz, so that
        # larger weights represent weaker coupling
        data['weight'] *= -1

        # goto edges
        if data['goto']:
            if show_gotos:
                if not edge_constraints:
                    data['constraint'] = False
                data['color'] = 'grey'
                data['fontcolor'] = 'grey'
            else:
                if edge[:2] not in path_edges:
                    remove_edges.append((source, sink, key))

        # jump edges
        if data.get('jump', False):
            if not edge_constraints:
                data['constraint'] = False
            data['style'] = 'dashed'
            data['color'] = 'red'

        # path edges
        if edge[:2] in path_edges:
            data['color'] = '/set19/5'
            data['penwidth'] = 5

    for edge in remove_edges:
        G.remove_edge(*edge)

    # do the graphviz conversion
    dot = to_pydot(G)

    return dot
