process_flow_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

# No longer importing add_nodes_and_edges from graph_generator_utils
# as process flow diagrams have a different structure and node addition logic.
# However, we'll re-implement the color gradient logic directly for consistency.

def generate_process_flow_diagram(json_input: str) -> str: # Removed base_color parameter
    """
    Generates a Process Flow Diagram (Flowchart) from JSON input.

    Args:
        json_input (str): A JSON string describing the process flow structure.
                          It must follow the Expected JSON Format Example below.

    Returns:
        str: The filepath to the generated PNG image file.

    Expected JSON Format Example:
    {
      "start_node": "Start Order Process",
      "nodes": [
        {"id": "customer_order", "label": "Customer Places Order", "type": "process"},
        {"id": "check_stock", "label": "Check Stock Availability", "type": "decision"},
        {"id": "stock_available", "label": "Stock Available?", "type": "decision"},
        {"id": "process_payment", "label": "Process Payment", "type": "process"},
        {"id": "order_confirmed", "label": "Order Confirmed", "type": "process"},
        {"id": "notify_customer_oos", "label": "Notify Customer (Out of Stock)", "type": "process"},
        {"id": "end_order_process", "label": "End Order Process", "type": "end"}
      ],
      "connections": [
        {"from": "start_node", "to": "customer_order", "label": "Initiate"},
        {"from": "customer_order", "to": "check_stock", "label": "Proceed"},
        {"from": "check_stock", "to": "stock_available", "label": "Result"},
        {"from": "stock_available", "to": "process_payment", "label": "Yes"},
        {"from": "stock_available", "to": "notify_customer_oos", "label": "No"},
        {"from": "process_payment", "to": "order_confirmed", "label": "Success"},
        {"from": "order_confirmed", "to": "end_order_process", "label": "Complete"},
        {"from": "notify_customer_oos", "to": "end_order_process", "label": "Finish"}
      ]
    }
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        # Validate required top-level keys for a flowchart
        if 'start_node' not in data or 'nodes' not in data or 'connections' not in data:
            raise ValueError("Missing required fields: 'start_node', 'nodes', or 'connections'")

        # Define specific node shapes for flowchart types
        node_shapes = {
            "process": "box",          # Rectangle for processes
            "decision": "diamond",     # Diamond for decisions
            "start": "oval",           # Oval for start
            "end": "oval",             # Oval for end
            "io": "parallelogram",     # Input/Output
            "document": "note",        # Document symbol
            "default": "box"           # Fallback
        }

        dot = graphviz.Digraph(
            name='ProcessFlowDiagram',
            format='png',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom flow is common for flowcharts
                'splines': 'ortho',     # Straight lines with 90-degree bends
                'bgcolor': 'white',     # White background
                'pad': '0.5',           # Padding around the graph
                'nodesep': '0.6',       # Spacing between nodes on same rank
                'ranksep': '0.8'        # Spacing between ranks
            }
        )
        
        base_color = '#19191a' # Hardcoded base color

        # Helper for color and font based on depth (simulated for flowcharts if needed, or consistent color)
        # For flowcharts, often all process/decision nodes are the same color.
        # If a gradient is desired, a 'depth' could be manually assigned or derived.
        # For simplicity and typical flowchart appearance, we'll use a consistent base_color for all nodes.
        fill_color_for_nodes = base_color
        font_color_for_nodes = 'white' if base_color == '#19191a' or base_color.lower() in ['#000000', '#19191a'] else 'black'
        
        # Store all nodes by ID for easy lookup
        all_defined_nodes = {node['id']: node for node in data['nodes']}
        
        # Add start node explicitly
        start_node_id = data['start_node']
        dot.node(
            start_node_id,
            start_node_id, # Label is typically the ID itself for start/end
            shape=node_shapes['start'],
            style='filled,rounded',
            fillcolor='#2196F3', # A distinct blue for Start
            fontcolor='white',
            fontsize='14'
        )

        # Add all other nodes (process, decision, etc.)
        for node_id, node_info in all_defined_nodes.items():
            if node_id == start_node_id: # Skip if it's the start node, already added
                continue

            node_type = node_info.get("type", "default")
            shape = node_shapes.get(node_type, "box")

            node_label = node_info['label']

            # Use distinct color for end node if it exists
            if node_type == 'end':
                 dot.node(
                    node_id,
                    node_label,
                    shape=shape,
                    style='filled,rounded',
                    fillcolor='#F44336', # A distinct red for End
                    fontcolor='white',
                    fontsize='14'
                )
            else: # Regular process, decision, etc. nodes use the selected base color
                dot.node(
                    node_id,
                    node_label,
                    shape=shape,
                    style='filled,rounded',
                    fillcolor=fill_color_for_nodes, 
                    fontcolor=font_color_for_nodes,
                    fontsize='14'
                )

        # Add connections (edges)
        for connection in data['connections']:
            dot.edge(
                connection['from'],
                connection['to'],
                label=connection.get('label', ''),
                color='#4a4a4a', # Dark gray for lines
                fontcolor='#4a4a4a',
                fontsize='10'
            )
        
        # Save to temporary file
        with NamedTemporaryFile(delete=False, suffix='.png') as tmp:
            dot.render(tmp.name, format='png', cleanup=True)
            return tmp.name + '.png'

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"