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

import re
import os
from pathlib import Path

def parse_trace_file(filepath):
    """Parse a trace output file and extract execution phases."""
    data = {}
    
    with open(filepath, 'r') as f:
        content = f.read()
    
    # Remove ANSI color codes
    content = re.sub(r'\x1b\[[0-9;]*m', '', content)
    
    # Find all lines with SUCCESS
    # Format: "Assembly     SUCCESS 8 ms (Startup: 0ms, Calc: 5ms, I/O: 0ms)"
    pattern = r'([\w\-]+)\s+SUCCESS\s+(\d+)\s+ms\s+\(Startup:\s+(\d+)ms,\s+Calc:\s+(\d+)ms,\s+I/O:\s+(\d+)ms\)'
    matches = re.findall(pattern, content)
    
    for match in matches:
        lang = match[0]
        total_time = int(match[1])
        startup_time = int(match[2])
        calc_time = int(match[3])
        io_time = int(match[4])
        
        data[lang] = {
            'total_time': total_time,
            'startup_time': startup_time,
            'calc_time': calc_time,
            'io_time': io_time
        }
    
    return data

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['startup_time']
        calc = metrics['calc_time']
        io = metrics['io_time']
        
        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]
    
    # Calculate max time for y-axis
    max_time = max([m['total_time'] for _, m in top_20]) + 10
    
    chart = f"""```mermaid
xychart-beta
    title "{title}"
    x-axis [{', '.join([f'"{lang}"' for lang, _ in top_20])}]
    y-axis "Tid (ms)" 0 --> {max_time}
    
    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."""
    
    reports_dir = Path("reports")
    if not reports_dir.exists():
        print("No reports directory found. Run run_trace.sh first.")
        return
    
    # Generate Gantt diagrams for each decimal level
    for decimals in [1, 2, 5, 10, 100, 1000, 2000]:
        trace_file = reports_dir / f"trace_{decimals}_output.txt"
        
        if not trace_file.exists():
            print(f"Warning: {trace_file} not found, skipping...")
            continue
        
        # Parse trace data
        data = parse_trace_file(trace_file)
        
        if not data:
            print(f"Warning: No data found in {trace_file}, skipping...")
            continue
        
        # Generate Gantt diagram
        gantt = generate_gantt_diagram(
            data, 
            f"Exekveringsfaser - {decimals} decimaler"
        )
        
        # Generate stacked bar chart
        stacked = generate_stacked_bar_chart(
            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()