wbs_diagram_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os
from graph_generator_utils import add_nodes_and_edges

def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
    """
    Generates a Work Breakdown Structure (WBS) Diagram from JSON input.

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

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

    Expected JSON Format Example:
    {
      "project_title": "Software Development Project",
      "phases": [
        {
          "id": "phase_prep",
          "label": "1. Preparation",
          "tasks": [
            {"id": "task_vision", "label": "1.1. Identify Vision"},
            {"id": "task_design", "label": "1.2. Design & Staffing"}
          ]
        },
        {
          "id": "phase_plan",
          "label": "2. Planning",
          "tasks": [
            {"id": "task_cost", "label": "2.1. Cost Analysis"},
            {"id": "task_benefit", "label": "2.2. Benefit Analysis"},
            {"id": "task_risk", "label": "2.3. Risk Assessment"}
          ]
        }
      ]
    }
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'project_title' not in data or 'phases' not in data:
            raise ValueError("Missing required fields: project_title or phases")

        dot = graphviz.Digraph(
            name='WBSDiagram',
            format='png',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom hierarchy
                'splines': 'ortho',     # Straight lines
                'bgcolor': 'white',     # White background
                'pad': '0.5',           # Padding
                'ranksep': '0.8',       # Adjust vertical separation between ranks
                'nodesep': '0.5'        # Adjust horizontal separation between nodes
            }
        )
        
        base_color = '#19191a' # Hardcoded base color

        # Project Title node (main node)
        dot.node(
            'project_root',
            data['project_title'],
            shape='box',
            style='filled,rounded',
            fillcolor=base_color, # Use the selected base color
            fontcolor='white',
            fontsize='18'
        )

        # Logic for color gradient within WBS specific nodes (Phases and Tasks)
        # This ensures the gradient works correctly with the hardcoded base_color
        def get_gradient_color(depth, base_hex_color, lightening_factor=0.12):
            base_r = int(base_hex_color[1:3], 16)
            base_g = int(base_hex_color[3:5], 16)
            base_b = int(base_hex_color[5:7], 16)

            current_r = base_r + int((255 - base_r) * depth * lightening_factor)
            current_g = base_g + int((255 - base_g) * depth * lightening_factor)
            current_b = base_b + int((255 - base_b) * depth * lightening_factor)
            
            return f'#{min(255, current_r):02x}{min(255, current_g):02x}{min(255, current_b):02x}'

        def get_font_color_for_background(depth, base_hex_color, lightening_factor=0.12):
            # Calculate brightness/lightness of the node color at this depth
            # and return black/white for text accordingly
            base_r = int(base_hex_color[1:3], 16)
            base_g = int(base_hex_color[3:5], 16)
            base_b = int(base_hex_color[5:7], 16)
            current_r = base_r + (255 - base_r) * depth * lightening_factor
            current_g = base_g + (255 - base_g) * depth * lightening_factor
            current_b = base_b + (255 - base_b) * depth * lightening_factor
            
            # Simple luminance check (ITU-R BT.709 coefficients)
            luminance = (0.2126 * current_r + 0.7152 * current_g + 0.0722 * current_b) / 255
            return 'white' if luminance < 0.5 else 'black'


        current_depth = 1 # Depth for phases
        for phase in data['phases']:
            phase_id = phase.get('id')
            phase_label = phase.get('label')
            tasks = phase.get('tasks', [])

            if not all([phase_id, phase_label]):
                raise ValueError(f"Invalid phase: {phase}")

            phase_fill_color = get_gradient_color(current_depth, base_color)
            phase_font_color = get_font_color_for_background(current_depth, base_color)
            
            dot.node(
                phase_id,
                phase_label,
                shape='box',
                style='filled,rounded',
                fillcolor=phase_fill_color,
                fontcolor=phase_font_color,
                fontsize='14'
            )
            dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none') # Connect to root

            task_depth = current_depth + 1 # Depth for tasks
            task_font_size = max(9, 14 - (task_depth * 2))

            task_nodes_in_phase = []
            for task in tasks:
                task_id = task.get('id')
                task_label = task.get('label')
                if not all([task_id, task_label]):
                    raise ValueError(f"Invalid task: {task}")
                
                task_fill_color = get_gradient_color(task_depth, base_color)
                task_font_color = get_font_color_for_background(task_depth, base_color)

                dot.node(
                    task_id,
                    task_label,
                    shape='box',
                    style='filled,rounded',
                    fillcolor=task_fill_color,
                    fontcolor=task_font_color,
                    fontsize=str(task_font_size)
                )
                dot.edge(phase_id, task_id, color='#4a4a4a', arrowhead='none') # Connect task to phase
                task_nodes_in_phase.append(task_id)
            
            # Use subgraph to enforce vertical alignment for tasks within a phase
            if task_nodes_in_phase: # Only create subgraph if there are tasks
                dot.subgraph(name=f'cluster_{phase_id}')
        # 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)}"