#!/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]))
                    # Memory is now in bytes, convert to MB for display
                    memories.append(int(parts[1]) / (1024 * 1024))  # Convert bytes 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 with memory and CPU."""
    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]
    
    # Find max values for scaling
    max_time = max(times)
    max_memory = max(memories) if max(memories) > 0 else 1
    max_cpu = max(cpus) if max(cpus) > 0 else 1
    
    chart = f"```mermaid\n"
    chart += f"xychart-beta\n"
    chart += f'    title "{language} - Tidslinje: Minne (MB) och CPU (%)"\n'
    chart += f'    x-axis "Tid (ms)" 0 --> {max_time}\n'
    chart += f'    y-axis "Minne (MB) / CPU (%)" 0 --> {max(max_memory, max_cpu) + 1}\n'
    # Memory line (in MB)
    chart += f'    line[{", ".join([f"{m:.2f}" for m in memories])}]\n'
    # CPU line (in %)
    chart += f'    line[{", ".join([f"{c}" for c in cpus])}]\n'
    chart += f"```\n"
    
    return chart

def generate_comparison_chart(languages_data, title, max_time=None):
    """Generate comparison chart for multiple languages showing memory and CPU."""
    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, cpus in languages_data:
        if len(times) > max_points:
            step = len(times) // max_points
            times = times[::step]
            memories = memories[::step]
            cpus = cpus[::step]
        sampled_data.append((lang, times, memories, cpus))
    
    # Find max values for scaling
    max_memory = max(max(m) if m else 0 for _, _, m, _ in sampled_data)
    max_cpu = max(max(c) if c else 0 for _, _, _, c in sampled_data)
    max_y = max(max_memory, max_cpu) + 1
    
    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) / CPU (%)" 0 --> {max_y}\n'
    
    # Add memory lines for each language
    for lang, times, memories, cpus in sampled_data:
        if len(memories) > 0:
            chart += f'    line[{", ".join([f"{m:.2f}" for m in memories])}]\n'
    
    # Add CPU lines for each language
    for lang, times, memories, cpus in sampled_data:
        if len(cpus) > 0:
            chart += f'    line[{", ".join([f"{c}" for c in cpus])}]\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 (MB) och CPU-användning (%) över tid.\n")
        f.write("Varje diagram har två linjer: minne (övre) och CPU (undre).\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, c) 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 - Minne och CPU över tid")
            if chart:
                f.write(chart)
                f.write("\n")
        
        # Medium languages (50-200ms)
        medium = [(lang, t, m, c) 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 - Minne och CPU över tid")
            if chart:
                f.write(chart)
                f.write("\n")
        
        # Slow languages (200ms+)
        slow = [(lang, t, m, c) 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 - Minne och CPU över tid")
            if chart:
                f.write(chart)
                f.write("\n")
    
    print(f"Generated Mermaid charts in {output_file}")

if __name__ == '__main__':
    main()