Diagram Generation

class DotGraphMachine:
    """
    Generates graphical representations of state machines using Graphviz.
    
    Provides customizable visualization with support for different layouts,
    colors, fonts, and styling options.
    
    Class Attributes:
    - graph_rankdir: Direction of the graph layout ("LR" for left-right, "TB" for top-bottom)
    - font_name: Font face name for graph text
    - state_font_size: Font size for state labels in points
    - state_active_penwidth: Line width for active state borders
    - state_active_fillcolor: Fill color for active states
    - transition_font_size: Font size for transition labels in points
    """
    
    graph_rankdir: str = "LR"
    """Direction of the graph. Defaults to "LR" (option "TB" for top bottom)."""
    
    font_name: str = "Arial"
    """Graph font face name."""
    
    state_font_size: str = "10"
    """State font size in points."""
    
    state_active_penwidth: int = 2
    """Active state external line width."""
    
    state_active_fillcolor: str = "turquoise"
    """Active state fill color."""
    
    transition_font_size: str = "9"
    """Transition font size in points."""
    
    def __init__(self, machine: StateMachine):
        """
        Initialize diagram generator for a state machine.
        
        Parameters:
        - machine: StateMachine instance to visualize
        """
    
    def get_graph(self) -> pydot.Dot:
        """
        Generate and return the Graphviz digraph object.
        
        Returns:
        pydot.Dot object that can be rendered to various formats
        """
    
    def create_digraph(self) -> pydot.Dot:
        """
        Create a Graphviz digraph object representing the state machine.
        
        Returns:
        pydot.Dot object that can be rendered to various formats
        """
    
    def __call__(self) -> pydot.Dot:
        """Alias for get_graph() - allows calling instance as function."""
    
    def write(self, filename: str, format: str = "png", prog: str = "dot"):
        """
        Write diagram to file.
        
        Parameters:
        - filename: Output file path
        - format: Output format (png, svg, pdf, etc.)
        - prog: Graphviz layout program (dot, neato, fdp, etc.)
        """

# Built-in diagram generation method
class StateMachine:
    def _graph(self) -> DotGraphMachine:
        """Get DotGraphMachine instance for this state machine."""

def quickchart_write_svg(sm: StateMachine, path: str):
    """
    Write state machine diagram as SVG using QuickChart service.
    
    Generates diagram without requiring local Graphviz installation
    by using the QuickChart.io online service.
    
    Parameters:
    - sm: StateMachine instance to visualize
    - path: Output SVG file path
    """

def write_image(qualname: str, out: str):
    """
    Write state machine diagram to file by class qualname.
    
    Parameters:
    - qualname: Fully qualified name of StateMachine class
    - out: Output file path
    """

def import_sm(qualname: str):
    """
    Import StateMachine class by fully qualified name.
    
    Parameters:
    - qualname: Fully qualified class name
    
    Returns:
    StateMachine class
    """

def main(argv=None):
    """
    Main entry point for command line diagram generation.
    
    Supports generating diagrams from command line using:
    python -m statemachine.contrib.diagram <qualname> <output_file>
    """

from statemachine import StateMachine, State
from statemachine.contrib.diagram import DotGraphMachine

class TrafficLight(StateMachine):
    green = State(initial=True)
    yellow = State()
    red = State()
    
    cycle = (
        green.to(yellow)
        | yellow.to(red)
        | red.to(green)
    )

# Create state machine and generate diagram
traffic = TrafficLight()

# Method 1: Using built-in _graph() method
graph = traffic._graph()
graph.write("traffic_light.png")

# Method 2: Using DotGraphMachine directly
diagram = DotGraphMachine(traffic)
diagram.write("traffic_light_custom.png", format="png")

# Method 3: Generate SVG
diagram.write("traffic_light.svg", format="svg")

# Method 4: Generate PDF
diagram.write("traffic_light.pdf", format="pdf")

class CustomStyledDiagram(DotGraphMachine):
    # Customize appearance
    graph_rankdir = "TB"  # Top to bottom layout
    font_name = "Helvetica"
    state_font_size = "12"
    state_active_penwidth = 3
    state_active_fillcolor = "lightblue"
    transition_font_size = "10"

