# -*- coding: utf-8 -*-
"""
State Machine Visualizer
~~~~~~~~~~~~~~~~~~~~~~~~

This code provides a module that can use graphviz to visualise the state
machines included in h2. These visualisations can be used as part of the
documentation of h2, and as a reference material to understand how the
state machines function.

The code in this module is heavily inspired by code in Automat, which can be
found here: https://github.com/glyph/automat. For details on the licensing of
Automat, please see the NOTICES.visualizer file in this folder.

This module is very deliberately not shipped with the rest of h2. This is
because it is of minimal value to users who are installing h2: its use
is only really for the developers of h2.
"""
import argparse
import collections
import sys

import graphviz
import graphviz.files

import h2.connection
import h2.stream


StateMachine = collections.namedtuple(
    'StateMachine', ['fqdn', 'machine', 'states', 'inputs', 'transitions']
)


# This is all the state machines we currently know about and will render.
# If any new state machines are added, they should be inserted here.
STATE_MACHINES = [
    StateMachine(
        fqdn='h2.connection.H2ConnectionStateMachine',
        machine=h2.connection.H2ConnectionStateMachine,
        states=h2.connection.ConnectionState,
        inputs=h2.connection.ConnectionInputs,
        transitions=h2.connection.H2ConnectionStateMachine._transitions,
    ),
    StateMachine(
        fqdn='h2.stream.H2StreamStateMachine',
        machine=h2.stream.H2StreamStateMachine,
        states=h2.stream.StreamState,
        inputs=h2.stream.StreamInputs,
        transitions=h2.stream._transitions,
    ),
]


def quote(s):
    return '"{}"'.format(s.replace('"', r'\"'))


def html(s):
    return '<{}>'.format(s)


def element(name, *children, **attrs):
    """
    Construct a string from the HTML element description.
    """
    formatted_attributes = ' '.join(
        '{}={}'.format(key, quote(str(value)))
        for key, value in sorted(attrs.items())
    )
    formatted_children = ''.join(children)
    return u'<{name} {attrs}>{children}</{name}>'.format(
        name=name,
        attrs=formatted_attributes,
        children=formatted_children
    )


def row_for_output(event, side_effect):
    """
    Given an output tuple (an event and its side effect), generates a table row
    from it.
    """
    point_size = {'point-size': '9'}
    event_cell = element(
        "td",
        element("font", enum_member_name(event), **point_size)
    )
    side_effect_name = (
        function_name(side_effect) if side_effect is not None else "None"
    )
    side_effect_cell = element(
        "td",
        element("font", side_effect_name, **point_size)
    )
    return element("tr", event_cell, side_effect_cell)


def table_maker(initial_state, final_state, outputs, port):
    """
    Construct an HTML table to label a state transition.
    """
    header = "{} -&gt; {}".format(
        enum_member_name(initial_state), enum_member_name(final_state)
    )
    header_row = element(
        "tr",
        element(
            "td",
            element(
                "font",
                header,
                face="menlo-italic"
            ),
            port=port,
            colspan="2",
        )
    )
    rows = [header_row]
    rows.extend(row_for_output(*output) for output in outputs)
    return element("table", *rows)


def enum_member_name(state):
    """
    All enum member names have the form <EnumClassName>.<EnumMemberName>. For
    our rendering we only want the member name, so we take their representation
    and split it.
    """
    return str(state).split('.', 1)[1]


def function_name(func):
    """
    Given a side-effect function, return its string name.
    """
    return func.__name__


def build_digraph(state_machine):
    """
    Produce a L{graphviz.Digraph} object from a state machine.
    """
    digraph = graphviz.Digraph(node_attr={'fontname': 'Menlo'},
                               edge_attr={'fontname': 'Menlo'},
                               graph_attr={'dpi': '200'})

    # First, add the states as nodes.
    seen_first_state = False
    for state in state_machine.states:
        if not seen_first_state:
            state_shape = "bold"
            font_name = "Menlo-Bold"
        else:
            state_shape = ""
            font_name = "Menlo"
        digraph.node(enum_member_name(state),
                     fontame=font_name,
                     shape="ellipse",
                     style=state_shape,
                     color="blue")
        seen_first_state = True

    # We frequently have vary many inputs that all trigger the same state
    # transition, and only differ in terms of their input and side-effect. It
    # would be polite to say that graphviz does not handle this very well. So
    # instead we *collapse* the state transitions all into the one edge, and
    # then provide a label that displays a table of all the inputs and their
    # associated side effects.
    transitions = collections.defaultdict(list)
    for transition in state_machine.transitions.items():
        initial_state, event = transition[0]
        side_effect, final_state = transition[1]
        transition_key = (initial_state, final_state)
        transitions[transition_key].append((event, side_effect))

    for n, (transition_key, outputs) in enumerate(transitions.items()):
        this_transition = "t{}".format(n)
        initial_state, final_state = transition_key

        port = "tableport"
        table = table_maker(
            initial_state=initial_state,
            final_state=final_state,
            outputs=outputs,
            port=port
        )

        digraph.node(this_transition,
                     label=html(table), margin="0.2", shape="none")

        digraph.edge(enum_member_name(initial_state),
                     '{}:{}:w'.format(this_transition, port),
                     arrowhead="none")
        digraph.edge('{}:{}:e'.format(this_transition, port),
                     enum_member_name(final_state))

    return digraph


def main():
    """
    Renders all the state machines in h2 into images.
    """
    program_name = sys.argv[0]
    argv = sys.argv[1:]

    description = """
    Visualize h2 state machines as graphs.
    """
    epilog = """
    You must have the graphviz tool suite installed.  Please visit
    http://www.graphviz.org for more information.
    """

    argument_parser = argparse.ArgumentParser(
        prog=program_name,
        description=description,
        epilog=epilog
    )
    argument_parser.add_argument(
        '--image-directory',
        '-i',
        help="Where to write out image files.",
        default=".h2_visualize"
    )
    argument_parser.add_argument(
        '--view',
        '-v',
        help="View rendered graphs with default image viewer",
        default=False,
        action="store_true"
    )
    args = argument_parser.parse_args(argv)

    for state_machine in STATE_MACHINES:
        print(state_machine.fqdn, '...discovered')

        digraph = build_digraph(state_machine)

        if args.image_directory:
            digraph.format = "png"
            digraph.render(filename="{}.dot".format(state_machine.fqdn),
                           directory=args.image_directory,
                           view=args.view,
                           cleanup=True)
            print(state_machine.fqdn, "...wrote image into", args.image_directory)


if __name__ == '__main__':
    main()