process_flow_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os
from graph_generator_utils import add_nodes_and_edges # Reusing common utility

def generate_process_flow_diagram(json_input: str, base_color: str) -> str:
    """
    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.
        base_color (str): The hexadecimal color string (e.g., '#19191a') for the base
                          color of the nodes, from which a gradient will be generated.

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

    Expected JSON Format Example:
    {
      "start_node": "Start Process",
      "nodes": [
        {
          "id": "step1",
          "label": "Gather Requirements",
          "type": "process",
          "relationship": "Next"
        },
        {
          "id": "decision1",
          "label": "Is Data Available?",
          "type": "decision",
          "relationship": "Check"
        },
        {
          "id": "path_yes",
          "label": "Process Data",
          "type": "process",
          "relationship": "Yes"
        },
        {
          "id": "path_no",
          "label": "Collect More Data",
          "type": "process",
          "relationship": "No"
        },
        {
          "id": "end_node",
          "label": "End Process",
          "type": "end"
        }
      ],
      "connections": [
        {"from": "start_node", "to": "step1", "label": "Start"},
        {"from": "step1", "to": "decision1", "label": "Continue"},
        {"from": "decision1", "to": "path_yes", "label": "Yes"},
        {"from": "decision1", "to": "path_no", "label": "No"},
        {"from": "path_yes", "to": "end_node", "label": "Done"},
        {"from": "path_no", "to": "end_node", "label": "Done"}
      ]
    }
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        # Determine 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
            }
        )
        
        # Ensure base_color is valid, fallback if not
        if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
            base_color = '#19191a' # Fallback to default dark if invalid

        # Add all nodes based on JSON structure
        all_nodes_data = {}
        # Special handling for start_node if defined separately
        if 'start_node' in data:
            all_nodes_data[data['start_node']] = {"label": data['start_node'], "type": "start"}
        
        for node_data in data.get('nodes', []):
            all_nodes_data[node_data['id']] = node_data

        for node_id, node_info in all_nodes_data.items():
            node_type = node_info.get("type", "default")
            shape = node_shapes.get(node_type, "box") # Default to box if type is unknown

            # Use base_color for all nodes for consistent flowchart look
            # You can adapt add_nodes_and_edges if you want color gradient per depth for flowcharts too
            fill_color_for_node = base_color
            font_color_for_node = 'white' if base_color == '#19191a' else 'black' # Keep text readable

            dot.node(
                node_id,
                node_info['label'],
                shape=shape,
                style='filled,rounded', # Keep rounded corners
                fillcolor=fill_color_for_node, 
                fontcolor=font_color_for_node,
                fontsize='14'
            )

        # Add connections (edges)
        for connection in data.get('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)}"