#!/usr/bin/env python3
"""
Generate Gantt diagrams from trace data showing execution phases.
"""

import json
import os
from pathlib import Path

def generate_gantt_diagram(data, title):
    """Generate a Mermaid Gantt diagram from trace data."""
    
    # Sort by total time
    sorted_data = sorted(data.items(), key=lambda x: x[1]['total_time'])
    
    # Take top 20 for readability
    top_20 = sorted_data[:20]
    
    gantt = f"""```mermaid
gantt
    title {title}
    dateFormat X
    axisFormat %L ms
    
"""
    
    # Add sections for each language
    for lang, metrics in top_20:
        startup = metrics.get('startup_time', 0)
        calc = metrics.get('calc_time', 0)
        io = metrics.get('io_time', 0)
        
        gantt += f"    section {lang}\n"
        gantt += f"    Startup     :0, {startup}\n"
        gantt += f"    Calculation :{startup}, {startup + calc}\n"
        gantt += f"    I/O         :{startup + calc}, {startup + calc + io}\n"
    
    gantt += "```\n"
    
    return gantt

def generate_stacked_bar_chart(data, title):
    """Generate a Mermaid stacked bar chart showing execution phases."""
    
    # Sort by total time
    sorted_data = sorted(data.items(), key=lambda x: x[1]['total_time'])
    
    # Take top 20 for readability
    top_20 = sorted_data[:20]
    
    chart = f"""```mermaid
xychart-beta
    title "{title}"
    x-axis [{', '.join([f'"{lang}"' for lang, _ in top_20])}]
    y-axis "Tid (ms)" 0 --> {max([m['total_time'] for _, m in top_20]) + 10}
    
    bar Startup [{', '.join([str(m['startup_time']) for _, m in top_20])}]
    bar Calculation [{', '.join([str(m['calc_time']) for _, m in top_20])}]
    bar I/O [{', '.join([str(m['io_time']) for _, m in top_20])}]
```

"""
    
    return chart

def main():
    """Main function to generate Gantt diagrams."""
    
    # Check if traces directory exists
    traces_dir = Path("traces")
    if not traces_dir.exists():
        print("No traces directory found. Run run_trace.sh first.")
        return
    
    # Load trace data
    trace_file = traces_dir / "trace_data.json"
    if not trace_file.exists():
        print("No trace data found. Run run_trace.sh first.")
        return
    
    with open(trace_file) as f:
        data = json.load(f)
    
    # Generate Gantt diagrams for each decimal level
    reports_dir = Path("reports")
    reports_dir.mkdir(exist_ok=True)
    
    for decimals in [1, 2, 5, 10, 100, 1000, 2000]:
        if str(decimals) in data:
            decimal_data = data[str(decimals)]
            
            # Generate Gantt diagram
            gantt = generate_gantt_diagram(
                decimal_data, 
                f"Exekveringsfaser - {decimals} decimaler"
            )
            
            # Generate stacked bar chart
            stacked = generate_stacked_bar_chart(
                decimal_data,
                f"Exekveringsfaser - {decimals} decimaler"
            )
            
            # Write to report
            report_file = reports_dir / f"{decimals}_decimals_gantt.md"
            with open(report_file, 'w') as f:
                f.write(f"# Exekveringsfaser - {decimals} decimaler\n\n")
                f.write("## Gantt-diagram\n\n")
                f.write(gantt)
                f.write("\n## Staplat stapeldiagram\n\n")
                f.write(stacked)
            
            print(f"Generated Gantt diagram for {decimals} decimals")
    
    print("Done!")

if __name__ == "__main__":
    main()