wbs_diagram_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_wbs_diagram(json_input: str, output_format: str) -> str:
    """
    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.

    Expected JSON Format Example:
    {
      "project_title": "AI Model Development Project",
      "phases": [
        {
          "id": "phase_prep",
          "label": "Preparation",
          "tasks": [
            {
              "id": "task_1_1_vision",
              "label": "Identify Vision",
              "subtasks": [
                {
                  "id": "subtask_1_1_1_design_staff",
                  "label": "Design & Staffing",
                  "sub_subtasks": [
                    {
                      "id": "ss_task_1_1_1_1_env_setup",
                      "label": "Environment Setup",
                      "sub_sub_subtasks": [
                        {
                          "id": "sss_task_1_1_1_1_1_lib_install",
                          "label": "Install Libraries",
                          "final_level_tasks": [
                            {"id": "ft_1_1_1_1_1_1_data_access", "label": "Grant Data Access"}
                          ]
                        }
                      ]
                    }
                  ]
                }
              ]
            }
          ]
        },
        {
          "id": "phase_plan",
          "label": "Planning",
          "tasks": [
            {
              "id": "task_2_1_cost_analysis",
              "label": "Cost Analysis",
              "subtasks": [
                {
                  "id": "subtask_2_1_1_benefit_analysis",
                  "label": "Benefit Analysis",
                  "sub_subtasks": [
                    {
                      "id": "ss_task_2_1_1_1_risk_assess",
                      "label": "AI Risk Assessment",
                      "sub_sub_subtasks": [
                        {
                          "id": "sss_task_2_1_1_1_1_model_selection",
                          "label": "Model Selection",
                          "final_level_tasks": [
                            {"id": "ft_2_1_1_1_1_1_data_strategy", "label": "Data Strategy"}
                          ]
                        }
                      ]
                    }
                  ]
                }
              ]
            }
          ]
        },
        {
          "id": "phase_dev",
          "label": "Development",
          "tasks": [
            {
              "id": "task_3_1_change_mgmt",
              "label": "Data Preprocessing",
              "subtasks": [
                {
                  "id": "subtask_3_1_1_implementation",
                  "label": "Feature Engineering",
                  "sub_subtasks": [
                    {
                      "id": "ss_task_3_1_1_1_beta_testing",
                      "label": "Model Training",
                      "sub_sub_subtasks": [
                        {
                          "id": "sss_task_3_1_1_1_1_other_task",
                          "label": "Model Evaluation",
                          "final_level_tasks": [
                            {"id": "ft_3_1_1_1_1_1_hyperparam_tune", "label": "Hyperparameter Tuning"}
                          ]
                        }
                      ]
                    }
                  ]
                }
              ]
            }
          ]
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    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")

        # 한글 폰트 설정
        # GDFONTPATH가 설정되어 있으면 폰트 파일명(확장자 제외) 사용
        korean_font = 'NanumGothic-Regular'

        dot = graphviz.Digraph(
            name='WBSDiagram',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom hierarchy
                'splines': 'polyline',  # polyline으로 변경 (ortho 대신)
                'bgcolor': 'white',     # White background
                'pad': '0.5',           # Padding
                'ranksep': '0.6',       # Adjust vertical separation between ranks
                'nodesep': '0.5',       # Adjust horizontal separation between nodes
                'fontname': korean_font,  # 그래프 전체 한글 폰트
                'charset': 'UTF-8'      # UTF-8 인코딩
            },
            node_attr={
                'fontname': korean_font  # 모든 노드의 기본 폰트
            },
            edge_attr={
                'fontname': korean_font  # 모든 엣지의 기본 폰트
            }
        )
        
        base_color = '#19191a' # Hardcoded base color
        
        # ID 정규화 함수 - 한글 ID를 안전한 형태로 변환
        def normalize_id(id_str):
            """노드 ID를 안전한 형태로 변환"""
            import re
            # 영문, 숫자, 언더스코어만 허용
            safe_id = re.sub(r'[^a-zA-Z0-9_]', '_', str(id_str))
            # 숫자로 시작하면 'n_' 접두사 추가
            if safe_id and safe_id[0].isdigit():
                safe_id = 'n_' + safe_id
            # 빈 문자열이면 기본값
            if not safe_id:
                safe_id = 'node_' + str(hash(id_str))
            return safe_id

        # Project Title node (main node)
        dot.node(
            'project_root',
            data['project_title'],
            shape='box',
            style='filled,rounded',
            fillcolor=base_color,
            fontcolor='white',
            fontsize='18',
            fontname=korean_font  # 한글 폰트 추가
        )

        # Helper for color and font based on depth for WBS
        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):
            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
            
            luminance = (0.2126 * current_r + 0.7152 * current_g + 0.0722 * current_b) / 255
            return 'white' if luminance < 0.5 else 'black'

        def _add_wbs_nodes_recursive(parent_id, current_level_tasks, current_depth):
            for task_data in current_level_tasks:
                task_id = task_data.get('id')
                task_label = task_data.get('label')
                
                if not all([task_id, task_label]):
                    raise ValueError(f"Invalid task data at depth {current_depth}: {task_data}")
                
                # ID 정규화
                safe_task_id = normalize_id(task_id)

                node_fill_color = get_gradient_color(current_depth, base_color)
                node_font_color = get_font_color_for_background(current_depth, base_color)
                font_size = max(9, 14 - (current_depth * 2))

                dot.node(
                    safe_task_id,
                    task_label,
                    shape='box',
                    style='filled,rounded',
                    fillcolor=node_fill_color,
                    fontcolor=node_font_color,
                    fontsize=str(font_size),
                    fontname=korean_font  # 한글 폰트 추가
                )
                dot.edge(parent_id, safe_task_id, color='#4a4a4a', arrowhead='none', fontname=korean_font)

                # Recursively call for next level of tasks (subtasks, sub_subtasks, etc.)
                # This handles arbitrary nested keys like 'subtasks', 'sub_subtasks', 'final_level_tasks'
                next_level_keys = ['tasks', 'subtasks', 'sub_subtasks', 'sub_sub_subtasks', 'final_level_tasks']
                for key_idx, key in enumerate(next_level_keys):
                    if key in task_data and isinstance(task_data[key], list):
                        _add_wbs_nodes_recursive(safe_task_id, task_data[key], current_depth + 1)
                        break # Only process the first found sub-level key

        # Process phases (level 1 from project_root)
        phase_depth = 1
        for phase in data['phases']:
            phase_id = phase.get('id')
            phase_label = phase.get('label')
            
            if not all([phase_id, phase_label]):
                raise ValueError(f"Invalid phase data: {phase}")
            
            # ID 정규화
            safe_phase_id = normalize_id(phase_id)

            phase_fill_color = get_gradient_color(phase_depth, base_color)
            phase_font_color = get_font_color_for_background(phase_depth, base_color)
            font_size_phase = max(9, 14 - (phase_depth * 2))

            dot.node(
                safe_phase_id,
                phase_label,
                shape='box',
                style='filled,rounded',
                fillcolor=phase_fill_color,
                fontcolor=phase_font_color,
                fontsize=str(font_size_phase),
                fontname=korean_font  # 한글 폰트 추가
            )
            dot.edge('project_root', safe_phase_id, color='#4a4a4a', arrowhead='none', fontname=korean_font)

            # Start recursion for tasks under this phase
            if 'tasks' in phase and isinstance(phase['tasks'], list):
                _add_wbs_nodes_recursive(safe_phase_id, phase['tasks'], phase_depth + 1)

        # 렌더링
        try:
            with NamedTemporaryFile(delete=False, suffix='.gv', prefix='wbs_') as tmp:
                # 파일 이름에서 .gv 확장자 제거
                output_filename = tmp.name[:-3]  # '.gv' 제거
                output_path = dot.render(output_filename, format=output_format, cleanup=True)
                return output_path
        except Exception as render_error:
            # 렌더링 실패 시 간단한 에러 메시지
            return f"Error: Failed to render diagram - {str(render_error).split(';')[0]}"

    except json.JSONDecodeError as e:
        return "Error: Invalid JSON format"
    except Exception as e:
        # 에러 메시지를 간단하게 유지
        error_msg = str(e).split('\n')[0]  # 첫 줄만 사용
        if len(error_msg) > 100:
            error_msg = error_msg[:100] + "..."
        return f"Error: {error_msg}"