# -*- coding: utf-8 -*-
import pygraphviz as gv  # type: ignore
import itertools as it

from typing import (
    List,
    Optional,
)

from pyfsa.lib.types import TransitionsTable


def get_state_graph(
    transitions: TransitionsTable,
    start: Optional[str] = None,
    end: Optional[str] = None,
    nodes: Optional[List[str]] = None,
    name: str = 'output.png',
    draw: bool = True,
    engine: str = 'circo',
) -> gv.AGraph:
    '''
    From a transition dictionary, creates a pygraphviz graph
    of all the possible states and how to reach the given state.

    Returns the resulting graph.
    '''
    graph = gv.AGraph(directed=True, strict=False, ranksep='1')
    key_num = it.count()
    if nodes is not None:
        graph.add_nodes_from(nodes)
    else:
        graph.add_nodes_from(transitions.keys())

    for node, transition_row in transitions.items():
        for label, targets in transition_row.items():
            for target in targets:
                graph.add_edge(
                    node,
                    target,
                    key=f'{next(key_num)}',
                    label=label,
                    weight=1,
                )

    if start:
        n: gv.Node = graph.get_node(start)
        n.attr['color'] = '#0000FF'
        n.attr['style'] = 'filled'

    if end:
        n = graph.get_node(end)
        n.attr['color'] = '#00FF00'
        n.attr['style'] = 'filled'

    if draw:
        graph.layout(prog=engine)
        graph.draw(name)

    return graph


def verify_string(
    string: str,
    starting_state: str,
    final_state: str,
    transitions: TransitionsTable,
) -> bool:
    '''
    Given a transitions table, a start and end state, and
    some string, verifies that executing the finite state machine
    on the given string produces the desired final state.
    '''
    current_state = starting_state
    for letter in string:
        transition = transitions[current_state]
        current_state = transition[letter][0]

    return current_state == final_state


def render_string_graph(
    string: str,
    start: str,
    end: str,
    transitions: TransitionsTable,
    name: str = 'output.png',
    draw: bool = True,
    engine: str = 'circo'
) -> gv.AGraph:
    '''
    Given a string, a start state, an end state, end a
    transitions table, produces the graph resulting in
    the traversal of the string through the states defined
    in the transitions table. By default, it will
    output a png file of the result, but that can be
    suppressed.
    '''
    graph = gv.AGraph(directed=True)

    graph.graph_attr['label'] = f'Evaluating {string}'

    node_names = it.count()
    current_state = start
    node_name = next(node_names)
    graph.add_node(node_name)
    current_node = gv.Node(graph, node_name)
    current_node.attr['label'] = current_state
    current_node.attr['fillcolor'] = '#0000FF'
    current_node.attr['style'] = 'filled'

    for letter in string:
        node_name = next(node_names)
        graph.add_node(node_name)
        next_node = gv.Node(graph, node_name)
        # TODO: The algorithm prioritizes just the first
        # found state, which may not produce a correct
        # answer. Needs to fix this
        next_state = transitions[current_state][letter][0]
        next_node.attr['label'] = next_state
        graph.add_edge(current_node, next_node, label=letter)
        current_node = next_node
        current_state = next_state

    if current_state == end:
        current_node.attr['style'] = 'filled'
        current_node.attr['fillcolor'] = '#00FF00'

    if draw:
        graph.layout(prog=engine)
        graph.draw(name)

    return graph

[(28, 'pygraphviz.AGraph', 'gv.AGraph', True, 'import pygraphviz as gv\n'), (29, 'itertools.count', 'it.count', True, 'import itertools as it\n'), (99, 'pygraphviz.AGraph', 'gv.AGraph', True, 'import pygraphviz as gv\n'), (103, 'itertools.count', 'it.count', True, 'import itertools as it\n'), (107, 'pygraphviz.Node', 'gv.Node', True, 'import pygraphviz as gv\n'), (115, 'pygraphviz.Node', 'gv.Node', True, 'import pygraphviz as gv\n')]