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 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.

    Expected JSON Format Example:
    {
      "title": "AI Development Timeline",
      "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 PNG 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': 'LR',        # Left-to-Right layout (horizontal timeline)
                'splines': 'ortho',     # Straight lines
                'bgcolor': 'white',     # White background
                'pad': '0.5',           # Padding around the graph
                'nodesep': '1.0',       # Spacing between nodes
                'ranksep': '2.0'        # Spacing between ranks
            }
        )
        
        base_color = '#19191a' # Hardcoded base color
        
        title = data.get('title', '')
        events = data.get('events', [])
        
        if not events:
            raise ValueError("Timeline must contain at least one event")

        # Add title node if provided
        if title:
            dot.node(
                'title',
                title,
                shape='plaintext',
                fontsize='18',
                fontweight='bold',
                fontcolor=base_color
            )
        
        # Add timeline events
        previous_event_id = None
        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', '')
            
            # Create full label with date and 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
            
            # Calculate color opacity based on position in timeline
            if total_events == 1:
                opacity = 'FF'
            else:
                opacity_value = int(255 * (1.0 - (i * 0.7 / (total_events - 1))))
                opacity = format(opacity_value, '02x')
            
            node_color = f"{base_color}{opacity}"
            font_color = 'white' if i < total_events * 0.7 else 'black'
            
            # Add the event node
            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'
            )
            
            # Connect to previous event if exists
            if previous_event_id:
                dot.edge(
                    previous_event_id,
                    event_id,
                    color='#666666',
                    arrowsize='0.8',
                    penwidth='2'
                )
            
            # Connect title to first event if title exists
            if title and i == 0:
                dot.edge('title', event_id, style='invis')
            
            previous_event_id = event_id

        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)}"