#!/usr/bin/env python3
"""
Generate Mermaid XY charts from timeline data.
Reads TSV files and creates Mermaid diagram code for README.md.
"""

import os
import sys
from pathlib import Path

def read_timeline_data(filepath):
    """Read timeline TSV file and return time, memory, cpu arrays."""
    times = []
    memories = []
    cpus = []
    
    with open(filepath, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 3:
                try:
                    times.append(int(parts[0]))
                    memories.append(int(parts[1]) / 1024)  # Convert KB to MB
                    cpus.append(int(parts[2]))
                except ValueError:
                    continue
    
    return times, memories, cpus

def generate_mermaid_chart(language, times, memories, cpus):
    """Generate Mermaid XY chart for a language."""
    if len(times) == 0:
        return None
    
    # Sample data if too many points (Mermaid has limits)
    max_points = 20
    if len(times) > max_points:
        step = len(times) // max_points
        times = times[::step]
        memories = memories[::step]
        cpus = cpus[::step]
    
    chart = f"```mermaid\n"
    chart += f"xychart-beta\n"
    chart += f'    title "{language} - Resursanvändning över tid"\n'
    chart += f'    x-axis "Tid (ms)" 0 --> {max(times)}\n'
    chart += f'    y-axis "Minne (MB)" 0 --> {max(memories) + 1}\n'
    chart += f'    line [{", ".join([f"{m:.1f}" for m in memories])}]\n'
    chart += f"```\n"
    
    return chart

def generate_comparison_chart(languages_data, title, max_time=None):
    """Generate comparison chart for multiple languages."""
    if not languages_data or all(len(data[1]) == 0 for data in languages_data):
        return None
    
    # Find max time for x-axis
    if max_time is None:
        max_time = max(max(data[1]) if len(data[1]) > 0 else 0 for data in languages_data)
    
    # Sample to max 10 points per language
    max_points = 10
    sampled_data = []
    for lang, times, memories in languages_data:
        if len(times) > max_points:
            step = len(times) // max_points
            times = times[::step]
            memories = memories[::step]
        sampled_data.append((lang, times, memories))
    
    chart = f"```mermaid\n"
    chart += f"xychart-beta\n"
    chart += f'    title "{title}"\n'
    chart += f'    x-axis "Tid (ms)" 0 --> {max_time}\n'
    chart += f'    y-axis "Minne (MB)" 0 --> {max(max(m) if m else 0 for _, _, m in sampled_data) + 1}\n'
    
    for lang, times, memories in sampled_data:
        if len(memories) > 0:
            chart += f'    line [{", ".join([f"{m:.1f}" for m in memories])}]\n'
    
    chart += f"```\n"
    
    return chart

def main():
    timelines_dir = Path('timelines')
    output_file = Path('mermaid_charts.md')
    
    if not timelines_dir.exists():
        print(f"Error: Directory {timelines_dir} does not exist")
        sys.exit(1)
    
    # Collect data for all languages
    all_languages = []
    
    # Process each language
    for lang_dir in sorted(timelines_dir.iterdir()):
        if lang_dir.is_dir():
            language = lang_dir.name
            
            # Find the best run (usually run_2.tsv)
            timeline_file = lang_dir / 'run_2.tsv'
            
            if timeline_file.exists():
                times, memories, cpus = read_timeline_data(timeline_file)
                
                if len(times) > 0:
                    all_languages.append((language, times, memories, cpus))
    
    # Generate output
    with open(output_file, 'w') as f:
        f.write("# Resursanvändning över tid\n\n")
        f.write("Följande diagram visar minnesanvändning över tid för varje språk.\n\n")
        
        # Individual charts for each language
        for language, times, memories, cpus in all_languages:
            chart = generate_mermaid_chart(language, times, memories, cpus)
            if chart:
                f.write(f"## {language}\n\n")
                f.write(chart)
                f.write("\n")
        
        # Comparison charts
        f.write("# Jämförelser\n\n")
        
        # Fast languages (under 50ms)
        fast = [(lang, t, m) for lang, t, m, c in all_languages if len(t) > 0 and t[-1] < 50]
        if fast:
            f.write("## Snabba språk (< 50ms)\n\n")
            chart = generate_comparison_chart(fast, "Snabba språk - Minnesanvändning")
            if chart:
                f.write(chart)
                f.write("\n")
        
        # Medium languages (50-200ms)
        medium = [(lang, t, m) for lang, t, m, c in all_languages 
                  if len(t) > 0 and 50 <= t[-1] < 200]
        if medium:
            f.write("## Medelsnabba språk (50-200ms)\n\n")
            chart = generate_comparison_chart(medium, "Medelsnabba språk - Minnesanvändning")
            if chart:
                f.write(chart)
                f.write("\n")
        
        # Slow languages (200ms+)
        slow = [(lang, t, m) for lang, t, m, c in all_languages if len(t) > 0 and t[-1] >= 200]
        if slow:
            f.write("## Långsamma språk (200ms+)\n\n")
            chart = generate_comparison_chart(slow, "Långsamma språk - Minnesanvändning")
            if chart:
                f.write(chart)
                f.write("\n")
    
    print(f"Generated Mermaid charts in {output_file}")

if __name__ == '__main__':
    main()