from ..core import State, ModelState
from .diagrams import GraphMachine
from .diagrams_base import BaseGraph
from logging import Logger
from typing import BinaryIO, Type, Optional, Dict, List, Union, DefaultDict, Any, Iterable
try:
    from graphviz import Digraph
    from graphviz.dot import SubgraphContext
except ImportError:
    class Digraph:  # type: ignore
        pass

    class SubgraphContext:  # type: ignore
        pass

_LOGGER: Logger

class Graph(BaseGraph):
    custom_styles: Dict[str, DefaultDict[str, Union[str, DefaultDict[str, str]]]]
    def __init__(self, machine: Type[GraphMachine]) -> None: ...
    def set_previous_transition(self, src: str, dst: str) -> None: ...
    def set_node_style(self, state: ModelState, style: str) -> None: ...
    def reset_styling(self) -> None: ...
    def _add_nodes(self, states: List[Dict[str, str]],   # type: ignore[no-any-unimported]
                   container: Union[Digraph, SubgraphContext]) -> None: ...
    def _add_edges(self, transitions: List[Dict[str, str]],  # type: ignore[no-any-unimported]
                   container: Union[Digraph, SubgraphContext]) -> None: ...
    def generate(self) -> None: ...
    def get_graph(self, title: Optional[str] = ...,  # type: ignore[no-any-unimported]
                  roi_state: Optional[str] = ...) -> Digraph: ...
    def draw(self, filename: Optional[Union[str, BinaryIO]], format:Optional[str] = ...,
             prog: Optional[str] = ..., args:str = ...) -> Optional[str]: ...

class NestedGraph(Graph):
    _cluster_states: List[str]
    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def set_previous_transition(self, src: str, dst: str) -> None: ...
    def _add_nodes(self, states: List[Dict[str, str]],  # type: ignore[no-any-unimported]
                   container: Union[Digraph, SubgraphContext]) -> None: ...
    def _add_nested_nodes(self,  # type: ignore[no-any-unimported]
                          states: List[Dict[str, Union[str, List[Dict[str, str]]]]],
                          container: Union[Digraph, SubgraphContext],
                          prefix: str, default_style: str) -> None: ...
    def _add_edges(self, transitions: List[Dict[str, str]],  # type: ignore[no-any-unimported]
                   container: Union[Digraph, SubgraphContext]) -> None: ...
    def _create_edge_attr(self, src: str, dst: str, transition: Dict[str, str]) -> Dict[str, Any]: ...

def filter_states(states: List[Dict[str, str]], state_names: Iterable[str], state_cls: Type[State],
                  prefix: Optional[List[str]] = ...) -> List[Dict[str, str]]: ...