class OrderWorkflow(StateMachine):
    pending = State("Pending Order", initial=True)
    paid = State("Payment Received")
    shipped = State("Order Shipped")
    delivered = State("Delivered", final=True)
    cancelled = State("Cancelled", final=True)
    
    pay = pending.to(paid)
    ship = paid.to(shipped)
    deliver = shipped.to(delivered)
    cancel = (
        pending.to(cancelled)
        | paid.to(cancelled)
        | shipped.to(cancelled)
    )

# Create workflow and generate custom styled diagram
workflow = OrderWorkflow()
custom_diagram = CustomStyledDiagram(workflow)
custom_diagram.write("order_workflow_custom.png")

# Set current state and regenerate to show active state
workflow.send("pay")
workflow.send("ship")
active_diagram = CustomStyledDiagram(workflow)
active_diagram.write("order_workflow_active.png")

from statemachine.contrib.diagram import quickchart_write_svg

class SimpleSwitch(StateMachine):
    off = State("Off", initial=True)
    on = State("On")
    
    turn_on = off.to(on)
    turn_off = on.to(off)

# Generate diagram using online service (no local Graphviz required)
switch = SimpleSwitch()
quickchart_write_svg(switch, "switch_diagram.svg")

print("Diagram generated using QuickChart service")

class ComplexWorkflow(StateMachine):
    # Define multiple states with descriptive names
    draft = State("Draft Document", initial=True)
    review = State("Under Review")
    revision = State("Needs Revision")
    approved = State("Approved")
    published = State("Published", final=True)
    rejected = State("Rejected", final=True)
    archived = State("Archived", final=True)
    
    # Define complex transition network
    submit = draft.to(review)
    approve = review.to(approved)
    request_revision = review.to(revision)
    resubmit = revision.to(review)
    reject = review.to(rejected) | approved.to(rejected)
    publish = approved.to(published)
    archive = (
        draft.to(archived)
        | rejected.to(archived)
        | published.to(archived)
    )
    
    # Add transition conditions for better diagram labeling
    urgent_publish = approved.to(published).cond("is_urgent")
    regular_publish = approved.to(published).unless("is_urgent")
    
    def is_urgent(self, priority: str = "normal"):
        return priority == "urgent"

# Generate comprehensive diagram
workflow = ComplexWorkflow()

# Create diagram with custom settings for complex visualization
class ComplexDiagram(DotGraphMachine):
    graph_rankdir = "LR"
    font_name = "Arial"
    state_font_size = "11"
    transition_font_size = "9"
    
    def create_digraph(self):
        """Override to add custom graph attributes."""
        dot = super().create_digraph()
        
        # Add graph-level styling
        dot.set_graph_defaults(
            fontname=self.font_name,
            fontsize="14",
            labelloc="t",
            label="Document Workflow State Machine"
        )
        
        # Add node-level styling
        dot.set_node_defaults(
            shape="box",
            style="rounded,filled",
            fillcolor="lightyellow",
            fontname=self.font_name
        )
        
        # Add edge-level styling
        dot.set_edge_defaults(
            fontname=self.font_name,
            fontsize=self.transition_font_size
        )
        
        return dot

complex_diagram = ComplexDiagram(workflow)
complex_diagram.write("complex_workflow.png", format="png")
complex_diagram.write("complex_workflow.svg", format="svg")

class MultiFormatExample(StateMachine):
    start = State("Start", initial=True)
    process = State("Processing")
    complete = State("Complete", final=True)
    error = State("Error", final=True)
    
    begin = start.to(process)
    finish = process.to(complete)
    fail = process.to(error)
    retry = error.to(process)

# Generate diagrams in multiple formats
example = MultiFormatExample()
diagram = DotGraphMachine(example)

# Common formats
formats = {
    "png": "Portable Network Graphics",
    "svg": "Scalable Vector Graphics", 
    "pdf": "Portable Document Format",
    "jpg": "JPEG Image",
    "gif": "Graphics Interchange Format",
    "dot": "DOT Source Code"
}

for fmt, description in formats.items():
    filename = f"multi_format_example.{fmt}"
    try:
        diagram.write(filename, format=fmt)
        print(f"Generated {description}: {filename}")
    except Exception as e:
        print(f"Failed to generate {fmt}: {e}")

# Command line diagram generation examples:

# Basic usage:
# python -m statemachine.contrib.diagram myapp.machines.OrderMachine order_diagram.png

