timeline_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_timeline_diagram(json_input: str, output_format: str) -> str:
    """
    Generates a serpentine timeline diagram from JSON input.

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

        output_format (str): The output format for the generated diagram.
                            Supported formats: "png" or "svg"

    Expected JSON Format Example:
    {
      "title": "AI Development Timeline",
      "events_per_row": 4,
      "events": [
        {
          "id": "event_1",
          "label": "Machine Learning Foundations",
          "date": "1950-1960",
          "description": "Early neural networks and perceptrons"
        },
        {
          "id": "event_2",
          "label": "Expert Systems Era",
          "date": "1970-1980",
          "description": "Rule-based AI systems"
        },
        {
          "id": "event_3",
          "label": "Neural Network Revival",
          "date": "1980-1990",
          "description": "Backpropagation algorithm"
        }
      ]
    }

    Returns:
        str: The filepath to the generated image file.
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'events' not in data:
            raise ValueError("Missing required field: events")

        dot = graphviz.Digraph(
            name='Timeline',
            format='png',
            graph_attr={
                'rankdir': 'TB',
                'splines': 'ortho',
                'bgcolor': 'white',
                'pad': '0.8',
                'nodesep': '3.0',
                'ranksep': '2.5'
            }
        )
        
        base_color = '#BEBEBE'
        
        title = data.get('title', '')
        events = data.get('events', [])
        events_per_row = data.get('events_per_row', 4)  
        
        if not events:
            raise ValueError("Timeline must contain at least one event")

        if title:
            dot.node(
                'title',
                title,
                shape='plaintext',
                fontsize='18',
                fontweight='bold',
                fontcolor='#000000',
                pos="6,2!" 
            )
        
        total_events = len(events)
        
        for i, event in enumerate(events):
            event_id = event.get('id', f'event_{i}')
            event_label = event.get('label', f'Event {i+1}')
            event_date = event.get('date', '')
            event_description = event.get('description', '')
            
            if event_date and event_description:
                full_label = f"{event_date}\\n{event_label}\\n{event_description}"
            elif event_date:
                full_label = f"{event_date}\\n{event_label}"
            elif event_description:
                full_label = f"{event_label}\\n{event_description}"
            else:
                full_label = event_label
            
            lightening_factor = 0.0417  
            base_r = int(base_color[1:3], 16)
            base_g = int(base_color[3:5], 16)
            base_b = int(base_color[5:7], 16)
            
            current_r = base_r + int((255 - base_r) * i * lightening_factor)
            current_g = base_g + int((255 - base_g) * i * lightening_factor)
            current_b = base_b + int((255 - base_b) * i * lightening_factor)
            
            current_r = min(255, current_r)
            current_g = min(255, current_g)
            current_b = min(255, current_b)
            
            node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
            font_color = 'black'
            
            row = i // events_per_row
            col = i % events_per_row
            
            if row % 2 == 1:
                visual_col = events_per_row - 1 - col
            else:
                visual_col = col
            
            dot.node(
                event_id,
                full_label,
                shape='box',
                style='filled,rounded',
                fillcolor=node_color,
                fontcolor=font_color,
                fontsize='12',
                width='2.5',
                height='1.2',
                pos=f"{visual_col * 4.5},{-row * 3}!"  
            )
        
        for i in range(len(events) - 1):
            current_event_id = events[i].get('id', f'event_{i}')
            next_event_id = events[i + 1].get('id', f'event_{i + 1}')
            
            dot.edge(
                current_event_id,
                next_event_id,
                color='#4a4a4a',
                arrowsize='0.8',
                penwidth='2'
            )
        
        dot.engine = 'neato'

        with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
            dot.render(tmp.name, format=output_format, cleanup=True)
            return f"{tmp.name}.{output_format}"

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