# Generate SVG:
# python -m statemachine.contrib.diagram myapp.machines.OrderMachine order_diagram.svg

# Example of programmatic command line invocation
import subprocess
import sys

def generate_diagram_cli(machine_qualname: str, output_path: str):
    """Generate diagram using command line interface."""
    try:
        result = subprocess.run([
            sys.executable, "-m", "statemachine.contrib.diagram",
            machine_qualname, output_path
        ], capture_output=True, text=True, check=True)
        
        print(f"Diagram generated successfully: {output_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error generating diagram: {e.stderr}")
        return False

# Usage
if __name__ == "__main__":
    # This would work if the machine class is importable
    success = generate_diagram_cli(
        "myproject.machines.WorkflowMachine",
        "workflow_diagram.png"
    )

# Jupyter notebook integration example
try:
    from IPython.display import Image, SVG, display
    import tempfile
    import os
    
    def display_state_machine(machine: StateMachine, format: str = "svg"):
        """Display state machine diagram in Jupyter notebook."""
        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) as tmp:
            temp_path = tmp.name
        
        try:
            # Generate diagram
            diagram = DotGraphMachine(machine)
            diagram.write(temp_path, format=format)
            
            # Display in notebook
            if format.lower() == "svg":
                with open(temp_path, 'r') as f:
                    display(SVG(f.read()))
            else:
                display(Image(temp_path))
                
        finally:
            # Clean up temporary file
            if os.path.exists(temp_path):
                os.unlink(temp_path)
    
    # Usage in Jupyter notebook
    class NotebookExample(StateMachine):
        idle = State("Idle", initial=True)
        working = State("Working")
        done = State("Done", final=True)
        
        start = idle.to(working)
        complete = working.to(done)
    
    notebook_machine = NotebookExample()
    
    # Display the diagram in notebook
    display_state_machine(notebook_machine, "svg")
    
    # Show state after transition
    notebook_machine.send("start")
    display_state_machine(notebook_machine, "svg")

except ImportError:
    print("IPython not available - skipping Jupyter notebook example")

import time
from pathlib import Path

def create_animated_sequence(machine: StateMachine, events: list, output_dir: str = "animation_frames"):
    """Create sequence of diagrams showing state transitions."""
    Path(output_dir).mkdir(exist_ok=True)
    
    frames = []
    
    # Initial state
    diagram = DotGraphMachine(machine)
    frame_path = f"{output_dir}/frame_000_initial.png"
    diagram.write(frame_path)
    frames.append(frame_path)
    
    # Process each event
    for i, event in enumerate(events, 1):
        try:
            machine.send(event)
            diagram = DotGraphMachine(machine)
            frame_path = f"{output_dir}/frame_{i:03d}_{event}.png"
            diagram.write(frame_path)
            frames.append(frame_path)
            print(f"Generated frame {i}: {event} -> {machine.current_state.id}")
        except Exception as e:
            print(f"Error processing event {event}: {e}")
            break
    
    return frames

# Create animated sequence
class AnimatedExample(StateMachine):
    state1 = State("State 1", initial=True)
    state2 = State("State 2")
    state3 = State("State 3")
    final = State("Final", final=True)
    
    next1 = state1.to(state2)
    next2 = state2.to(state3)
    finish = state3.to(final)
    reset = (state2.to(state1) | state3.to(state1))

animated_machine = AnimatedExample()
event_sequence = ["next1", "next2", "reset", "next1", "next2", "finish"]

frames = create_animated_sequence(animated_machine, event_sequence)
print(f"Generated {len(frames)} animation frames")

# Note: To create actual animated GIF, you would need additional tools like PIL or ImageIO
# Example with PIL (if available):
try:
    from PIL import Image
    
    def create_gif(frame_paths: list, output_path: str, duration: int = 1000):
        """Create animated GIF from frame images."""
        images = []
        for path in frame_paths:
            images.append(Image.open(path))
        
        images[0].save(
            output_path,
            save_all=True,
            append_images=images[1:],
            duration=duration,
            loop=0
        )
        print(f"Animated GIF created: {output_path}")
    
    # create_gif(frames, "state_machine_animation.gif", duration=1500)

except ImportError:
    print("PIL not available - cannot create animated GIF")